diff --git a/lucene/benchmark-jmh/src/java/module-info.java b/lucene/benchmark-jmh/src/java/module-info.java index 8090c7554739..bebcd829e56a 100644 --- a/lucene/benchmark-jmh/src/java/module-info.java +++ b/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; diff --git a/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/FunctionScoreWANDMainVsFeatureBenchmark.java b/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/FunctionScoreWANDMainVsFeatureBenchmark.java new file mode 100644 index 000000000000..b281a7d8fe7e --- /dev/null +++ b/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); + } +} diff --git a/lucene/core/src/java/org/apache/lucene/search/DoubleValues.java b/lucene/core/src/java/org/apache/lucene/search/DoubleValues.java index 4ccf8eb9df04..33b6f07e654a 100644 --- a/lucene/core/src/java/org/apache/lucene/search/DoubleValues.java +++ b/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); + } }; } diff --git a/lucene/core/src/java/org/apache/lucene/search/DoubleValuesSource.java b/lucene/core/src/java/org/apache/lucene/search/DoubleValuesSource.java index 6a3e638539ee..a6750cbd9864 100644 --- a/lucene/core/src/java/org/apache/lucene/search/DoubleValuesSource.java +++ b/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,10 +479,13 @@ 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 @@ -466,7 +493,9 @@ 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; + } }; } diff --git a/lucene/queries/src/java/org/apache/lucene/queries/function/FunctionScoreQuery.java b/lucene/queries/src/java/org/apache/lucene/queries/function/FunctionScoreQuery.java index d16749f73868..a2b972c039dd 100644 --- a/lucene/queries/src/java/org/apache/lucene/queries/function/FunctionScoreQuery.java +++ b/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); diff --git a/lucene/queries/src/test/org/apache/lucene/queries/function/TestFunctionScoreQuery.java b/lucene/queries/src/test/org/apache/lucene/queries/function/TestFunctionScoreQuery.java index 7aed6cbe9c11..0cec2c9fd677 100644 --- a/lucene/queries/src/test/org/apache/lucene/queries/function/TestFunctionScoreQuery.java +++ b/lucene/queries/src/test/org/apache/lucene/queries/function/TestFunctionScoreQuery.java @@ -35,6 +35,7 @@ import org.apache.lucene.search.BooleanClause; import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.BoostQuery; +import org.apache.lucene.search.DoubleValues; import org.apache.lucene.search.DoubleValuesSource; import org.apache.lucene.search.Explanation; import org.apache.lucene.search.IndexSearcher; @@ -42,6 +43,8 @@ import org.apache.lucene.search.PhraseQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.ScorerSupplier; import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.Weight; @@ -378,4 +381,211 @@ public void testQueryMatchesCount() throws Exception { } assertEquals(searchCount, weightCount); } + + public void testMaxScoreDelegationWithDocValuesSkipper() throws Exception { + try (Directory dir = newDirectory()) { + IndexWriterConfig conf = newIndexWriterConfig(); + try (IndexWriter indexWriter = new IndexWriter(dir, conf)) { + for (int i = 1; i <= 500; i++) { + Document doc = new Document(); + doc.add(new TextField(TEXT_FIELD, "test doc", Field.Store.NO)); + doc.add(new NumericDocValuesField("val", i * 10)); + indexWriter.addDocument(doc); + } + indexWriter.commit(); + } + + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + DoubleValuesSource valueSource = DoubleValuesSource.fromLongField("val"); + Query q = new FunctionScoreQuery(new TermQuery(new Term(TEXT_FIELD, "test")), valueSource); + + Weight weight = searcher.createWeight(q, ScoreMode.TOP_SCORES, 1f); + LeafReaderContext ctx = reader.leaves().get(0); + var scorerSupplier = weight.scorerSupplier(ctx); + assertNotNull(scorerSupplier); + var scorer = scorerSupplier.get(Long.MAX_VALUE); + assertNotNull(scorer); + + int maxDoc = ctx.reader().maxDoc(); + int shallowEnd = scorer.advanceShallow(0); + assertTrue(shallowEnd >= 0); + + float maxScore = scorer.getMaxScore(maxDoc); + if (ctx.reader().getDocValuesSkipper("val") != null) { + assertFalse(Float.isInfinite(maxScore)); + assertTrue(maxScore >= 5000f); + } else { + assertTrue(Float.isInfinite(maxScore)); + } + } + } + } + + public void testMaxScorePruningTopDocs() throws Exception { + try (Directory dir = newDirectory()) { + IndexWriterConfig conf = newIndexWriterConfig(); + try (IndexWriter indexWriter = new IndexWriter(dir, conf)) { + for (int i = 1; i <= 200; i++) { + Document doc = new Document(); + doc.add(new TextField(TEXT_FIELD, "prune search", Field.Store.NO)); + doc.add(new NumericDocValuesField("score_val", i)); + indexWriter.addDocument(doc); + } + indexWriter.commit(); + } + + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + Query baseQuery = new TermQuery(new Term(TEXT_FIELD, "prune")); + DoubleValuesSource valSource = DoubleValuesSource.fromLongField("score_val"); + Query scriptQuery = new FunctionScoreQuery(baseQuery, valSource); + + TopDocs topDocs = searcher.search(scriptQuery, 5); + assertEquals(5, topDocs.scoreDocs.length); + // Highest numeric values (200, 199, 198, 197, 196) must be returned + assertTrue(topDocs.scoreDocs[0].score >= 196f); + } + } + } + + public void testMaxScoreMonotonicDecreasingFunction() throws Exception { + try (Directory dir = newDirectory()) { + IndexWriterConfig conf = newIndexWriterConfig(); + try (IndexWriter indexWriter = new IndexWriter(dir, conf)) { + for (int i = 1; i <= 100; i++) { + Document doc = new Document(); + doc.add(new TextField(TEXT_FIELD, "decreasing", Field.Store.NO)); + doc.add(new NumericDocValuesField("val", i * 10)); // 10..1000 + indexWriter.addDocument(doc); + } + indexWriter.commit(); + } + + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + Query baseQuery = new TermQuery(new Term(TEXT_FIELD, "decreasing")); + // Monotonically decreasing function: f(x) = 10000.0 / x (Monotonicity.DECREASING) + DoubleValuesSource valSource = + DoubleValuesSource.fromField( + "val", (v) -> 10000.0 / v, DoubleValuesSource.Monotonicity.DECREASING); + Query scriptQuery = new FunctionScoreQuery(baseQuery, valSource); + + LeafReaderContext ctx = reader.leaves().get(0); + Weight weight = scriptQuery.createWeight(searcher, ScoreMode.TOP_SCORES, 1f); + ScorerSupplier supplier = weight.scorerSupplier(ctx); + assertNotNull(supplier); + Scorer scorer = supplier.get(Long.MAX_VALUE); + + int maxDoc = ctx.reader().maxDoc(); + scorer.advanceShallow(0); + float maxScore = scorer.getMaxScore(maxDoc); + + if (ctx.reader().getDocValuesSkipper("val") != null) { + // Minimum raw value (10) produces max score: 10000 / 10 = 1000.0 + assertFalse(Float.isInfinite(maxScore)); + assertTrue("Expected maxScore >= 1000f but got " + maxScore, maxScore >= 1000f); + } + } + } + } + + public void testCustomHomemadeIncreasingAndDecreasingFunction() throws Exception { + try (Directory dir = newDirectory()) { + IndexWriterConfig conf = newIndexWriterConfig(); + try (IndexWriter indexWriter = new IndexWriter(dir, conf)) { + for (int i = 1; i <= 100; i++) { + Document doc = new Document(); + doc.add(new TextField(TEXT_FIELD, "custom test", Field.Store.NO)); + doc.add(new NumericDocValuesField("val", i * 5)); // 5..500 + indexWriter.addDocument(doc); + } + indexWriter.commit(); + } + + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + Query baseQuery = new TermQuery(new Term(TEXT_FIELD, "custom")); + + // Homemade custom DoubleValuesSource class: f(x) = sqrt(x) [Increasing] + DoubleValuesSource customIncreasing = + new DoubleValuesSource() { + @Override + public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) + throws IOException { + DoubleValues in = DoubleValuesSource.fromLongField("val").getValues(ctx, scores); + return new DoubleValues() { + @Override + public double doubleValue() throws IOException { + return Math.sqrt(in.doubleValue()); + } + + @Override + public boolean advanceExact(int doc) throws IOException { + return in.advanceExact(doc); + } + + @Override + public int advanceShallow(int target) throws IOException { + return in.advanceShallow(target); + } + + @Override + public float getMaxScore(int upTo) throws IOException { + float innerMax = in.getMaxScore(upTo); + return Float.isInfinite(innerMax) + ? Float.POSITIVE_INFINITY + : (float) Math.sqrt(innerMax); + } + }; + } + + @Override + public boolean needsScores() { + return false; + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return true; + } + + @Override + public DoubleValuesSource rewrite(IndexSearcher searcher) { + return this; + } + + @Override + public boolean equals(Object o) { + return o == this; + } + + @Override + public int hashCode() { + return System.identityHashCode(this); + } + + @Override + public String toString() { + return "customSqrt(val)"; + } + }; + + Query q = new FunctionScoreQuery(baseQuery, customIncreasing); + LeafReaderContext ctx = reader.leaves().get(0); + Weight weight = q.createWeight(searcher, ScoreMode.TOP_SCORES, 1f); + ScorerSupplier supplier = weight.scorerSupplier(ctx); + assertNotNull(supplier); + Scorer scorer = supplier.get(Long.MAX_VALUE); + scorer.advanceShallow(0); + float maxScore = scorer.getMaxScore(ctx.reader().maxDoc()); + + if (ctx.reader().getDocValuesSkipper("val") != null) { + // Raw max is 500, sqrt(500) ~ 22.36 + assertFalse(Float.isInfinite(maxScore)); + assertTrue(maxScore >= 22.3f); + } + } + } + } }