LUCENE-7054: add newDistanceQuery to sandbox LatLonPoint

lucene/sandbox/src/java/org/apache/lucene/document/LatLonPoint.java

@@ -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>

lucene/sandbox/src/java/org/apache/lucene/search/PointDistanceQuery.java

@@ -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();
+  }
+}

lucene/sandbox/src/test/org/apache/lucene/document/TestLatLonPoint.java

@@ -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();
+  }
 }

lucene/sandbox/src/test/org/apache/lucene/search/TestLatLonPointQueries.java

@@ -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;
   }