LUCENE-9373: FunctionMatchQuery: add "matchCost" param

lucene/CHANGES.txt

@@ -162,7 +162,6 @@ Other
 
 API Changes
 ---------------------
-(No changes)
 
 * LUCENE-9437: Lucene's facet module's DocValuesOrdinalsReader.decode method
   is now public, making it easier for applications to decode facet
@@ -195,6 +194,9 @@ Optimizations
 * LUCENE-9395: ConstantValuesSource now shares a single DoubleValues
   instance across all segments (Tony Xu)
 
+* LUCENE-9373: FunctionMatchQuery now accepts a "matchCost" optimization hint.
+  (Maxim Glazkov, David Smiley)
+
 Bug Fixes
 ---------------------
 

lucene/queries/src/java/org/apache/lucene/queries/function/FunctionMatchQuery.java

@@ -44,17 +44,32 @@ import org.apache.lucene.search.Weight;
  */
 public final class FunctionMatchQuery extends Query {
 
+  static final float DEFAULT_MATCH_COST = 100;
+
   private final DoubleValuesSource source;
   private final DoublePredicate filter;
+  private final float matchCost; // not used in equals/hashCode
 
   /**
-   * Create a FunctionMatchQuery
+   * Create a FunctionMatchQuery with default TwoPhaseIterator matchCost -
+   * {@link #DEFAULT_MATCH_COST} = {@value #DEFAULT_MATCH_COST}
    * @param source  a {@link DoubleValuesSource} to use for values
    * @param filter  the predicate to match against
    */
   public FunctionMatchQuery(DoubleValuesSource source, DoublePredicate filter) {
+    this(source, filter, DEFAULT_MATCH_COST);
+  }
+
+  /**
+   * Create a FunctionMatchQuery
+   * @param source     a {@link DoubleValuesSource} to use for values
+   * @param filter     the predicate to match against
+   * @param matchCost  to be returned by {@link TwoPhaseIterator#matchCost()}
+   */
+  public FunctionMatchQuery(DoubleValuesSource source, DoublePredicate filter, float matchCost) {
     this.source = source;
     this.filter = filter;
+    this.matchCost = matchCost;
   }
 
   @Override
@@ -83,7 +98,7 @@ public final class FunctionMatchQuery extends Query {
 
           @Override
           public float matchCost() {
-            return 100; // TODO maybe DoubleValuesSource should have a matchCost?
+            return matchCost; // TODO maybe DoubleValuesSource should have a matchCost?
           }
         };
         return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);

lucene/queries/src/test/org/apache/lucene/queries/function/TestFunctionMatchQuery.java

@@ -18,20 +18,26 @@
 package org.apache.lucene.queries.function;
 
 import java.io.IOException;
+import java.util.function.DoublePredicate;
 
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.DoubleValuesSource;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.QueryUtils;
+import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.TopDocs;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 
+import static org.apache.lucene.queries.function.FunctionMatchQuery.DEFAULT_MATCH_COST;
+
 public class TestFunctionMatchQuery extends FunctionTestSetup {
 
   static IndexReader reader;
   static IndexSearcher searcher;
+  private static final DoubleValuesSource in = DoubleValuesSource.fromFloatField(FLOAT_FIELD);
 
   @BeforeClass
   public static void beforeClass() throws Exception {
@@ -46,7 +52,6 @@ public class TestFunctionMatchQuery extends FunctionTestSetup {
   }
 
   public void testRangeMatching() throws IOException {
-    DoubleValuesSource in = DoubleValuesSource.fromFloatField(FLOAT_FIELD);
     FunctionMatchQuery fmq = new FunctionMatchQuery(in, d -> d >= 2 && d < 4);
     TopDocs docs = searcher.search(fmq, 10);
 
@@ -58,4 +63,23 @@ public class TestFunctionMatchQuery extends FunctionTestSetup {
 
   }
 
+  public void testTwoPhaseIteratorMatchCost() throws IOException {
+    DoublePredicate predicate = d -> true;
+
+    // should use default match cost
+    FunctionMatchQuery fmq = new FunctionMatchQuery(in, predicate);
+    assertEquals(DEFAULT_MATCH_COST, getMatchCost(fmq), 0.1);
+
+    // should use client defined match cost
+    fmq = new FunctionMatchQuery(in, predicate, 200);
+    assertEquals(200, getMatchCost(fmq), 0.1);
+  }
+
+  private static float getMatchCost(FunctionMatchQuery fmq) throws IOException {
+    LeafReaderContext ctx = reader.leaves().get(0);
+    return fmq.createWeight(searcher, ScoreMode.TOP_DOCS, 1)
+      .scorer(ctx)
+      .twoPhaseIterator()
+      .matchCost();
+  }
 }