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,150 @@
+/*
+ * 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.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

@@ -32,6 +32,30 @@ 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;
+  }
+
+  /**
+   * Return the minimum score that documents between the current position and {@code upTo} can
+   * produce. Default implementation returns {@link Float#NEGATIVE_INFINITY}.
+   */
+  public float getMinScore(int upTo) throws IOException {
+    return Float.NEGATIVE_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 +72,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);
+      }
     };
   }
 

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;
@@ -46,6 +47,16 @@
  */
 public abstract class DoubleValuesSource implements SegmentCacheable {
 
+  /** Monotonicity direction of a decoder function */
+  public enum Monotonicity {
+    /** Increasing function */
+    INCREASING,
+    /** Decreasing function */
+    DECREASING,
+    /** Non-monotonic or unknown direction function */
+    NONE
+  }
+
   /**
    * Returns a {@link DoubleValues} instance for the passed-in LeafReaderContext and scores
    *
@@ -278,28 +289,41 @@ public static DoubleValues similarityToQueryVector(
   }
 
   /**
-   * Creates a DoubleValuesSource that wraps a generic NumericDocValues field
+   * Creates a DoubleValuesSource that wraps a generic NumericDocValues field. Assumes no monotonic
+   * direction by default (Monotonicity.NONE).
    *
    * @param field the field to wrap, must have NumericDocValues
    * @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, Monotonicity.NONE);
+  }
+
+  /**
+   * 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 monotonicity the monotonicity direction of the decoder function
+   */
+  public static DoubleValuesSource fromField(
+      String field, LongToDoubleFunction decoder, Monotonicity monotonicity) {
+    return new FieldValuesSource(field, decoder, monotonicity);
   }
 
   /** Creates a DoubleValuesSource that wraps a double-valued field */
   public static DoubleValuesSource fromDoubleField(String field) {
-    return fromField(field, Double::longBitsToDouble);
+    return fromField(field, Double::longBitsToDouble, Monotonicity.INCREASING);
   }
 
   /** Creates a DoubleValuesSource that wraps a float-valued field */
   public static DoubleValuesSource fromFloatField(String field) {
-    return fromField(field, (v) -> (double) Float.intBitsToFloat((int) v));
+    return fromField(field, (v) -> (double) Float.intBitsToFloat((int) v), Monotonicity.INCREASING);
   }
 
   /** Creates a DoubleValuesSource that wraps a long-valued field */
   public static DoubleValuesSource fromLongField(String field) {
-    return fromField(field, (v) -> (double) v);
+    return fromField(field, (v) -> (double) v, Monotonicity.INCREASING);
   }
 
   /** Creates a DoubleValuesSource that wraps an int-valued field */
@@ -455,18 +479,23 @@ private static class FieldValuesSource extends DoubleValuesSource {
 
     final String field;
     final LongToDoubleFunction decoder;
+    final Monotonicity monotonicity;
 
-    private FieldValuesSource(String field, LongToDoubleFunction decoder) {
+    private FieldValuesSource(
+        String field, LongToDoubleFunction decoder, Monotonicity monotonicity) {
       this.field = field;
       this.decoder = decoder;
+      this.monotonicity = Objects.requireNonNull(monotonicity);
     }
 
     @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)
+          && monotonicity == that.monotonicity;
     }
 
     @Override
@@ -476,12 +505,13 @@ public String toString() {
 
     @Override
     public int hashCode() {
-      return Objects.hash(field, decoder);
+      return Objects.hash(field, decoder, monotonicity);
     }
 
     @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 +522,43 @@ 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 && monotonicity != Monotonicity.NONE) {
+            if (skipper.minDocID(0) <= upTo) {
+              long rawBound =
+                  (monotonicity == Monotonicity.INCREASING)
+                      ? skipper.maxValue(0)
+                      : skipper.minValue(0);
+              return (float) decoder.applyAsDouble(rawBound);
+            }
+          }
+          return Float.POSITIVE_INFINITY;
+        }
+
+        @Override
+        public float getMinScore(int upTo) throws IOException {
+          if (skipper != null && monotonicity != Monotonicity.NONE) {
+            if (skipper.minDocID(0) <= upTo) {
+              long rawBound =
+                  (monotonicity == Monotonicity.INCREASING)
+                      ? skipper.minValue(0)
+                      : skipper.maxValue(0);
+              return (float) decoder.applyAsDouble(rawBound);
+            }
+          }
+          return Float.NEGATIVE_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);