Add BM25 k3 query-term frequency saturation (#15818)

Author: sgup432

lucene/CHANGES.txt

@@ -85,6 +85,8 @@ New Features
 
 * GITHUB#15508: Use native vectorization in Lucene. (Ankur Goel, Shubham Chaudhary, Dawid Weiss)
 
+* GITHUB#15818: Add BM25 k3 query-term frequency saturation to BM25Similarity. (Sagar Upadhyaya)
+
 Improvements
 ---------------------
 

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

@@ -31,6 +31,7 @@
 import java.util.function.Predicate;
 import org.apache.lucene.index.TermStates;
 import org.apache.lucene.search.BooleanClause.Occur;
+import org.apache.lucene.search.similarities.Similarity;
 
 /**
  * A Query that matches documents matching boolean combinations of other queries, e.g. {@link
@@ -424,24 +425,41 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
       }
     }
 
-    // Deduplicate SHOULD clauses by summing up their boosts
+    // Deduplicate SHOULD clauses. Tracks total boost per query and which clauses have
+    // explicit boosts. Uses Similarity.computeQueryTermWeight for unweighted duplicates
+    // (e.g. "a a a"), falls back to linear boost sum when explicit boosts are present
+    // (e.g. "a^2 a^3").
     if (clauseSets.get(Occur.SHOULD).size() > 0 && minimumNumberShouldMatch <= 1) {
-      Map<Query, Double> shouldClauses = new HashMap<>();
+      Map<Query, Double> shouldBoosts = new HashMap<>();
+      Set<Query> boostedClauses = new HashSet<>();
       for (Query query : clauseSets.get(Occur.SHOULD)) {
         double boost = 1;
+        boolean hasBoosted = false;
         while (query instanceof BoostQuery) {
           BoostQuery bq = (BoostQuery) query;
           boost *= bq.getBoost();
           query = bq.getQuery();
+          hasBoosted = true;
+        }
+        shouldBoosts.merge(query, boost, Double::sum);
+        if (hasBoosted) {
+          boostedClauses.add(query);
         }
-        shouldClauses.put(query, shouldClauses.getOrDefault(query, 0d) + boost);
       }
-      if (shouldClauses.size() != clauseSets.get(Occur.SHOULD).size()) {
+      if (shouldBoosts.size() != clauseSets.get(Occur.SHOULD).size()) {
+        Similarity similarity = indexSearcher.getSimilarity();
         BooleanQuery.Builder builder =
             new BooleanQuery.Builder().setMinimumNumberShouldMatch(minimumNumberShouldMatch);
-        for (Map.Entry<Query, Double> entry : shouldClauses.entrySet()) {
+        for (Map.Entry<Query, Double> entry : shouldBoosts.entrySet()) {
           Query query = entry.getKey();
-          float boost = entry.getValue().floatValue();
+          double boostSum = entry.getValue();
+          // Only apply similarity-based query term weighting when all occurrences have the
+          // default boost of 1.0 (the common case from query parsers). When explicit boosts
+          // are present (e.g. programmatic BoostQuery), preserve the original linear sum.
+          float boost =
+              boostedClauses.contains(query)
+                  ? (float) boostSum
+                  : similarity.computeQueryTermWeight(Math.toIntExact(Math.round(boostSum)));
           if (boost != 1f) {
             query = new BoostQuery(query, boost);
           }
@@ -456,24 +474,38 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
       }
     }
 
-    // Deduplicate MUST clauses by summing up their boosts
+    // Deduplicate MUST clauses — same approach as SHOULD above.
     if (clauseSets.get(Occur.MUST).size() > 0) {
-      Map<Query, Double> mustClauses = new HashMap<>();
+      Map<Query, Double> mustBoosts = new HashMap<>();
+      Set<Query> boostedClauses = new HashSet<>();
       for (Query query : clauseSets.get(Occur.MUST)) {
         double boost = 1;
+        boolean hasBoosted = false;
         while (query instanceof BoostQuery) {
           BoostQuery bq = (BoostQuery) query;
           boost *= bq.getBoost();
           query = bq.getQuery();
+          hasBoosted = true;
+        }
+        mustBoosts.merge(query, boost, Double::sum);
+        if (hasBoosted) {
+          boostedClauses.add(query);
         }
-        mustClauses.put(query, mustClauses.getOrDefault(query, 0d) + boost);
       }
-      if (mustClauses.size() != clauseSets.get(Occur.MUST).size()) {
+      if (mustBoosts.size() != clauseSets.get(Occur.MUST).size()) {
+        Similarity similarity = indexSearcher.getSimilarity();
         BooleanQuery.Builder builder =
             new BooleanQuery.Builder().setMinimumNumberShouldMatch(minimumNumberShouldMatch);
-        for (Map.Entry<Query, Double> entry : mustClauses.entrySet()) {
+        for (Map.Entry<Query, Double> entry : mustBoosts.entrySet()) {
           Query query = entry.getKey();
-          float boost = entry.getValue().floatValue();
+          // Only apply similarity-based query term weighting when all occurrences have the
+          // default boost of 1.0 (the common case from query parsers). When explicit boosts
+          // are present (e.g. programmatic BoostQuery), preserve the original linear sum.
+          double boostSum = entry.getValue();
+          float boost =
+              boostedClauses.contains(query)
+                  ? (float) boostSum
+                  : similarity.computeQueryTermWeight(Math.toIntExact(Math.round(boostSum)));
           if (boost != 1f) {
             query = new BoostQuery(query, boost);
           }

lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java

@@ -32,6 +32,7 @@
 public class BM25Similarity extends Similarity {
   private final float k1;
   private final float b;
+  private final float k3;
 
   /**
    * BM25 with the supplied parameter values.
@@ -40,10 +41,14 @@ public class BM25Similarity extends Similarity {
    * @param b Controls to what degree document length normalizes tf values.
    * @param discountOverlaps True if overlap tokens (tokens with a position of increment of zero)
    *     are discounted from the document's length.
-   * @throws IllegalArgumentException if {@code k1} is infinite or negative, or if {@code b} is not
-   *     within the range {@code [0..1]}
+   * @param k3 Controls query-side term frequency saturation. When duplicate terms appear in a
+   *     query, their boost is computed as ((k3 + 1) * qtf) / (k3 + qtf) instead of linear summing.
+   *     A negative value disables saturation (linear behavior). Common values are 7 or 8.
+   * @throws IllegalArgumentException if {@code k1} is infinite or negative
+   * @throws IllegalArgumentException if {@code b} is not within the range {@code [0..1]}
+   * @throws IllegalArgumentException if {@code k3} is NaN or infinite
    */
-  public BM25Similarity(float k1, float b, boolean discountOverlaps) {
+  public BM25Similarity(float k1, float b, boolean discountOverlaps, float k3) {
     super(discountOverlaps);
     if (Float.isFinite(k1) == false || k1 < 0) {
       throw new IllegalArgumentException(
@@ -52,8 +57,27 @@ public BM25Similarity(float k1, float b, boolean discountOverlaps) {
     if (Float.isNaN(b) || b < 0 || b > 1) {
       throw new IllegalArgumentException("illegal b value: " + b + ", must be between 0 and 1");
     }
+    if (Float.isNaN(k3) || Float.isInfinite(k3)) {
+      throw new IllegalArgumentException(
+          "illegal k3 value: " + k3 + ", must be a finite value (negative to disable)");
+    }
     this.k1 = k1;
     this.b = b;
+    this.k3 = k3;
+  }
+
+  /**
+   * BM25 with the supplied parameter values.
+   *
+   * @param k1 Controls non-linear term frequency normalization (saturation).
+   * @param b Controls to what degree document length normalizes tf values.
+   * @param discountOverlaps True if overlap tokens (tokens with a position of increment of zero)
+   *     are discounted from the document's length.
+   * @throws IllegalArgumentException if {@code k1} is infinite or negative, or if {@code b} is not
+   *     within the range {@code [0..1]}
+   */
+  public BM25Similarity(float k1, float b, boolean discountOverlaps) {
+    this(k1, b, discountOverlaps, -1f);
   }
 
   /**
@@ -65,7 +89,7 @@ public BM25Similarity(float k1, float b, boolean discountOverlaps) {
    *     within the range {@code [0..1]}
    */
   public BM25Similarity(float k1, float b) {
-    this(k1, b, true);
+    this(k1, b, true, -1f);
   }
 
   /**
@@ -82,7 +106,7 @@ public BM25Similarity(float k1, float b) {
    *     are discounted from the document's length.
    */
   public BM25Similarity(boolean discountOverlaps) {
-    this(1.2f, 0.75f, discountOverlaps);
+    this(1.2f, 0.75f, discountOverlaps, -1f);
   }
 
   /**
@@ -95,7 +119,20 @@ public BM25Similarity(boolean discountOverlaps) {
    * </ul>
    */
   public BM25Similarity() {
-    this(1.2f, 0.75f, true);
+    this(1.2f, 0.75f, true, -1f);
+  }
+
+  /**
+   * Computes the query-term weight using BM25's k3 saturation formula: ((k3 + 1) * qtf) / (k3 +
+   * qtf). When k3 is negative (disabled), falls back to linear weighting where the weight equals
+   * the query term frequency.
+   */
+  @Override
+  public float computeQueryTermWeight(int queryTermFrequency) {
+    if (k3 < 0) {
+      return (float) queryTermFrequency;
+    }
+    return ((k3 + 1f) * queryTermFrequency) / (k3 + queryTermFrequency);
   }
 
   /** Implemented as <code>log(1 + (docCount - docFreq + 0.5)/(docFreq + 0.5))</code>. */
@@ -307,7 +344,7 @@ private List<Explanation> explainConstantFactors() {
 
   @Override
   public String toString() {
-    return "BM25(k1=" + k1 + ",b=" + b + ")";
+    return "BM25(k1=" + k1 + ",b=" + b + (k3 >= 0 ? ",k3=" + k3 : "") + ")";
   }
 
   /**
@@ -327,4 +364,12 @@ public final float getK1() {
   public final float getB() {
     return b;
   }
+
+  /**
+   * Returns the <code>k3</code> parameter for query-side term frequency saturation. A negative
+   * value means saturation is disabled.
+   */
+  public final float getK3() {
+    return k3;
+  }
 }

lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java

@@ -162,6 +162,22 @@ public long computeNorm(FieldInvertState state) {
     return SmallFloat.intToByte4(numTerms);
   }
 
+  /**
+   * Computes the weight for a query term based on how many times it appears in the query. This is
+   * used during query rewriting to compute the boost for duplicate query terms.
+   *
+   * <p>The default implementation returns {@code queryTermFrequency} as a float, which preserves
+   * the existing linear boost behavior. Subclasses may override this to apply saturation (e.g.
+   * BM25's k3 parameter).
+   *
+   * @param queryTermFrequency the number of times a term appears in the query
+   * @return the computed weight for this query term frequency
+   * @lucene.experimental
+   */
+  public float computeQueryTermWeight(int queryTermFrequency) {
+    return (float) queryTermFrequency;
+  }
+
   /**
    * Compute any collection-level weight (e.g. IDF, average document length, etc) needed for scoring
    * a query.

lucene/core/src/test/org/apache/lucene/search/TestBooleanRewrites.java

@@ -32,6 +32,7 @@
 import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BooleanClause.Occur;
+import org.apache.lucene.search.similarities.BM25Similarity;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.apache.lucene.tests.search.RandomApproximationQuery;
@@ -1097,6 +1098,75 @@ public void testShouldClausesLessThanOrEqualToMinimumNumberShouldMatch() throws
     assertEquals(MatchNoDocsQuery.INSTANCE, searcher.rewrite(query));
   }
 
+  public void testDeduplicateShouldClausesWithK3Saturation() throws IOException {
+    IndexSearcher searcher = newSearcher(new MultiReader());
+    searcher.setSimilarity(new BM25Similarity(1.2f, 0.75f, true, 8f));
+
+    // 3 duplicate unweighted SHOULD clauses: similarity applies k3 saturation
+    Query query =
+        new BooleanQuery.Builder()
+            .add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD)
+            .add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD)
+            .add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD)
+            .build();
+    float expectedWeight = ((8f + 1f) * 3f) / (8f + 3f); // 27/11 ≈ 2.4545
+    Query expected = new BoostQuery(new TermQuery(new Term("foo", "bar")), expectedWeight);
+    assertEquals(expected, searcher.rewrite(query));
+  }
+
+  public void testDeduplicateMustClausesWithK3Saturation() throws IOException {
+    IndexSearcher searcher = newSearcher(new MultiReader());
+    searcher.setSimilarity(new BM25Similarity(1.2f, 0.75f, true, 7f));
+
+    // 2 duplicate unweighted MUST clauses
+    Query query =
+        new BooleanQuery.Builder()
+            .add(new TermQuery(new Term("foo", "bar")), Occur.MUST)
+            .add(new TermQuery(new Term("foo", "bar")), Occur.MUST)
+            .add(new TermQuery(new Term("foo", "quux")), Occur.MUST)
+            .build();
+    float expectedWeight = ((7f + 1f) * 2f) / (7f + 2f); // 16/9 ≈ 1.7778
+    Query expected =
+        new BooleanQuery.Builder()
+            .add(new BoostQuery(new TermQuery(new Term("foo", "bar")), expectedWeight), Occur.MUST)
+            .add(new TermQuery(new Term("foo", "quux")), Occur.MUST)
+            .build();
+    assertEquals(expected, searcher.rewrite(query));
+  }
+
+  public void testDeduplicateWithExplicitBoostsFallsBackToLinearSum() throws IOException {
+    IndexSearcher searcher = newSearcher(new MultiReader());
+    searcher.setSimilarity(new BM25Similarity(1.2f, 0.75f, true, 8f));
+
+    // Explicit boosts: should fall back to linear sum (3.0), not k3 saturation
+    Query query =
+        new BooleanQuery.Builder()
+            .add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD)
+            .add(new BoostQuery(new TermQuery(new Term("foo", "bar")), 2), Occur.SHOULD)
+            .add(new TermQuery(new Term("foo", "quux")), Occur.SHOULD)
+            .build();
+    Query expected =
+        new BooleanQuery.Builder()
+            .add(new BoostQuery(new TermQuery(new Term("foo", "bar")), 3), Occur.SHOULD)
+            .add(new TermQuery(new Term("foo", "quux")), Occur.SHOULD)
+            .build();
+    assertEquals(expected, searcher.rewrite(query));
+  }
+
+  public void testDefaultSimilarityPreservesLinearDedup() throws IOException {
+    // Default similarity: computeQueryTermWeight returns qtf linearly
+    IndexSearcher searcher = newSearcher(new MultiReader());
+
+    Query query =
+        new BooleanQuery.Builder()
+            .add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD)
+            .add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD)
+            .add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD)
+            .build();
+    Query expected = new BoostQuery(new TermQuery(new Term("foo", "bar")), 3);
+    assertEquals(expected, searcher.rewrite(query));
+  }
+
   // Test query to count number of rewrites for its lifetime.
   private static final class TestRewriteQuery extends Query {
 

lucene/core/src/test/org/apache/lucene/search/similarities/TestBM25Similarity.java

@@ -47,6 +47,46 @@ public void testIllegalK1() {
     assertTrue(expected.getMessage().contains("illegal k1 value"));
   }
 
+  public void testIllegalK3() {
+    IllegalArgumentException expected =
+        expectThrows(
+            IllegalArgumentException.class,
+            () -> {
+              new BM25Similarity(1.2f, 0.75f, true, Float.POSITIVE_INFINITY);
+            });
+    assertTrue(expected.getMessage().contains("illegal k3 value"));
+
+    expected =
+        expectThrows(
+            IllegalArgumentException.class,
+            () -> {
+              new BM25Similarity(1.2f, 0.75f, true, Float.NEGATIVE_INFINITY);
+            });
+    assertTrue(expected.getMessage().contains("illegal k3 value"));
+
+    expected =
+        expectThrows(
+            IllegalArgumentException.class,
+            () -> {
+              new BM25Similarity(1.2f, 0.75f, true, Float.NaN);
+            });
+    assertTrue(expected.getMessage().contains("illegal k3 value"));
+  }
+
+  public void testValidK3() {
+    // negative k3 disables saturation — should be allowed
+    BM25Similarity sim = new BM25Similarity(1.2f, 0.75f, true, -1f);
+    assertEquals(-1f, sim.getK3(), 0f);
+
+    // zero k3 — maximum saturation
+    sim = new BM25Similarity(1.2f, 0.75f, true, 0f);
+    assertEquals(0f, sim.getK3(), 0f);
+
+    // typical k3 value
+    sim = new BM25Similarity(1.2f, 0.75f, true, 8f);
+    assertEquals(8f, sim.getK3(), 0f);
+  }
+
   public void testIllegalB() {
     IllegalArgumentException expected =
         expectThrows(
@@ -126,6 +166,6 @@ protected Similarity getSimilarity(Random random) {
         b = random.nextFloat();
         break;
     }
-    return new BM25Similarity(k1, b);
+    return new BM25Similarity(k1, b, true, random.nextBoolean() ? -1f : random.nextFloat() * 20f);
   }
 }