Defer term collection when query term count is unknown

lucene/CHANGES.txt

@@ -390,6 +390,13 @@ Improvements
 
 Optimizations
 ---------------------
+* GITHUB#16222: MultiTermQuery constant-score wrapper now defers term collection to ScorerSupplier#get()
+  for queries with an unknown term count (automaton queries such as wildcard/regexp/prefix/range),
+  instead of scanning the term dictionary while building the ScorerSupplier. This keeps the
+  ScorerSupplier "planning" phase cheap so a parent conjunction can short-circuit (e.g. a sibling
+  clause matching no documents) before a non-seekable scan, such as a leading wildcard, runs.
+  (Tianxiao Wei)
+
 * GITHUB#16176: Restore WANDScorer for TOP_SCORES + minShouldMatch > 1. (Tianxiao Wei)
 
 * GITHUB#16153: Use TernaryLongHeap in UpdateGraphsUtils for faster HNSW graph merging. (Prithvi S)

lucene/core/src/java/org/apache/lucene/search/AbstractMultiTermQueryConstantScoreWrapper.java

@@ -225,40 +225,69 @@ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOExcepti
       final TermsEnum termsEnum = q.getTermsEnum(terms);
       assert termsEnum != null;
 
-      List<TermAndState> collectedTerms = new ArrayList<>();
-      boolean collectResult = collectTerms(fieldDocCount, termsEnum, collectedTerms);
-
       final long cost;
-      if (collectResult) {
-        // Return a null supplier if no query terms were in the segment:
-        if (collectedTerms.isEmpty()) {
-          return null;
-        }
+      final IOLongFunction<WeightOrDocIdSetIterator> weightOrIteratorSupplier;
+
+      // Only collect terms while building the ScorerSupplier when the query exposes a known,
+      // bounded term count (e.g. TermInSetQuery, getTermsCount() >= 0). There, collecting is
+      // cheap and lets us return a null supplier up-front so a parent BooleanQuery can
+      // short-circuit.
+      //
+      // For queries with an unknown term count (e.g. automaton queries: wildcard / regexp /
+      // prefix / range), collecting eagerly can scan the whole term dictionary during
+      // ScorerSupplier construction -- a leading wildcard such as "*foo*" cannot seek and must
+      // visit every term. That is supposed to be the cheap "planning" phase, and doing it there
+      // defeats a parent conjunction's ability to short-circuit (a sibling clause matching no
+      // documents can no longer skip this clause before the scan runs). So for an unknown term
+      // count we estimate the cost and defer term collection to ScorerSupplier#get().
+      if (q.getTermsCount() >= 0) {
+        List<TermAndState> collectedTerms = new ArrayList<>();
+        boolean collectResult = collectTerms(fieldDocCount, termsEnum, collectedTerms);
+        if (collectResult) {
+          // Return a null supplier if no query terms were in the segment:
+          if (collectedTerms.isEmpty()) {
+            return null;
+          }
 
-        // TODO: Instead of replicating the cost logic of a BooleanQuery we could consider rewriting
-        // to a BQ eagerly at this point and delegating to its cost method (instead of lazily
-        // rewriting on #get). Not sure what the performance hit would be of doing this though.
-        long sumTermCost = 0;
-        for (TermAndState collectedTerm : collectedTerms) {
-          sumTermCost += collectedTerm.docFreq;
+          // TODO: Instead of replicating the cost logic of a BooleanQuery we could consider
+          // rewriting to a BQ eagerly at this point and delegating to its cost method (instead of
+          // lazily rewriting on #get). Not sure what the performance hit would be of doing this
+          // though.
+          long sumTermCost = 0;
+          for (TermAndState collectedTerm : collectedTerms) {
+            sumTermCost += collectedTerm.docFreq;
+          }
+          cost = sumTermCost;
+        } else {
+          cost = estimateCost(terms, q.getTermsCount());
         }
-        cost = sumTermCost;
+        weightOrIteratorSupplier =
+            leadCost -> {
+              if (collectResult) {
+                return rewriteAsBooleanQuery(context, collectedTerms);
+              } else {
+                // Too many terms to rewrite as a simple bq.
+                // Invoke rewriteInner logic to handle rewriting:
+                return rewriteInner(
+                    context, fieldDocCount, terms, termsEnum, collectedTerms, leadCost);
+              }
+            };
       } else {
         cost = estimateCost(terms, q.getTermsCount());
+        weightOrIteratorSupplier =
+            leadCost -> {
+              List<TermAndState> collectedTerms = new ArrayList<>();
+              if (collectTerms(fieldDocCount, termsEnum, collectedTerms)) {
+                return rewriteAsBooleanQuery(context, collectedTerms);
+              } else {
+                // Too many terms to rewrite as a simple bq.
+                // Invoke rewriteInner logic to handle rewriting:
+                return rewriteInner(
+                    context, fieldDocCount, terms, termsEnum, collectedTerms, leadCost);
+              }
+            };
       }
 
-      IOLongFunction<WeightOrDocIdSetIterator> weightOrIteratorSupplier =
-          leadCost -> {
-            if (collectResult) {
-              return rewriteAsBooleanQuery(context, collectedTerms);
-            } else {
-              // Too many terms to rewrite as a simple bq.
-              // Invoke rewriteInner logic to handle rewriting:
-              return rewriteInner(
-                  context, fieldDocCount, terms, termsEnum, collectedTerms, leadCost);
-            }
-          };
-
       return new ScorerSupplier() {
         @Override
         public Scorer get(long leadCost) throws IOException {

lucene/core/src/test/org/apache/lucene/search/TestWildcardQuery.java

@@ -20,19 +20,26 @@
 import static org.hamcrest.Matchers.not;
 
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Field;
 import org.apache.lucene.index.AutomatonTermsEnum;
 import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.FilterDirectoryReader;
+import org.apache.lucene.index.FilterLeafReader;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.MultiTerms;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.index.Terms;
+import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.analysis.MockAnalyzer;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.automaton.CompiledAutomaton;
 
 /** TestWildcardQuery tests the '*' and '?' wildcard characters. */
 public class TestWildcardQuery extends LuceneTestCase {
@@ -438,7 +445,10 @@ public void testCostEstimate() throws IOException {
     Query rewritten = searcher.rewrite(query);
     Weight weight = rewritten.createWeight(searcher, ScoreMode.COMPLETE_NO_SCORES, 1.0f);
     ScorerSupplier supplier = weight.scorerSupplier(lrc);
-    assertEquals(2000, supplier.cost()); // Sum the terms doc freqs
+    // Automaton queries have an unknown term count, so term collection is deferred to get() and the
+    // cost is the worst-case estimate (sum of doc freqs across all terms) rather than the sum over
+    // the matching terms only.
+    assertEquals(3000, supplier.cost());
 
     query = new WildcardQuery(new Term("body", "bar*"));
     rewritten = searcher.rewrite(query);
@@ -449,4 +459,117 @@ public void testCostEstimate() throws IOException {
     reader.close();
     dir.close();
   }
+
+  // A leading wildcard is an automaton MultiTermQuery with an unknown term count (getTermsCount()
+  // == -1). Building its ScorerSupplier must not scan the term dictionary -- that is the cheap
+  // "planning" phase, and a leading wildcard such as "*foo*" cannot seek, so collecting terms
+  // there would walk the whole dictionary. The scan must be deferred to ScorerSupplier#get(), so a
+  // parent conjunction can short-circuit (a sibling clause matching no documents) before it runs.
+  public void testScorerSupplierDoesNotScanTermsEagerly() throws IOException {
+    Directory dir = newDirectory();
+    RandomIndexWriter writer = new RandomIndexWriter(random(), dir);
+    for (int i = 0; i < 1000; i++) {
+      Document doc = new Document();
+      doc.add(newStringField("body", "foo " + i, Field.Store.NO));
+      writer.addDocument(doc);
+    }
+    writer.flush();
+    writer.forceMerge(1);
+    writer.close();
+
+    AtomicInteger termsEnumNextCalls = new AtomicInteger();
+    DirectoryReader reader =
+        new NextCountingReaderWrapper(DirectoryReader.open(dir), termsEnumNextCalls);
+    IndexSearcher searcher = new IndexSearcher(reader);
+    LeafReaderContext lrc = reader.leaves().get(0);
+
+    // Leading wildcard => automaton query, getTermsCount() == -1.
+    WildcardQuery query = new WildcardQuery(new Term("body", "*foo*"));
+    Query rewritten = searcher.rewrite(query);
+    Weight weight = rewritten.createWeight(searcher, ScoreMode.COMPLETE_NO_SCORES, 1.0f);
+
+    termsEnumNextCalls.set(0);
+    ScorerSupplier supplier = weight.scorerSupplier(lrc);
+    assertNotNull(supplier);
+    assertEquals(
+        "scorerSupplier() must not scan the term dictionary for an automaton MultiTermQuery",
+        0,
+        termsEnumNextCalls.get());
+
+    // The scan is deferred to get(): building the scorer is where the terms are actually walked.
+    assertNotNull(supplier.get(Long.MAX_VALUE));
+    assertTrue("get() should scan the term dictionary", termsEnumNextCalls.get() > 0);
+
+    reader.close();
+    dir.close();
+  }
+
+  private static TermsEnum nextCountingTermsEnum(TermsEnum in, AtomicInteger counter) {
+    return new FilterLeafReader.FilterTermsEnum(in) {
+      @Override
+      public BytesRef next() throws IOException {
+        counter.incrementAndGet();
+        return super.next();
+      }
+    };
+  }
+
+  /**
+   * Wraps a reader so every {@link TermsEnum#next()} (via iterator() or intersect()) is counted.
+   */
+  private static class NextCountingReaderWrapper extends FilterDirectoryReader {
+    private final AtomicInteger counter;
+
+    NextCountingReaderWrapper(DirectoryReader in, AtomicInteger counter) throws IOException {
+      super(
+          in,
+          new SubReaderWrapper() {
+            @Override
+            public LeafReader wrap(LeafReader reader) {
+              return new FilterLeafReader(reader) {
+                @Override
+                public Terms terms(String field) {
+                  Terms terms = super.terms(field);
+                  if (terms == null) {
+                    return null;
+                  }
+                  return new FilterTerms(terms) {
+                    @Override
+                    public TermsEnum iterator() throws IOException {
+                      return nextCountingTermsEnum(in.iterator(), counter);
+                    }
+
+                    @Override
+                    public TermsEnum intersect(CompiledAutomaton automaton, BytesRef startTerm)
+                        throws IOException {
+                      return nextCountingTermsEnum(in.intersect(automaton, startTerm), counter);
+                    }
+                  };
+                }
+
+                @Override
+                public CacheHelper getCoreCacheHelper() {
+                  return null;
+                }
+
+                @Override
+                public CacheHelper getReaderCacheHelper() {
+                  return null;
+                }
+              };
+            }
+          });
+      this.counter = counter;
+    }
+
+    @Override
+    protected DirectoryReader doWrapDirectoryReader(DirectoryReader in) throws IOException {
+      return new NextCountingReaderWrapper(in, counter);
+    }
+
+    @Override
+    public CacheHelper getReaderCacheHelper() {
+      return null;
+    }
+  }
 }