LUCENE-6513: Add FrequentTermsSpanBooleanQueryRewrite

lucene/queries/src/java/org/apache/lucene/queries/spans/SpanMultiTermQueryWrapper.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import org.apache.lucene.index.Term;
@@ -30,6 +31,7 @@
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.ScoringRewrite;
 import org.apache.lucene.search.TopTermsRewrite;
+import org.apache.lucene.util.PriorityQueue;
 
 /**
  * Wraps any {@link MultiTermQuery} as a {@link SpanQuery}, so it can be nested within other
@@ -60,7 +62,6 @@ public class SpanMultiTermQueryWrapper<Q extends MultiTermQuery> extends SpanQue
    *
    * @param query Query to wrap.
    */
-  @SuppressWarnings({"rawtypes", "unchecked"})
   public SpanMultiTermQueryWrapper(Q query) {
     this.query = Objects.requireNonNull(query);
     this.rewriteMethod = selectRewriteMethod(query);
@@ -252,4 +253,114 @@ public boolean equals(Object obj) {
       return delegate.equals(other.delegate);
     }
   }
+
+  /**
+   * A rewrite method that translates each term into a SpanTermQuery within a {@link SpanOrQuery},
+   * but retains only the most frequent terms so it will not overflow the boolean max clause count.
+   *
+   * <p>Unlike {@link TopTermsSpanBooleanQueryRewrite}, which keeps the top terms by boost, this
+   * method ranks terms by their index statistics &mdash; document frequency, optionally breaking
+   * ties by total term frequency &mdash; and keeps the {@code maxSize} most frequent ones.
+   *
+   * @see #setRewriteMethod
+   */
+  public static final class FrequentTermsSpanBooleanQueryRewrite extends SpanRewriteMethod {
+    /*
+    This is built on ScoringRewrite rather than TopTermsRewrite because TopTermsRewrite caps terms while collecting
+    them per segment, which requires a ranking key that is stable across segments (such as boost). Document and total
+    term frequencies are only per-segment during collection, so instead ScoringRewrite first collects every matching
+    term — aggregating TermStates across all segments — and this rewrite then retains the maxSize most frequent using
+    those aggregated statistics.
+    */
+
+    /** A collected term paired with its {@link TermStates}, used for ranking. */
+    public static final class ScoreTerm {
+      public final Term term;
+      public final TermStates termState;
+
+      public ScoreTerm(Term term, TermStates termState) {
+        this.term = term;
+        this.termState = termState;
+      }
+    }
+
+    /** Ranks terms by document frequency; the least frequent term is dropped first. */
+    public static final Comparator<ScoreTerm> DF_ORDER =
+        Comparator.comparingInt(st -> st.termState.docFreq());
+
+    /** Ranks terms by document frequency, breaking ties by total term frequency. */
+    public static final Comparator<ScoreTerm> DF_THEN_TTF_ORDER =
+        Comparator.comparingInt((ScoreTerm st) -> st.termState.docFreq())
+            .thenComparingLong(st -> st.termState.totalTermFreq());
+
+    private final int maxSize;
+    private final Comparator<ScoreTerm> order;
+    private final ScoringRewrite<PriorityQueue<ScoreTerm>> delegate;
+
+    /** Create a rewrite that keeps at most {@code maxSize} terms, ranked by {@link #DF_ORDER}. */
+    public FrequentTermsSpanBooleanQueryRewrite(int maxSize) {
+      this(maxSize, DF_ORDER);
+    }
+
+    /** Create a rewrite that keeps at most {@code maxSize} terms, ranked by {@code order}. */
+    public FrequentTermsSpanBooleanQueryRewrite(int maxSize, Comparator<ScoreTerm> order) {
+      this.maxSize = maxSize;
+      this.order = order;
+      this.delegate =
+          new ScoringRewrite<>() {
+            @Override
+            protected PriorityQueue<ScoreTerm> getTopLevelBuilder() {
+              return PriorityQueue.usingComparator(maxSize, order);
+            }
+
+            @Override
+            protected Query build(PriorityQueue<ScoreTerm> builder) {
+              SpanQuery[] result = new SpanQuery[builder.size()];
+              for (int pos = 0; pos < result.length; pos++) {
+                ScoreTerm st = builder.pop();
+                result[pos] = new SpanTermQuery(st.term, st.termState);
+              }
+              return new SpanOrQuery(result);
+            }
+
+            @Override
+            protected void checkMaxClauseCount(int count) {
+              // no-op: the priority queue already bounds the number of retained terms
+            }
+
+            @Override
+            protected void addClause(
+                PriorityQueue<ScoreTerm> topLevel,
+                Term term,
+                int docCount,
+                float boost,
+                TermStates states) {
+              topLevel.insertWithOverflow(new ScoreTerm(term, states));
+            }
+          };
+    }
+
+    /** Returns the maximum number of terms retained by this rewrite. */
+    public int getMaxSize() {
+      return maxSize;
+    }
+
+    @Override
+    public SpanQuery rewrite(IndexSearcher indexSearcher, MultiTermQuery query) throws IOException {
+      return (SpanQuery) delegate.rewrite(indexSearcher, query);
+    }
+
+    @Override
+    public int hashCode() {
+      return 31 * (31 + maxSize) + order.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (this == obj) return true;
+      if (obj == null || getClass() != obj.getClass()) return false;
+      final FrequentTermsSpanBooleanQueryRewrite other = (FrequentTermsSpanBooleanQueryRewrite) obj;
+      return maxSize == other.maxSize && order.equals(other.order);
+    }
+  }
 }

lucene/queries/src/test/org/apache/lucene/queries/spans/TestSpanMultiTermQueryWrapper.java

@@ -17,6 +17,10 @@
 package org.apache.lucene.queries.spans;
 
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Field;
 import org.apache.lucene.index.IndexReader;
@@ -28,6 +32,7 @@
 import org.apache.lucene.search.RegexpQuery;
 import org.apache.lucene.search.WildcardQuery;
 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;
 
@@ -169,32 +174,32 @@ public void testNoSuchMultiTermsInOr() throws Exception {
     FuzzyQuery fuzzyNoSuch = new FuzzyQuery(new Term("field", "noSuch"), 1, 0, 1, false);
     SpanQuery spanNoSuch = new SpanMultiTermQueryWrapper<>(fuzzyNoSuch);
     SpanQuery term = new SpanTermQuery(new Term("field", "brown"));
-    SpanOrQuery near = new SpanOrQuery(new SpanQuery[] {term, spanNoSuch});
+    SpanOrQuery near = new SpanOrQuery(term, spanNoSuch);
     assertEquals(1, searcher.count(near));
 
     // flip
-    near = new SpanOrQuery(new SpanQuery[] {spanNoSuch, term});
+    near = new SpanOrQuery(spanNoSuch, term);
     assertEquals(1, searcher.count(near));
 
     WildcardQuery wcNoSuch = new WildcardQuery(new Term("field", "noSuch*"));
     SpanQuery spanWCNoSuch = new SpanMultiTermQueryWrapper<>(wcNoSuch);
-    near = new SpanOrQuery(new SpanQuery[] {term, spanWCNoSuch});
+    near = new SpanOrQuery(term, spanWCNoSuch);
     assertEquals(1, searcher.count(near));
 
     RegexpQuery rgxNoSuch = new RegexpQuery(new Term("field", "noSuch"));
     SpanQuery spanRgxNoSuch = new SpanMultiTermQueryWrapper<>(rgxNoSuch);
-    near = new SpanOrQuery(new SpanQuery[] {term, spanRgxNoSuch});
+    near = new SpanOrQuery(term, spanRgxNoSuch);
     assertEquals(1, searcher.count(near));
 
     PrefixQuery prfxNoSuch = new PrefixQuery(new Term("field", "noSuch"));
     SpanQuery spanPrfxNoSuch = new SpanMultiTermQueryWrapper<>(prfxNoSuch);
-    near = new SpanOrQuery(new SpanQuery[] {term, spanPrfxNoSuch});
+    near = new SpanOrQuery(term, spanPrfxNoSuch);
     assertEquals(1, searcher.count(near));
 
-    near = new SpanOrQuery(new SpanQuery[] {spanPrfxNoSuch});
+    near = new SpanOrQuery(spanPrfxNoSuch);
     assertEquals(0, searcher.count(near));
 
-    near = new SpanOrQuery(new SpanQuery[] {spanPrfxNoSuch, spanPrfxNoSuch});
+    near = new SpanOrQuery(spanPrfxNoSuch, spanPrfxNoSuch);
     assertEquals(0, searcher.count(near));
   }
 
@@ -237,4 +242,52 @@ public SpanQuery rewrite(IndexSearcher indexSearcher, MultiTermQuery query)
         });
     assertEquals(pqHash, pq.hashCode());
   }
+
+  /** Test is inspired by the patch submitted in LUCENE-6513 GITHUB#7571. */
+  public void testFrequentTermsRewrite() throws Exception {
+    Directory dir = newDirectory();
+    RandomIndexWriter writer = new RandomIndexWriter(random(), dir, new MockAnalyzer(random()));
+
+    // term8 and term54 occur in every document; termN (odd) accumulates so that lower odd N are
+    // more frequent than higher odd N; termN (even) occurs in a single document only.
+    StringBuilder terms = new StringBuilder(" term1");
+    for (int i = 0; i < 300; i++) {
+      Document doc = new Document();
+      String value = "term8 term54";
+      if (i % 2 == 0) {
+        value += " term" + i;
+      } else {
+        terms.append(" term").append(i);
+        value += terms.toString();
+      }
+      doc.add(newTextField("field", value, Field.Store.NO));
+      writer.addDocument(doc);
+    }
+    IndexReader r = writer.getReader();
+    writer.close();
+    IndexSearcher s = newSearcher(r);
+
+    final int size = 8;
+    MultiTermQuery wildcardQuery = new WildcardQuery(new Term("field", "term*"));
+    SpanMultiTermQueryWrapper.FrequentTermsSpanBooleanQueryRewrite rewrite =
+        new SpanMultiTermQueryWrapper.FrequentTermsSpanBooleanQueryRewrite(
+            size, SpanMultiTermQueryWrapper.FrequentTermsSpanBooleanQueryRewrite.DF_THEN_TTF_ORDER);
+    SpanOrQuery query = (SpanOrQuery) rewrite.rewrite(s, wildcardQuery);
+
+    List<SpanQuery> clauses = Arrays.asList(query.getClauses());
+    assertEquals(size, clauses.size());
+
+    Set<String> actual = new HashSet<>();
+    for (SpanQuery clause : clauses) {
+      actual.add(clause.toString("field"));
+    }
+    Set<String> expected =
+        new HashSet<>(
+            Arrays.asList(
+                "term8", "term54", "term1", "term3", "term5", "term7", "term9", "term11"));
+    assertEquals(expected, actual);
+
+    r.close();
+    dir.close();
+  }
 }