SOLR-3120: spatial post-filter

git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/trunk@1242832 13f79535-47bb-0310-9956-ffa450edef68

solr/CHANGES.txt

@@ -222,6 +222,10 @@ New Features
       MaxFieldValueUpdateProcessorFactory
   (hossman, janhoy)
 
+* SOLR-3120: Optional post filtering for spatial queries bbox and geofilt
+  for LatLonType. (yonik)
+
+
 Optimizations
 ----------------------
 

solr/core/src/java/org/apache/solr/schema/LatLonType.java

@@ -22,15 +22,16 @@ import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.AtomicReaderContext;
 import org.apache.lucene.queries.function.FunctionValues;
 import org.apache.lucene.queries.function.ValueSource;
+import org.apache.lucene.queries.function.ValueSourceScorer;
 import org.apache.lucene.queries.function.valuesource.VectorValueSource;
 import org.apache.lucene.search.*;
 import org.apache.lucene.spatial.DistanceUtils;
 import org.apache.lucene.spatial.tier.InvalidGeoException;
 import org.apache.lucene.util.Bits;
 import org.apache.solr.common.SolrException;
+import org.apache.solr.common.params.SolrParams;
 import org.apache.solr.response.TextResponseWriter;
-import org.apache.solr.search.QParser;
+import org.apache.solr.search.*;
-import org.apache.solr.search.SpatialOptions;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -60,7 +61,7 @@ public class LatLonType extends AbstractSubTypeFieldType implements SpatialQuery
     IndexableField[] f = new IndexableField[(field.indexed() ? 2 : 0) + (field.stored() ? 1 : 0)];
     if (field.indexed()) {
       int i = 0;
-      double[] latLon = new double[0];
+      double[] latLon;
       try {
         latLon = DistanceUtils.parseLatitudeLongitude(null, externalVal);
       } catch (InvalidGeoException e) {
@@ -194,8 +195,11 @@ public class LatLonType extends AbstractSubTypeFieldType implements SpatialQuery
     SchemaField latField = subField(options.field, LAT);
     SchemaField lonField = subField(options.field, LONG);
 
+    SpatialDistanceQuery spatial = new SpatialDistanceQuery();
+
+
     if (options.bbox) {
-      BooleanQuery result = new BooleanQuery();  // only used if box==true
+      BooleanQuery result = new BooleanQuery();
 
       Query latRange = latField.getType().getRangeQuery(parser, latField,
                 String.valueOf(latMin),
@@ -225,11 +229,10 @@ public class LatLonType extends AbstractSubTypeFieldType implements SpatialQuery
         result.add(lonRange, BooleanClause.Occur.MUST);
       }
 
-      return result;
+      spatial.bboxQuery = result;
     }
 
 
-    SpatialDistanceQuery spatial = new SpatialDistanceQuery();
     spatial.origField = options.field.getName();
     spatial.latSource = latField.getType().getValueSource(latField, parser);
     spatial.lonSource = lonField.getType().getValueSource(lonField, parser);
@@ -307,10 +310,12 @@ class LatLonValueSource extends VectorValueSource {
 }
 
 
+
+
 ////////////////////////////////////////////////////////////////////////////////////////////
 // TODO: recast as a value source that doesn't have to match all docs
 
-class SpatialDistanceQuery extends Query {
+class SpatialDistanceQuery extends ExtendedQueryBase implements PostFilter {
   String origField;
   ValueSource latSource;
   ValueSource lonSource;
@@ -318,6 +323,7 @@ class SpatialDistanceQuery extends Query {
   boolean lon2;
 
   boolean calcDist;  // actually calculate the distance with haversine
+  Query bboxQuery;
 
   double latCenter;
   double lonCenter;
@@ -327,12 +333,13 @@ class SpatialDistanceQuery extends Query {
 
   @Override
   public Query rewrite(IndexReader reader) throws IOException {
-    return this;
+    return bboxQuery != null ? bboxQuery.rewrite(reader) : this;
   }
 
   @Override
   public void extractTerms(Set terms) {}
 
+
   protected class SpatialWeight extends Weight {
     protected IndexSearcher searcher;
     protected float queryNorm;
@@ -385,7 +392,7 @@ class SpatialDistanceQuery extends Query {
     int doc=-1;
     final FunctionValues latVals;
     final FunctionValues lonVals;
-    final Bits liveDocs;
+    final Bits acceptDocs;
 
 
     final double lonMin, lonMax, lon2Min, lon2Max, latMin, latMax;
@@ -407,7 +414,7 @@ class SpatialDistanceQuery extends Query {
       this.qWeight = qWeight;
       this.reader = readerContext.reader();
       this.maxDoc = reader.maxDoc();
-      this.liveDocs = acceptDocs;
+      this.acceptDocs = acceptDocs;
       latVals = latSource.getValues(weight.latContext, readerContext);
       lonVals = lonSource.getValues(weight.lonContext, readerContext);
 
@@ -485,7 +492,7 @@ class SpatialDistanceQuery extends Query {
         if (doc>=maxDoc) {
           return doc=NO_MORE_DOCS;
         }
-        if (liveDocs != null && !liveDocs.get(doc)) continue;
+        if (acceptDocs != null && !acceptDocs.get(doc)) continue;
         if (!match()) continue;
         return doc;
       }
@@ -524,9 +531,44 @@ class SpatialDistanceQuery extends Query {
     }
   }
 
+  @Override
+  public DelegatingCollector getFilterCollector(IndexSearcher searcher) {
+    try {
+      return new SpatialCollector(new SpatialWeight(searcher));
+    } catch (IOException e) {
+      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, e);
+    }
+  }
+
+
+  class SpatialCollector extends DelegatingCollector {
+    final SpatialWeight weight;
+    SpatialScorer spatialScorer;
+    int maxdoc;
+    
+    public SpatialCollector(SpatialWeight weight) {
+      this.weight = weight;
+    }
+
+    @Override
+    public void collect(int doc) throws IOException {
+      spatialScorer.doc = doc;
+      if (spatialScorer.match()) delegate.collect(doc);
+    }
+
+    @Override
+    public void setNextReader(AtomicReaderContext context) throws IOException {
+      maxdoc = context.reader().maxDoc();
+      spatialScorer = new SpatialScorer(context, null, weight, 1.0f);
+      super.setNextReader(context);
+    }
+  }
+
 
   @Override
   public Weight createWeight(IndexSearcher searcher) throws IOException {
+    // if we were supposed to use bboxQuery, then we should have been rewritten using that query
+    assert bboxQuery == null;
     return new SpatialWeight(searcher);
   }
 
@@ -537,7 +579,7 @@ class SpatialDistanceQuery extends Query {
   {
     float boost = getBoost();
     return (boost!=1.0?"(":"") +
-            "geofilt(latlonSource="+origField +"(" + latSource + "," + lonSource + ")"
+            (calcDist ? "geofilt" : "bbox") + "(latlonSource="+origField +"(" + latSource + "," + lonSource + ")"
             +",latCenter="+latCenter+",lonCenter="+lonCenter
             +",dist=" + dist
             +",latMin=" + latMin + ",latMax="+latMax
@@ -545,6 +587,7 @@ class SpatialDistanceQuery extends Query {
             +",lon2Min=" + lon2Min + ",lon2Max" + lon2Max
             +",calcDist="+calcDist
             +",planetRadius="+planetRadius
+            // + (bboxQuery == null ? "" : ",bboxQuery="+bboxQuery)
             +")"
             + (boost==1.0 ? "" : ")^"+boost);
   }

solr/core/src/java/org/apache/solr/search/DelegatingCollector.java

@@ -28,7 +28,9 @@ import java.io.IOException;
 
 /** A simple delegating collector where one can set the delegate after creation */
 public class DelegatingCollector extends Collector {
-  static int setLastDelegateCount; // for testing purposes only to determine the number of times a delegating collector chain was used
+
+  /* for internal testing purposes only to determine the number of times a delegating collector chain was used */
+  public static int setLastDelegateCount;
 
   protected Collector delegate;
   protected Scorer scorer;
@@ -74,3 +76,4 @@ public class DelegatingCollector extends Collector {
     return delegate.acceptsDocsOutOfOrder();
   }
 }
+

solr/core/src/test/org/apache/solr/search/SpatialFilterTest.java

@@ -149,10 +149,29 @@ public class SpatialFilterTest extends SolrTestCaseJ4 {
     }
 
     String method = exact ? "geofilt" : "bbox";
+    int postFilterCount = DelegatingCollector.setLastDelegateCount;
 
-    assertQ(req("fl", "id", "q","*:*", "rows", "1000", "fq", "{!"+method+" sfield=" +fieldName +"}",
+    // throw in a random into the main query to prevent most cache hits
+    assertQ(req("fl", "id", "q","*:* OR foo_i:" + random.nextInt(100), "rows", "1000", "fq", "{!"+method+" sfield=" +fieldName +"}",
               "pt", pt, "d", String.valueOf(distance)),
               tests);
+    assertEquals(postFilterCount, DelegatingCollector.setLastDelegateCount);    // post filtering shouldn't be used
+    
+    // try uncached
+    assertQ(req("fl", "id", "q","*:* OR foo_i:" + random.nextInt(100), "rows", "1000", "fq", "{!"+method+" sfield=" +fieldName + " cache=false" + "}",
+        "pt", pt, "d", String.valueOf(distance)),
+        tests);
+    assertEquals(postFilterCount, DelegatingCollector.setLastDelegateCount);      // post filtering shouldn't be used
+
+    // try post filtered for fields that support it
+    if (fieldName.endsWith("ll")) {
+
+    assertQ(req("fl", "id", "q","*:* OR foo_i:" + random.nextInt(100)+100, "rows", "1000", "fq", "{!"+method+" sfield=" +fieldName + " cache=false cost=150" + "}",
+        "pt", pt, "d", String.valueOf(distance)),
+        tests);
+    assertEquals(postFilterCount + 1, DelegatingCollector.setLastDelegateCount);      // post filtering shouldn't be used
+
+    }
   }