From f8b5d439bd2390f97ea586af5fac0b2febff551b Mon Sep 17 00:00:00 2001 From: Tommaso Teofili Date: Mon, 23 Jun 2025 18:10:05 +0200 Subject: [PATCH 1/6] Make it possible to extend PatienceKnnQuery --- .../lucene/search/PatienceKnnVectorQuery.java | 26 +++++++++++++++---- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java index ad7332e4ca8a..d205d1508733 100644 --- 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 @@ -123,8 +123,24 @@ public static PatienceKnnVectorQuery fromSeededQuery(SeededKnnVectorQuery knnQue knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery)); } - PatienceKnnVectorQuery( - AbstractKnnVectorQuery knnQuery, double saturationThreshold, int patience) { + public PatienceKnnVectorQuery( + SeededKnnVectorQuery knnQuery, double saturationThreshold, int patience) { + super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy); + this.delegate = knnQuery; + this.saturationThreshold = saturationThreshold; + this.patience = patience; + } + + public PatienceKnnVectorQuery( + KnnFloatVectorQuery knnQuery, double saturationThreshold, int patience) { + super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy); + this.delegate = knnQuery; + this.saturationThreshold = saturationThreshold; + this.patience = patience; + } + + public PatienceKnnVectorQuery( + KnnByteVectorQuery knnQuery, double saturationThreshold, int patience) { super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy); this.delegate = knnQuery; this.saturationThreshold = saturationThreshold; From dea0b894b762a2bb2953504dd54bc29d9767d2c7 Mon Sep 17 00:00:00 2001 From: Tommaso Teofili Date: Tue, 24 Jun 2025 10:49:14 +0200 Subject: [PATCH 2/6] refactoring to use common package private ctor --- .../lucene/search/PatienceKnnVectorQuery.java | 48 ++++++++++++++----- 1 file changed, 37 insertions(+), 11 deletions(-) diff --git a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java index d205d1508733..373d30151319 100644 --- a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java +++ b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java @@ -123,28 +123,54 @@ public static PatienceKnnVectorQuery fromSeededQuery(SeededKnnVectorQuery knnQue knnQuery, DEFAULT_SATURATION_THRESHOLD, defaultPatience(knnQuery)); } - public PatienceKnnVectorQuery( - SeededKnnVectorQuery knnQuery, double saturationThreshold, int patience) { - super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy); + PatienceKnnVectorQuery( + 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) { - super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy); - this.delegate = knnQuery; - this.saturationThreshold = saturationThreshold; - this.patience = patience; + this( + knnQuery, + knnQuery.field, + knnQuery.k, + knnQuery.filter, + knnQuery.searchStrategy, + saturationThreshold, + patience); } public PatienceKnnVectorQuery( KnnByteVectorQuery knnQuery, double saturationThreshold, int patience) { - super(knnQuery.field, knnQuery.k, knnQuery.filter, knnQuery.searchStrategy); - this.delegate = knnQuery; - this.saturationThreshold = saturationThreshold; - this.patience = patience; + this( + knnQuery, + knnQuery.field, + knnQuery.k, + knnQuery.filter, + knnQuery.searchStrategy, + saturationThreshold, + patience); } private static int defaultPatience(AbstractKnnVectorQuery delegate) { From 8b0933fbbc90229c393bbcf59dcb2c6b5545118a Mon Sep 17 00:00:00 2001 From: Tommaso Teofili Date: Tue, 24 Jun 2025 10:51:37 +0200 Subject: [PATCH 3/6] changelog entry --- lucene/CHANGES.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index e453c697790d..2ff471cd7a2a 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -188,6 +188,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 PatienceKnnQuery (Tommaso Teofili) + Build --------------------- * Upgrade forbiddenapis to version 3.9. (Uwe Schindler) From 75d5a7040501f463cd6be3ea21458d65eb2d6c64 Mon Sep 17 00:00:00 2001 From: Tommaso Teofili Date: Tue, 24 Jun 2025 11:14:43 +0200 Subject: [PATCH 4/6] make the patience collector ctor public as well --- .../org/apache/lucene/search/HnswQueueSaturationCollector.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java b/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java index 1d46baba4ff1..68e710238bce 100644 --- 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,7 @@ 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; From aaa94fdbc0e695ca112c540997d89913868ca13f Mon Sep 17 00:00:00 2001 From: Tommaso Teofili Date: Tue, 24 Jun 2025 15:57:09 +0200 Subject: [PATCH 5/6] tidy --- .../org/apache/lucene/search/HnswQueueSaturationCollector.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java b/lucene/core/src/java/org/apache/lucene/search/HnswQueueSaturationCollector.java index 68e710238bce..a67b28689f83 100644 --- 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; - public HnswQueueSaturationCollector(KnnCollector delegate, double saturationThreshold, int patience) { + public HnswQueueSaturationCollector( + KnnCollector delegate, double saturationThreshold, int patience) { super(delegate); this.delegate = delegate; this.previousQueueSize = 0; From ecc37de92b9c78fecfdf5214ba219692c873c4ad Mon Sep 17 00:00:00 2001 From: Tommaso Teofili Date: Tue, 24 Jun 2025 17:52:14 +0200 Subject: [PATCH 6/6] extend same approach to SeededKnnQuery --- lucene/CHANGES.txt | 2 +- .../lucene/search/PatienceKnnVectorQuery.java | 6 ++- .../lucene/search/SeededKnnVectorQuery.java | 42 +++++++++++++++++-- .../search/TestSeededKnnByteVectorQuery.java | 2 +- 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 2ff471cd7a2a..33a9d696a5e8 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -188,7 +188,7 @@ 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 PatienceKnnQuery (Tommaso Teofili) +* GITHUB#14838: Make it possible to extend Patience/Seeded knn queries (Tommaso Teofili) Build --------------------- diff --git a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java index 373d30151319..170a26a714ec 100644 --- a/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java +++ b/lucene/core/src/java/org/apache/lucene/search/PatienceKnnVectorQuery.java @@ -285,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 index 9fcd34e3d99b..2b43b6c0b638 100644 --- 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 index ba4ba46bf3e8..169f09db2f64 100644 --- 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; }