LUCENE-10415: FunctionScoreQuery and IndexOrDocValuesQuery delegate W…

…eight#count. (apache#685) These query wrappers do not modify the set of matching documents so they can delegate Weight#count.
SeaseLtd · Mar 10, 2022 · 331a41d · 331a41d
1 parent bd4f094
commit 331a41d
Show file tree

Hide file tree

Showing 5 changed files with 73 additions and 2 deletions.
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
@@ -154,6 +154,8 @@ New Features
   the number of matching range docs when each doc has at-most one point and the points are 1-dimensional.
   (Gautam Worah, Ignacio Vera, Adrien Grand)
 
+* LUCENE-10415: FunctionScoreQuery and IndexOrDocValuesQuery delegate Weight#count. (Ignacio Vera)     
+
 Improvements
 ---------------------
 
@@ -181,7 +183,7 @@ Improvements
 
 * LUCENE-10371: Make IndexRearranger able to arrange segment in a determined order.
   (Patrick Zhai)
-
+  
 Optimizations
 ---------------------
 
@@ -614,7 +616,7 @@ Improvements
   (David Smiley)
 
 * LUCENE-10062: Switch taxonomy faceting to use numeric doc values for storing ordinals instead of binary doc values
-  with its own custom encoding. (Greg Miller)
+  with its own custom encoding. (Greg Miller) 
 
 Bug fixes
 ---------------------

diff --git a/lucene/core/src/java/org/apache/lucene/search/IndexOrDocValuesQuery.java b/lucene/core/src/java/org/apache/lucene/search/IndexOrDocValuesQuery.java
@@ -142,6 +142,15 @@ public BulkScorer bulkScorer(LeafReaderContext context) throws IOException {
         return indexWeight.bulkScorer(context);
       }
 
+      @Override
+      public int count(LeafReaderContext context) throws IOException {
+        final int count = indexWeight.count(context);
+        if (count != -1) {
+          return count;
+        }
+        return dvWeight.count(context);
+      }
+
       @Override
       public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
         final ScorerSupplier indexScorerSupplier = indexWeight.scorerSupplier(context);

diff --git a/lucene/core/src/test/org/apache/lucene/search/TestIndexOrDocValuesQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestIndexOrDocValuesQuery.java
@@ -26,6 +26,7 @@
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BooleanClause.Occur;
 import org.apache.lucene.store.Directory;
@@ -182,4 +183,42 @@ public void testUseIndexForSelectiveMultiValueQueries() throws IOException {
     w.close();
     dir.close();
   }
+
+  // Weight#count is delegated to the inner weight
+  public void testQueryMatchesCount() throws Exception {
+    Directory dir = newDirectory();
+    IndexWriter w =
+        new IndexWriter(
+            dir,
+            newIndexWriterConfig()
+                // relies on costs and PointValues.estimateCost so we need the default codec
+                .setCodec(TestUtil.getDefaultCodec()));
+    final int numDocs = random().nextInt(5000);
+    for (int i = 0; i < numDocs; ++i) {
+      Document doc = new Document();
+      doc.add(new LongPoint("f2", 42L));
+      doc.add(new SortedNumericDocValuesField("f2", 42L));
+      w.addDocument(doc);
+    }
+    w.forceMerge(1);
+    IndexReader reader = DirectoryReader.open(w);
+    IndexSearcher searcher = newSearcher(reader);
+
+    final IndexOrDocValuesQuery query =
+        new IndexOrDocValuesQuery(
+            LongPoint.newExactQuery("f2", 42),
+            SortedNumericDocValuesField.newSlowRangeQuery("f2", 42, 42L));
+
+    final int searchCount = searcher.count(query);
+    final Weight weight = searcher.createWeight(query, ScoreMode.COMPLETE, 1);
+    int weightCount = 0;
+    for (LeafReaderContext leafReaderContext : reader.leaves()) {
+      weightCount += weight.count(leafReaderContext);
+    }
+    assertEquals(searchCount, weightCount);
+
+    reader.close();
+    w.close();
+    dir.close();
+  }
 }
diff --git a/lucene/queries/src/java/org/apache/lucene/queries/function/FunctionScoreQuery.java b/lucene/queries/src/java/org/apache/lucene/queries/function/FunctionScoreQuery.java
@@ -253,6 +253,11 @@ public float getMaxScore(int upTo) throws IOException {
       };
     }
 
+    @Override
+    public int count(LeafReaderContext context) throws IOException {
+      return inner.count(context);
+    }
+
     @Override
     public boolean isCacheable(LeafReaderContext ctx) {
       return inner.isCacheable(ctx) && valueSource.isCacheable(ctx);

diff --git a/lucene/queries/src/test/org/apache/lucene/queries/function/TestFunctionScoreQuery.java b/lucene/queries/src/test/org/apache/lucene/queries/function/TestFunctionScoreQuery.java
@@ -30,6 +30,7 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
@@ -361,4 +362,19 @@ public void testScoreCalledTwice() throws Exception {
       }
     }
   }
+
+  // Weight#count is delegated to the inner weight
+  public void testQueryMatchesCount() throws Exception {
+    TermQuery query = new TermQuery(new Term(TEXT_FIELD, "first"));
+    FunctionScoreQuery fq =
+        FunctionScoreQuery.boostByValue(query, DoubleValuesSource.fromIntField("iii"));
+
+    final int searchCount = searcher.count(fq);
+    final Weight weight = searcher.createWeight(fq, ScoreMode.COMPLETE, 1);
+    int weightCount = 0;
+    for (LeafReaderContext leafReaderContext : reader.leaves()) {
+      weightCount += weight.count(leafReaderContext);
+    }
+    assertEquals(searchCount, weightCount);
+  }
 }