public static GeoSearchConfig getGeoSearchConfig() {
GeoSearchConfig geoConfig = new GeoSearchConfig();
geoConfig.addFieldBitMask(LOCATION_FIELD, (byte) 1);
return geoConfig;
}
public GeoIndexer(GeoSearchConfig config) {
geoUtil = config.getGeoUtil();
geoConverter = config.getGeoConverter();
fieldTree = geoUtil.getBinaryTreeOrderedByBitMag();
geoRecordSerializer = new CartesianGeoRecordSerializer();
this.config = config;
}
GeoSegmentReader<IDGeoRecord> getGeoSegmentReader() throws IOException {
String fileName = indexName + "." + config.getGeoFileExtension();
return new GeoSegmentReader<IDGeoRecord>(
directory,
fileName,
-1,
config.getBufferSizePerGeoSegmentReader(),
geoRecordSerializer,
geoComparator);
}
GeoSegmentWriter<IDGeoRecord> getGeoSegmentWriter(Set<IDGeoRecord> dataToFlush)
throws IOException {
String fileName = indexName + "." + config.getGeoFileExtension();
return new GeoSegmentWriter<IDGeoRecord>(
dataToFlush, directory, fileName, buildGeoSegmentInfo(indexName), geoRecordSerializer);
}
@Override
public void flush(Directory directory, String segmentName) throws IOException {
Set<CartesianGeoRecord> treeToFlush;
Set<String> fieldNamesToFlush;
synchronized (treeLock) {
fieldNamesToFlush = fieldNames;
fieldNames = new HashSet<String>();
treeToFlush = fieldTree;
fieldTree = geoUtil.getBinaryTreeOrderedByBitMag();
}
GeoSegmentWriter<CartesianGeoRecord> geoRecordBTree = null;
GeoSegmentInfo geoSegmentInfo = buildGeoSegmentInfo(fieldNamesToFlush, segmentName);
boolean success = false;
try {
String fileName = config.getGeoFileName(segmentName);
geoRecordBTree =
new GeoSegmentWriter<CartesianGeoRecord>(
treeToFlush, directory, fileName, geoSegmentInfo, geoRecordSerializer);
success = true;
} finally {
// see https://issues.apache.org/jira/browse/LUCENE-3405
if (success) {
IOUtils.close(geoRecordBTree);
} else {
IOUtils.closeWhileHandlingException(geoRecordBTree);
}
}
}
@SuppressWarnings("unchecked")
@Before
public void setUp() {
directory = context.mock(Directory.class);
mockOutput = context.mock(IndexOutput.class);
fieldNameFilterConverter = context.mock(IFieldNameFilterConverter.class);
geoRecordSerializer = context.mock(IGeoRecordSerializer.class);
config.setGeoFileExtension("gto");
treeSet =
new TreeSet<GeoRecord>(
new Comparator<GeoRecord>() {
@Override
public int compare(GeoRecord o1, GeoRecord o2) {
return o1.lowOrder - o2.lowOrder;
}
});
String segmentName = "01";
info = new GeoSegmentInfo();
info.setFieldNameFilterConverter(fieldNameFilterConverter);
info.setGeoVersion(GeoVersion.CURRENT_VERSION);
info.setSegmentName(segmentName);
}
/**
* Adds a record to the geo only index
*
* @param uuid
* @param field
*/
public void index(byte[] uuid, GeoCoordinateField field) {
if (uuid.length != config.getBytesForId()) {
throw new IllegalArgumentException(
"invalid uuid length: "
+ uuid.length
+ ". Expected uuid to be of length "
+ config.getBytesForId()
+ ".");
}
IGeoConverter converter = config.getGeoConverter();
GeoCoordinate geoCoordinate = field.getGeoCoordinate();
IDGeoRecord geoRecord =
converter.toIDGeoRecord(geoCoordinate.getLatitude(), geoCoordinate.getLongitude(), uuid);
newRecords.add(geoRecord);
}
void buildGeoIndexWriter() throws CorruptIndexException, LockObtainFailedException, IOException {
geoComparator = new CartesianGeoRecordComparator();
geoRecordSerializer = new CartesianGeoRecordSerializer();
directory = new RAMDirectory();
config =
new IndexWriterConfig(Version.LUCENE_CURRENT, new StandardAnalyzer(Version.LUCENE_CURRENT));
config.setMergePolicy(new MergeOnOptimizeOnly());
geoConfig = getGeoSearchConfig();
geoConfig.addFieldBitMask(LOCATION_FIELD, LOCATION_BIT_MASK);
geoConfig.addFieldBitMask(IMAGE_LOCATION_FIELD, IMAGE_LOCATION_BIT_MASK);
writer = new GeoIndexWriter(directory, config, geoConfig);
}
private GeoSegmentInfo buildGeoSegmentInfo(String segmentName) throws IOException {
IGeoConverter converter = config.getGeoConverter();
// write version
GeoSegmentInfo info = new GeoSegmentInfo();
info.setGeoVersion(GeoVersion.CURRENT_GEOONLY_VERSION);
info.setSegmentName(segmentName);
info.setBytesPerRecord(IDGeoRecordSerializer.INTERLACE_BYTES + config.getBytesForId());
// now write field -> filterByte mapping info
IFieldNameFilterConverter fieldNameFilterConverter = converter.makeFieldNameFilterConverter();
if (fieldNameFilterConverter != null) {
info.setFieldNameFilterConverter(fieldNameFilterConverter);
}
return info;
}
private void doCreateAndTest(final int docsToAdd, final int version) throws IOException {
final byte[] byteBuf = new byte[10];
final Class<?> clazz = byteBuf.getClass();
context.assertIsSatisfied(); // we should have no calls to mock Objects before we flush
context.checking(
new Expectations() {
{
ignoring(mockOutput).getFilePointer();
will(returnValue(1L));
// get output
one(directory).createOutput(info.getSegmentName() + "." + config.getGeoFileExtension());
will(returnValue(mockOutput));
inSequence(outputSequence);
// write file header
one(mockOutput).writeVInt(version);
inSequence(outputSequence);
one(mockOutput).writeInt(0);
inSequence(outputSequence);
one(mockOutput).writeVInt(docsToAdd);
inSequence(outputSequence);
if (version > GeoVersion.VERSION_0) {
one(mockOutput).writeVInt(GeoSegmentInfo.BYTES_PER_RECORD_V1);
inSequence(outputSequence);
}
one(fieldNameFilterConverter).writeToOutput(mockOutput);
inSequence(outputSequence);
one(mockOutput).seek(1);
inSequence(outputSequence);
one(mockOutput).writeInt(1);
inSequence(outputSequence);
// fill zeroes
one(mockOutput).length();
will(returnValue(7L));
inSequence(outputSequence);
one(mockOutput).seek(with(any(Long.class)));
inSequence(outputSequence);
one(mockOutput)
.writeBytes(
with(any(byteBuf.getClass())),
with(any(Integer.TYPE)),
with(any(Integer.TYPE)));
inSequence(outputSequence);
one(mockOutput).seek(with(any(Long.class)));
inSequence(outputSequence);
one(mockOutput).length();
will(returnValue((long) (7 + 13 * docsToAdd)));
inSequence(outputSequence);
// write actual tree
for (int i = 0; i < docsToAdd; i++) {
one(mockOutput).seek(with(any(Long.class)));
inSequence(outputSequence);
one(geoRecordSerializer)
.writeGeoRecord(
with(mockOutput), with(any(GeoRecord.class)), with(any(Integer.class)));
inSequence(outputSequence);
}
// close
one(mockOutput).close();
}
});
info.setGeoVersion(version);
String fileName = config.getGeoFileName(info.getSegmentName());
GeoSegmentWriter<GeoRecord> bTree =
new GeoSegmentWriter<GeoRecord>(treeSet, directory, fileName, info, geoRecordSerializer);
bTree.close();
context.assertIsSatisfied();
}
