Make it possible to extend Patience/Seeded knn queries (apache#14838)

tteofili · tteofili · commit 52597fa4cacb · 2025-06-25T11:00:50.000+02:00
(cherry picked from commit 2b47cd3)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
@@ -104,6 +104,8 @@ Bug Fixes
 
 * GITHUB#14755: Fix too many documents collected when only bool-filter condition is present. (Ke Wei)
 
+* GITHUB#14838: Make it possible to extend Patience/Seeded knn queries (Tommaso Teofili)
+
 Build
 ---------------------
 * Upgrade forbiddenapis to version 3.9.  (Uwe Schindler)
diff --git a/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java b/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java
@@ -37,7 +37,8 @@ public class HnswQueueSaturationCollector extends KnnCollector.Decorator {
   private int previousQueueSize;
   private int currentQueueSize;
 
-  HnswQueueSaturationCollector(KnnCollector delegate, double saturationThreshold, int patience) {
+  public HnswQueueSaturationCollector(
+      KnnCollector delegate, double saturationThreshold, int patience) {
     super(delegate);
     this.delegate = delegate;
     this.previousQueueSize = 0;
diff --git a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java
@@ -48,7 +48,7 @@ public class PatienceKnnVectorQuery extends AbstractKnnVectorQuery {
   /**
    * Construct a new PatienceKnnVectorQuery instance for a float vector field
    *
-   * @param knnQuery the knn query to be seeded
+   * @param knnQuery the knn query to be wrapped
    * @param saturationThreshold the early exit saturation threshold
    * @param patience the patience parameter
    * @return a new PatienceKnnVectorQuery instance
@@ -62,7 +62,7 @@ public static PatienceKnnVectorQuery fromFloatQuery(
   /**
    * Construct a new PatienceKnnVectorQuery instance for a float vector field
    *
-   * @param knnQuery the knn query to be seeded
+   * @param knnQuery the knn query to be wrapped
    * @return a new PatienceKnnVectorQuery instance
    * @lucene.experimental
    */
@@ -74,7 +74,7 @@ public static PatienceKnnVectorQuery fromFloatQuery(KnnFloatVectorQuery knnQuery
   /**
    * Construct a new PatienceKnnVectorQuery instance for a byte vector field
    *
-   * @param knnQuery the knn query to be seeded
+   * @param knnQuery the knn query to be wrapped
    * @param saturationThreshold the early exit saturation threshold
    * @param patience the patience parameter
    * @return a new PatienceKnnVectorQuery instance
@@ -124,13 +124,55 @@ public static PatienceKnnVectorQuery fromSeededQuery(SeededKnnVectorQuery knnQue
   }
 
   PatienceKnnVectorQuery(
-      AbstractKnnVectorQuery knnQuery, double saturationThreshold, int patience) {
-    super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy);
+      AbstractKnnVectorQuery knnQuery,
+      String field,
+      int k,
+      Query filter,
+      KnnSearchStrategy searchStrategy,
+      double saturationThreshold,
+      int patience) {
+    super(field, k, filter, searchStrategy);
     this.delegate = knnQuery;
     this.saturationThreshold = saturationThreshold;
     this.patience = patience;
   }
 
+  public PatienceKnnVectorQuery(
+      SeededKnnVectorQuery knnQuery, double saturationThreshold, int patience) {
+    this(
+        knnQuery,
+        knnQuery.field,
+        knnQuery.k,
+        knnQuery.filter,
+        knnQuery.searchStrategy,
+        saturationThreshold,
+        patience);
+  }
+
+  public PatienceKnnVectorQuery(
+      KnnFloatVectorQuery knnQuery, double saturationThreshold, int patience) {
+    this(
+        knnQuery,
+        knnQuery.field,
+        knnQuery.k,
+        knnQuery.filter,
+        knnQuery.searchStrategy,
+        saturationThreshold,
+        patience);
+  }
+
+  public PatienceKnnVectorQuery(
+      KnnByteVectorQuery knnQuery, double saturationThreshold, int patience) {
+    this(
+        knnQuery,
+        knnQuery.field,
+        knnQuery.k,
+        knnQuery.filter,
+        knnQuery.searchStrategy,
+        saturationThreshold,
+        patience);
+  }
+
   private static int defaultPatience(AbstractKnnVectorQuery delegate) {
     return Math.max(7, (int) (delegate.k * 0.3));
   }
@@ -243,7 +285,11 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
           new SeededKnnVectorQuery(
               seededKnnVectorQuery.delegate,
               seededKnnVectorQuery.seed,
-              seededKnnVectorQuery.createSeedWeight(indexSearcher));
+              seededKnnVectorQuery.createSeedWeight(indexSearcher),
+              delegate.field,
+              delegate.k,
+              delegate.filter,
+              delegate.searchStrategy);
     }
     return super.rewrite(indexSearcher);
   }
diff --git a/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/SeededKnnVectorQuery.java
@@ -69,13 +69,42 @@ public static SeededKnnVectorQuery fromByteQuery(KnnByteVectorQuery knnQuery, Qu
     return new SeededKnnVectorQuery(knnQuery, seed, null);
   }
 
-  SeededKnnVectorQuery(AbstractKnnVectorQuery knnQuery, Query seed, Weight seedWeight) {
-    super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy);
+  SeededKnnVectorQuery(
+      AbstractKnnVectorQuery knnQuery,
+      Query seed,
+      Weight seedWeight,
+      String field,
+      int k,
+      Query filter,
+      KnnSearchStrategy searchStrategy) {
+    super(field, k, filter, searchStrategy);
     this.delegate = knnQuery;
     this.seed = Objects.requireNonNull(seed);
     this.seedWeight = seedWeight;
   }
 
+  public SeededKnnVectorQuery(KnnFloatVectorQuery knnQuery, Query seed, Weight seedWeight) {
+    this(
+        knnQuery,
+        seed,
+        seedWeight,
+        knnQuery.field,
+        knnQuery.k,
+        knnQuery.filter,
+        knnQuery.searchStrategy);
+  }
+
+  public SeededKnnVectorQuery(KnnByteVectorQuery knnQuery, Query seed, Weight seedWeight) {
+    this(
+        knnQuery,
+        seed,
+        seedWeight,
+        knnQuery.field,
+        knnQuery.k,
+        knnQuery.filter,
+        knnQuery.searchStrategy);
+  }
+
   @Override
   public String toString(String field) {
     return "SeededKnnVectorQuery{"
@@ -94,7 +123,14 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
       return super.rewrite(indexSearcher);
     }
     SeededKnnVectorQuery rewritten =
-        new SeededKnnVectorQuery(delegate, seed, createSeedWeight(indexSearcher));
+        new SeededKnnVectorQuery(
+            delegate,
+            seed,
+            createSeedWeight(indexSearcher),
+            delegate.field,
+            delegate.k,
+            delegate.filter,
+            delegate.searchStrategy);
     return rewritten.rewrite(indexSearcher);
   }
 
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestSeededKnnByteVectorQuery.java
@@ -238,7 +238,7 @@ static class AssertingSeededKnnVectorQuery extends SeededKnnVectorQuery {
 
     public AssertingSeededKnnVectorQuery(
         AbstractKnnVectorQuery query, Query seed, Weight seedWeight, AtomicInteger seedCalls) {
-      super(query, seed, seedWeight);
+      super(query, seed, seedWeight, query.field, query.k, query.filter, query.searchStrategy);
       this.seedCalls = seedCalls;
     }
 

Original file line number	Diff line number	Diff line change
`@@ -238,7 +238,7 @@ static class AssertingSeededKnnVectorQuery extends SeededKnnVectorQuery {`
`238`	`238`
`239`	`239`	`public AssertingSeededKnnVectorQuery(`
`240`	`240`	`AbstractKnnVectorQuery query, Query seed, Weight seedWeight, AtomicInteger seedCalls) {`
`241`		`- super(query, seed, seedWeight);`
	`241`	`+ super(query, seed, seedWeight, query.field, query.k, query.filter, query.searchStrategy);`
`242`	`242`	`this.seedCalls = seedCalls;`
`243`	`243`	`}`
`244`	`244`