From 423573759f74645e0f2cf4a092d8e2d51b75b559 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Thu, 17 Feb 2022 08:03:47 +0100 Subject: [PATCH] LUCENE-10415: FunctionScoreQuery and IndexOrDocValuesQuery delegate Weight#count. (#685) These query wrappers do not modify the set of matching documents so they can delegate Weight#count. --- lucene/CHANGES.txt | 6 ++- .../lucene/search/IndexOrDocValuesQuery.java | 9 +++++ .../search/TestIndexOrDocValuesQuery.java | 39 +++++++++++++++++++ .../queries/function/FunctionScoreQuery.java | 5 +++ .../function/TestFunctionScoreQuery.java | 16 ++++++++ 5 files changed, 73 insertions(+), 2 deletions(-) diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index f652d0dc67de..377b2e22ad9f 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -114,6 +114,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 --------------------- @@ -141,7 +143,7 @@ Improvements * LUCENE-10371: Make IndexRearranger able to arrange segment in a determined order. (Patrick Zhai) - + Optimizations --------------------- @@ -574,7 +576,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 index c9338b71ef85..1a12addf488b 100644 --- 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 index 8fb6be0930c5..0277e4a2dd9d 100644 --- 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 index aed12ff4d3fe..2d207690153c 100644 --- 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 index a7751c753def..d74c3823e9d7 100644 --- 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); + } }