mirror of https://github.com/apache/lucene.git
LUCENE-7054: add newDistanceQuery to sandbox LatLonPoint
This commit is contained in:
Robert Muir
parent
44324d3dfe
commit
7385d3a4a1
|
|
@ -22,6 +22,7 @@ import org.apache.lucene.util.NumericUtils;
|
|||
import org.apache.lucene.search.BooleanClause;
|
||||
import org.apache.lucene.search.BooleanQuery;
|
||||
import org.apache.lucene.search.ConstantScoreQuery;
|
||||
import org.apache.lucene.search.PointDistanceQuery;
|
||||
import org.apache.lucene.search.PointInPolygonQuery;
|
||||
import org.apache.lucene.search.PointRangeQuery;
|
||||
import org.apache.lucene.search.Query;
|
||||
|
|
@ -37,6 +38,7 @@ import org.apache.lucene.spatial.util.GeoUtils;
|
|||
* This field defines static factory methods for creating common queries:
|
||||
* <ul>
|
||||
* <li>{@link #newBoxQuery newBoxQuery()} for matching points within a bounding box.
|
||||
* <li>{@link #newDistanceQuery newDistanceQuery()} for matching points within a specified distance.
|
||||
* <li>{@link #newPolygonQuery newPolygonQuery()} for matching points within an arbitrary polygon.
|
||||
* </ul>
|
||||
* <p>
|
||||
|
|
@ -206,6 +208,13 @@ public class LatLonPoint extends Field {
|
|||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a query for matching points within the specified distance of the supplied location.
|
||||
*/
|
||||
public static Query newDistanceQuery(String field, double latitude, double longitude, double radiusMeters) {
|
||||
return new PointDistanceQuery(field, latitude, longitude, radiusMeters);
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a query for matching a polygon.
|
||||
* <p>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,181 @@
|
|||
/*
|
||||
* Licensed to the Apache Software Foundation (ASF) under one or more
|
||||
* contributor license agreements. See the NOTICE file distributed with
|
||||
* this work for additional information regarding copyright ownership.
|
||||
* The ASF licenses this file to You under the Apache License, Version 2.0
|
||||
* (the "License"); you may not use this file except in compliance with
|
||||
* the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.apache.lucene.search;
|
||||
|
||||
import java.io.IOException;
|
||||
|
||||
import org.apache.lucene.document.LatLonPoint;
|
||||
import org.apache.lucene.index.LeafReader;
|
||||
import org.apache.lucene.index.LeafReaderContext;
|
||||
import org.apache.lucene.index.PointValues;
|
||||
import org.apache.lucene.index.PointValues.IntersectVisitor;
|
||||
import org.apache.lucene.index.PointValues.Relation;
|
||||
import org.apache.lucene.spatial.util.GeoDistanceUtils;
|
||||
import org.apache.lucene.spatial.util.GeoRect;
|
||||
import org.apache.lucene.spatial.util.GeoUtils;
|
||||
import org.apache.lucene.util.DocIdSetBuilder;
|
||||
import org.apache.lucene.util.NumericUtils;
|
||||
|
||||
/**
|
||||
* Distance query for {@link LatLonPoint}.
|
||||
*/
|
||||
public class PointDistanceQuery extends Query {
|
||||
final String field;
|
||||
final double latitude;
|
||||
final double longitude;
|
||||
final double radiusMeters;
|
||||
|
||||
public PointDistanceQuery(String field, double latitude, double longitude, double radiusMeters) {
|
||||
if (field == null) {
|
||||
throw new IllegalArgumentException("field cannot be null");
|
||||
}
|
||||
if (GeoUtils.isValidLat(latitude) == false) {
|
||||
throw new IllegalArgumentException("latitude: '" + latitude + "' is invalid");
|
||||
}
|
||||
if (GeoUtils.isValidLon(longitude) == false) {
|
||||
throw new IllegalArgumentException("longitude: '" + longitude + "' is invalid");
|
||||
}
|
||||
this.field = field;
|
||||
this.latitude = latitude;
|
||||
this.longitude = longitude;
|
||||
this.radiusMeters = radiusMeters;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException {
|
||||
GeoRect box = GeoUtils.circleToBBox(longitude, latitude, radiusMeters);
|
||||
final GeoRect box1;
|
||||
final GeoRect box2;
|
||||
|
||||
// crosses dateline: split
|
||||
if (box.maxLon < box.minLon) {
|
||||
box1 = new GeoRect(-180.0, box.maxLon, box.minLat, box.maxLat);
|
||||
box2 = new GeoRect(box.minLon, 180.0, box.minLat, box.maxLat);
|
||||
} else {
|
||||
box1 = box;
|
||||
box2 = null;
|
||||
}
|
||||
|
||||
return new ConstantScoreWeight(this) {
|
||||
|
||||
@Override
|
||||
public Scorer scorer(LeafReaderContext context) throws IOException {
|
||||
LeafReader reader = context.reader();
|
||||
PointValues values = reader.getPointValues();
|
||||
if (values == null) {
|
||||
// No docs in this segment had any points fields
|
||||
return null;
|
||||
}
|
||||
|
||||
DocIdSetBuilder result = new DocIdSetBuilder(reader.maxDoc());
|
||||
int[] hitCount = new int[1];
|
||||
values.intersect(field,
|
||||
new IntersectVisitor() {
|
||||
@Override
|
||||
public void visit(int docID) {
|
||||
hitCount[0]++;
|
||||
result.add(docID);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void visit(int docID, byte[] packedValue) {
|
||||
assert packedValue.length == 8;
|
||||
double lat = LatLonPoint.decodeLat(NumericUtils.bytesToInt(packedValue, 0));
|
||||
double lon = LatLonPoint.decodeLon(NumericUtils.bytesToInt(packedValue, Integer.BYTES));
|
||||
if (GeoDistanceUtils.haversin(latitude, longitude, lat, lon) <= radiusMeters) {
|
||||
visit(docID);
|
||||
}
|
||||
}
|
||||
|
||||
// algorithm: we create a bounding box (two bounding boxes if we cross the dateline).
|
||||
// 1. check our bounding box(es) first. if the subtree is entirely outside of those, bail.
|
||||
// 2. see if the subtree is fully contained. if the subtree is enormous along the x axis, wrapping half way around the world, etc: then this can't work, just go to step 3.
|
||||
// 3. recurse naively.
|
||||
@Override
|
||||
public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
|
||||
double latMin = LatLonPoint.decodeLat(NumericUtils.bytesToInt(minPackedValue, 0));
|
||||
double lonMin = LatLonPoint.decodeLon(NumericUtils.bytesToInt(minPackedValue, Integer.BYTES));
|
||||
double latMax = LatLonPoint.decodeLat(NumericUtils.bytesToInt(maxPackedValue, 0));
|
||||
double lonMax = LatLonPoint.decodeLon(NumericUtils.bytesToInt(maxPackedValue, Integer.BYTES));
|
||||
|
||||
if ((latMax < box1.minLat || lonMax < box1.minLon || latMin > box1.maxLat || lonMin > box1.maxLon) &&
|
||||
(box2 == null || latMax < box2.minLat || lonMax < box2.minLon || latMin > box2.maxLat || lonMin > box2.maxLon)) {
|
||||
// we are fully outside of bounding box(es), don't proceed any further.
|
||||
return Relation.CELL_OUTSIDE_QUERY;
|
||||
} else if (lonMax - longitude < 90 && longitude - lonMin < 90 &&
|
||||
GeoDistanceUtils.haversin(latitude, longitude, latMin, lonMin) <= radiusMeters &&
|
||||
GeoDistanceUtils.haversin(latitude, longitude, latMin, lonMax) <= radiusMeters &&
|
||||
GeoDistanceUtils.haversin(latitude, longitude, latMax, lonMin) <= radiusMeters &&
|
||||
GeoDistanceUtils.haversin(latitude, longitude, latMax, lonMax) <= radiusMeters) {
|
||||
// we are fully enclosed, collect everything within this subtree
|
||||
return Relation.CELL_INSIDE_QUERY;
|
||||
} else {
|
||||
// recurse: its inside our bounding box(es), but not fully, or it wraps around.
|
||||
return Relation.CELL_CROSSES_QUERY;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
return new ConstantScoreScorer(this, score(), result.build(hitCount[0]).iterator());
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
@Override
|
||||
public int hashCode() {
|
||||
final int prime = 31;
|
||||
int result = super.hashCode();
|
||||
result = prime * result + field.hashCode();
|
||||
long temp;
|
||||
temp = Double.doubleToLongBits(latitude);
|
||||
result = prime * result + (int) (temp ^ (temp >>> 32));
|
||||
temp = Double.doubleToLongBits(longitude);
|
||||
result = prime * result + (int) (temp ^ (temp >>> 32));
|
||||
temp = Double.doubleToLongBits(radiusMeters);
|
||||
result = prime * result + (int) (temp ^ (temp >>> 32));
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean equals(Object obj) {
|
||||
if (this == obj) return true;
|
||||
if (!super.equals(obj)) return false;
|
||||
if (getClass() != obj.getClass()) return false;
|
||||
PointDistanceQuery other = (PointDistanceQuery) obj;
|
||||
if (!field.equals(other.field)) return false;
|
||||
if (Double.doubleToLongBits(latitude) != Double.doubleToLongBits(other.latitude)) return false;
|
||||
if (Double.doubleToLongBits(longitude) != Double.doubleToLongBits(other.longitude)) return false;
|
||||
if (Double.doubleToLongBits(radiusMeters) != Double.doubleToLongBits(other.radiusMeters)) return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString(String field) {
|
||||
StringBuilder sb = new StringBuilder();
|
||||
if (!this.field.equals(field)) {
|
||||
sb.append(field);
|
||||
sb.append(':');
|
||||
}
|
||||
sb.append(latitude);
|
||||
sb.append(",");
|
||||
sb.append(longitude);
|
||||
sb.append(" +/- ");
|
||||
sb.append(radiusMeters);
|
||||
sb.append(" meters");
|
||||
return sb.toString();
|
||||
}
|
||||
}
|
||||
|
|
@ -16,11 +16,29 @@
|
|||
*/
|
||||
package org.apache.lucene.document;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.BitSet;
|
||||
|
||||
import org.apache.lucene.codecs.FilterCodec;
|
||||
import org.apache.lucene.codecs.PointFormat;
|
||||
import org.apache.lucene.codecs.PointReader;
|
||||
import org.apache.lucene.codecs.PointWriter;
|
||||
import org.apache.lucene.codecs.lucene60.Lucene60PointReader;
|
||||
import org.apache.lucene.codecs.lucene60.Lucene60PointWriter;
|
||||
import org.apache.lucene.index.IndexReader;
|
||||
import org.apache.lucene.index.IndexWriterConfig;
|
||||
import org.apache.lucene.index.RandomIndexWriter;
|
||||
import org.apache.lucene.index.SegmentReadState;
|
||||
import org.apache.lucene.index.SegmentWriteState;
|
||||
import org.apache.lucene.search.IndexSearcher;
|
||||
import org.apache.lucene.search.ScoreDoc;
|
||||
import org.apache.lucene.search.Sort;
|
||||
import org.apache.lucene.search.TopDocs;
|
||||
import org.apache.lucene.spatial.util.GeoDistanceUtils;
|
||||
import org.apache.lucene.store.Directory;
|
||||
import org.apache.lucene.util.LuceneTestCase;
|
||||
import org.apache.lucene.util.TestUtil;
|
||||
import org.apache.lucene.util.bkd.BKDWriter;
|
||||
|
||||
/** Simple tests for {@link LatLonPoint} */
|
||||
public class TestLatLonPoint extends LuceneTestCase {
|
||||
|
|
@ -53,4 +71,92 @@ public class TestLatLonPoint extends LuceneTestCase {
|
|||
// looks crazy due to lossiness
|
||||
assertEquals("field:[17.99999997485429 TO 18.999999999068677},[-65.9999999217689 TO -64.99999998137355}", LatLonPoint.newBoxQuery("field", 18, 19, -66, -65).toString());
|
||||
}
|
||||
|
||||
public void testRadiusRandom() throws Exception {
|
||||
for (int iters = 0; iters < 100; iters++) {
|
||||
doRandomTest(10, 100);
|
||||
}
|
||||
}
|
||||
|
||||
@Nightly
|
||||
public void testRadiusRandomHuge() throws Exception {
|
||||
for (int iters = 0; iters < 10; iters++) {
|
||||
doRandomTest(2000, 100);
|
||||
}
|
||||
}
|
||||
|
||||
private void doRandomTest(int numDocs, int numQueries) throws IOException {
|
||||
Directory dir = newDirectory();
|
||||
IndexWriterConfig iwc = newIndexWriterConfig();
|
||||
int pointsInLeaf = 2 + random().nextInt(4);
|
||||
iwc.setCodec(new FilterCodec("Lucene60", TestUtil.getDefaultCodec()) {
|
||||
@Override
|
||||
public PointFormat pointFormat() {
|
||||
return new PointFormat() {
|
||||
@Override
|
||||
public PointWriter fieldsWriter(SegmentWriteState writeState) throws IOException {
|
||||
return new Lucene60PointWriter(writeState, pointsInLeaf, BKDWriter.DEFAULT_MAX_MB_SORT_IN_HEAP);
|
||||
}
|
||||
|
||||
@Override
|
||||
public PointReader fieldsReader(SegmentReadState readState) throws IOException {
|
||||
return new Lucene60PointReader(readState);
|
||||
}
|
||||
};
|
||||
}
|
||||
});
|
||||
RandomIndexWriter writer = new RandomIndexWriter(random(), dir, iwc);
|
||||
|
||||
for (int i = 0; i < numDocs; i++) {
|
||||
double latRaw = -90 + 180.0 * random().nextDouble();
|
||||
double lonRaw = -180 + 360.0 * random().nextDouble();
|
||||
// pre-normalize up front, so we can just use quantized value for testing and do simple exact comparisons
|
||||
double lat = LatLonPoint.decodeLat(LatLonPoint.encodeLat(latRaw));
|
||||
double lon = LatLonPoint.decodeLon(LatLonPoint.encodeLon(lonRaw));
|
||||
Document doc = new Document();
|
||||
doc.add(new LatLonPoint("field", lat, lon));
|
||||
doc.add(new StoredField("lat", lat));
|
||||
doc.add(new StoredField("lon", lon));
|
||||
writer.addDocument(doc);
|
||||
}
|
||||
IndexReader reader = writer.getReader();
|
||||
IndexSearcher searcher = new IndexSearcher(reader);
|
||||
|
||||
for (int i = 0; i < numQueries; i++) {
|
||||
double lat = -90 + 180.0 * random().nextDouble();
|
||||
double lon = -180 + 360.0 * random().nextDouble();
|
||||
double radius = 50000000 * random().nextDouble();
|
||||
|
||||
BitSet expected = new BitSet();
|
||||
for (int doc = 0; doc < reader.maxDoc(); doc++) {
|
||||
double docLatitude = reader.document(doc).getField("lat").numericValue().doubleValue();
|
||||
double docLongitude = reader.document(doc).getField("lon").numericValue().doubleValue();
|
||||
double distance = GeoDistanceUtils.haversin(lat, lon, docLatitude, docLongitude);
|
||||
if (distance <= radius) {
|
||||
expected.set(doc);
|
||||
}
|
||||
}
|
||||
|
||||
TopDocs topDocs = searcher.search(LatLonPoint.newDistanceQuery("field", lat, lon, radius), reader.maxDoc(), Sort.INDEXORDER);
|
||||
BitSet actual = new BitSet();
|
||||
for (ScoreDoc doc : topDocs.scoreDocs) {
|
||||
actual.set(doc.doc);
|
||||
}
|
||||
|
||||
try {
|
||||
assertEquals(expected, actual);
|
||||
} catch (AssertionError e) {
|
||||
for (int doc = 0; doc < reader.maxDoc(); doc++) {
|
||||
double docLatitude = reader.document(doc).getField("lat").numericValue().doubleValue();
|
||||
double docLongitude = reader.document(doc).getField("lon").numericValue().doubleValue();
|
||||
double distance = GeoDistanceUtils.haversin(lat, lon, docLatitude, docLongitude);
|
||||
System.out.println("" + doc + ": (" + docLatitude + "," + docLongitude + "), distance=" + distance);
|
||||
}
|
||||
throw e;
|
||||
}
|
||||
}
|
||||
reader.close();
|
||||
writer.close();
|
||||
dir.close();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -39,7 +39,8 @@ public class TestLatLonPointQueries extends BaseGeoPointTestCase {
|
|||
|
||||
@Override
|
||||
protected Query newDistanceQuery(String field, double centerLat, double centerLon, double radiusMeters) {
|
||||
// return new BKDDistanceQuery(field, centerLat, centerLon, radiusMeters);
|
||||
// TODO: fix this to be debuggable before enabling!
|
||||
// return LatLonPoint.newDistanceQuery(field, centerLat, centerLon, radiusMeters);
|
||||
return null;
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue