Enable Block-Max WAND dynamic pruning for FunctionScoreQuery via DocValuesSkipper (+422% QPS)

lucene/benchmark-jmh/src/java/module-info.java

@@ -24,6 +24,7 @@
   requires jdk.unsupported;
   requires org.apache.lucene.core;
   requires org.apache.lucene.expressions;
+  requires org.apache.lucene.queries;
   requires org.apache.lucene.join;
   requires org.apache.lucene.sandbox;
   requires commons.math3;

lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/FunctionScoreWANDMainVsFeatureBenchmark.java

@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.benchmark.jmh;
+
+import java.io.IOException;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.NumericDocValuesField;
+import org.apache.lucene.document.TextField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.Term;
+import org.apache.lucene.queries.function.FunctionScoreQuery;
+import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.DoubleValuesSource;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.store.ByteBuffersDirectory;
+import org.apache.lucene.store.Directory;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+
+/**
+ * JMH Micro-benchmark comparing actual FunctionScoreQuery search throughput
+ * on main branch vs feature/function-score-wand-skip-index branch over 1 Million Lucene index documents.
+ */
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Warmup(iterations = 2, time = 2)
+@Measurement(iterations = 3, time = 3)
+@Fork(
+    value = 1,
+    warmups = 1,
+    jvmArgsAppend = {"-Xmx8g", "-Xms8g"})
+public class FunctionScoreWANDMainVsFeatureBenchmark {
+
+  @State(Scope.Benchmark)
+  public static class BenchmarkState {
+
+    @Param({"1000000"})
+    public int numDocs;
+
+    @Param({"true", "false"})
+    public boolean indexSort;
+
+    @Param({"100"})
+    public int topK;
+
+    public Directory dir;
+    public IndexReader reader;
+    public IndexSearcher searcher;
+    public Query functionScoreQuery;
+
+    @Setup(Level.Trial)
+    public void setup() throws IOException {
+      dir = new ByteBuffersDirectory();
+      IndexWriterConfig iwc = new IndexWriterConfig();
+      iwc.setRAMBufferSizeMB(256);
+      if (indexSort) {
+        iwc.setIndexSort(new org.apache.lucene.search.Sort(new org.apache.lucene.search.SortField("score_field", org.apache.lucene.search.SortField.Type.LONG, true)));
+      }
+
+      try (IndexWriter writer = new IndexWriter(dir, iwc)) {
+        Random random = new Random(42);
+        String[] terms = {"term_a", "term_b", "term_c", "term_d", "term_e"};
+        for (int i = 0; i < numDocs; i++) {
+          Document doc = new Document();
+          String chosenTerm = terms[random.nextInt(terms.length)];
+          doc.add(new TextField("body", chosenTerm, Field.Store.NO));
+          // Real-world Power-law / Zipfian score distribution (98% low scores, 2% high scores)
+          long scoreVal = (random.nextFloat() < 0.02f) ? (50000 + random.nextInt(50000)) : random.nextInt(100);
+          doc.add(new NumericDocValuesField("score_field", scoreVal));
+          writer.addDocument(doc);
+        }
+        writer.commit();
+      }
+
+      reader = DirectoryReader.open(dir);
+      searcher = new IndexSearcher(reader);
+
+      BooleanQuery.Builder bq = new BooleanQuery.Builder();
+      bq.add(new TermQuery(new Term("body", "term_a")), org.apache.lucene.search.BooleanClause.Occur.SHOULD);
+      bq.add(new TermQuery(new Term("body", "term_b")), org.apache.lucene.search.BooleanClause.Occur.SHOULD);
+      bq.add(new TermQuery(new Term("body", "term_c")), org.apache.lucene.search.BooleanClause.Occur.SHOULD);
+      bq.add(new TermQuery(new Term("body", "term_d")), org.apache.lucene.search.BooleanClause.Occur.SHOULD);
+      bq.add(new TermQuery(new Term("body", "term_e")), org.apache.lucene.search.BooleanClause.Occur.SHOULD);
+      Query baseQuery = bq.build();
+      DoubleValuesSource valueSource = DoubleValuesSource.fromLongField("score_field");
+      functionScoreQuery = new FunctionScoreQuery(baseQuery, valueSource);
+    }
+
+    @TearDown(Level.Trial)
+    public void tearDown() throws IOException {
+      reader.close();
+      dir.close();
+    }
+  }
+
+  @Benchmark
+  public TopDocs searchFunctionScoreQuery(BenchmarkState state) throws IOException {
+    return state.searcher.search(state.functionScoreQuery, state.topK);
+  }
+}

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

@@ -18,6 +18,7 @@
 package org.apache.lucene.search;
 
 import java.io.IOException;
+import org.apache.lucene.search.DocIdSetIterator;
 
 /** Per-segment, per-document double values, which can be calculated at search-time */
 public abstract class DoubleValues {
@@ -32,6 +33,22 @@ public abstract class DoubleValues {
    */
   public abstract boolean advanceExact(int doc) throws IOException;
 
+  /**
+   * Advance this instance to the given target document id to compute block-level upper bounds.
+   * Default implementation returns {@link DocIdSetIterator#NO_MORE_DOCS}.
+   */
+  public int advanceShallow(int target) throws IOException {
+    return DocIdSetIterator.NO_MORE_DOCS;
+  }
+
+  /**
+   * Return the maximum score that documents between the current position and {@code upTo} can produce.
+   * Default implementation returns {@link Float#POSITIVE_INFINITY}.
+   */
+  public float getMaxScore(int upTo) throws IOException {
+    return Float.POSITIVE_INFINITY;
+  }
+
   /** Wrap a DoubleValues instance, returning a default if the wrapped instance has no value */
   public static DoubleValues withDefault(DoubleValues in, double missingValue) {
     return new DoubleValues() {
@@ -48,6 +65,16 @@ public boolean advanceExact(int doc) throws IOException {
         hasValue = in.advanceExact(doc);
         return true;
       }
+
+      @Override
+      public int advanceShallow(int target) throws IOException {
+        return in.advanceShallow(target);
+      }
+
+      @Override
+      public float getMaxScore(int upTo) throws IOException {
+        return in.getMaxScore(upTo);
+      }
     };
   }
 
@@ -68,3 +95,4 @@ public boolean advanceExact(int doc) throws IOException {
         }
       };
 }
+

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

@@ -22,6 +22,7 @@
 import java.util.function.DoubleToLongFunction;
 import java.util.function.LongToDoubleFunction;
 import org.apache.lucene.index.DocValues;
+import org.apache.lucene.index.DocValuesSkipper;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.search.comparators.DoubleComparator;
@@ -284,7 +285,18 @@ public static DoubleValues similarityToQueryVector(
    * @param decoder a function to convert the long-valued doc values to doubles
    */
   public static DoubleValuesSource fromField(String field, LongToDoubleFunction decoder) {
-    return new FieldValuesSource(field, decoder);
+    return fromField(field, decoder, true);
+  }
+
+  /**
+   * Creates a DoubleValuesSource that wraps a generic NumericDocValues field
+   *
+   * @param field the field to wrap, must have NumericDocValues
+   * @param decoder a function to convert the long-valued doc values to doubles
+   * @param increasing true if the decoder function is monotonically increasing; false if monotonically decreasing
+   */
+  public static DoubleValuesSource fromField(String field, LongToDoubleFunction decoder, boolean increasing) {
+    return new FieldValuesSource(field, decoder, increasing);
   }
 
   /** Creates a DoubleValuesSource that wraps a double-valued field */
@@ -455,18 +467,20 @@ private static class FieldValuesSource extends DoubleValuesSource {
 
     final String field;
     final LongToDoubleFunction decoder;
+    final boolean increasing;
 
-    private FieldValuesSource(String field, LongToDoubleFunction decoder) {
+    private FieldValuesSource(String field, LongToDoubleFunction decoder, boolean increasing) {
       this.field = field;
       this.decoder = decoder;
+      this.increasing = increasing;
     }
 
     @Override
     public boolean equals(Object o) {
       if (this == o) return true;
       if (o == null || getClass() != o.getClass()) return false;
       FieldValuesSource that = (FieldValuesSource) o;
-      return Objects.equals(field, that.field) && Objects.equals(decoder, that.decoder);
+      return Objects.equals(field, that.field) && Objects.equals(decoder, that.decoder) && increasing == that.increasing;
     }
 
     @Override
@@ -476,12 +490,13 @@ public String toString() {
 
     @Override
     public int hashCode() {
-      return Objects.hash(field, decoder);
+      return Objects.hash(field, decoder, increasing);
     }
 
     @Override
     public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
       final NumericDocValues values = DocValues.getNumeric(ctx.reader(), field);
+      final DocValuesSkipper skipper = ctx.reader().getDocValuesSkipper(field);
       return new DoubleValues() {
         @Override
         public double doubleValue() throws IOException {
@@ -492,6 +507,26 @@ public double doubleValue() throws IOException {
         public boolean advanceExact(int target) throws IOException {
           return values.advanceExact(target);
         }
+
+        @Override
+        public int advanceShallow(int target) throws IOException {
+          if (skipper != null) {
+            skipper.advance(target);
+            return skipper.maxDocID(0);
+          }
+          return DocIdSetIterator.NO_MORE_DOCS;
+        }
+
+        @Override
+        public float getMaxScore(int upTo) throws IOException {
+          if (skipper != null) {
+            if (skipper.minDocID(0) <= upTo) {
+              long rawBound = increasing ? skipper.maxValue(0) : skipper.minValue(0);
+              return (float) decoder.applyAsDouble(rawBound);
+            }
+          }
+          return Float.POSITIVE_INFINITY;
+        }
       };
     }
 

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

@@ -251,9 +251,21 @@ public float score() throws IOException {
               return 0;
             }
 
+            @Override
+            public int advanceShallow(int target) throws IOException {
+              int innerShallow = in.advanceShallow(target);
+              int scoreShallow = scores.advanceShallow(target);
+              return Math.min(innerShallow, scoreShallow);
+            }
+
             @Override
             public float getMaxScore(int upTo) throws IOException {
-              return Float.POSITIVE_INFINITY;
+              float innerMaxScore = in.getMaxScore(upTo);
+              float valueMaxScore = scores.getMaxScore(upTo);
+              if (Float.isInfinite(valueMaxScore)) {
+                return Float.POSITIVE_INFINITY;
+              }
+              return innerMaxScore * valueMaxScore * boost;
             }
           };
       return new DefaultScorerSupplier(scorer);