Adding filter to the toString() method of KnnFloatVectorQuery (apache…

…#13990) * Adding filter to toString() of KnnFloatVectorQuery when it's present (addresses apache#13983) * addressing review comments * adding knnbytevectorquery * unit test improvements * tidy * adding changes entry for the bug fix
vigyasharma · Nov 18, 2024 · a991a9a · a991a9a
1 parent 4400d55
commit a991a9a
Show file tree

Hide file tree

Showing 10 changed files with 79 additions and 5 deletions.
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
@@ -106,6 +106,8 @@ Bug Fixes
 * GITHUB#13944: Ensure deterministic order of clauses for `DisjunctionMaxQuery#toString`. (Laurent Jakubina)
 * GITHUB#13841: Improve Tessellatorlogic when two holes share the same vertex with the polygon which was failing
   in valid polygons. (Ignacio Vera)
+* GITHUB#13990: Added filter to the toString() method of Knn[Float|Byte]VectorQuery
+  and DiversifyingChildren[Float|Byte]KnnVectorQuery. (Viswanath Kuchibhotla)
 
 Build
 ---------------------

diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -55,7 +55,7 @@ abstract class AbstractKnnVectorQuery extends Query {
 
   protected final String field;
   protected final int k;
-  private final Query filter;
+  protected final Query filter;
 
   public AbstractKnnVectorQuery(String field, int k, Query filter) {
     this.field = Objects.requireNonNull(field, "field");

diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnByteVectorQuery.java
@@ -111,7 +111,14 @@ VectorScorer createVectorScorer(LeafReaderContext context, FieldInfo fi) throws
 
   @Override
   public String toString(String field) {
-    return getClass().getSimpleName() + ":" + this.field + "[" + target[0] + ",...][" + k + "]";
+    StringBuilder buffer = new StringBuilder();
+    buffer.append(getClass().getSimpleName() + ":");
+    buffer.append(this.field + "[" + target[0] + ",...]");
+    buffer.append("[" + k + "]");
+    if (this.filter != null) {
+      buffer.append("[" + this.filter + "]");
+    }
+    return buffer.toString();
   }
 
   @Override

diff --git a/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/KnnFloatVectorQuery.java
@@ -112,7 +112,14 @@ VectorScorer createVectorScorer(LeafReaderContext context, FieldInfo fi) throws
 
   @Override
   public String toString(String field) {
-    return getClass().getSimpleName() + ":" + this.field + "[" + target[0] + ",...][" + k + "]";
+    StringBuilder buffer = new StringBuilder();
+    buffer.append(getClass().getSimpleName() + ":");
+    buffer.append(this.field + "[" + target[0] + ",...]");
+    buffer.append("[" + k + "]");
+    if (this.filter != null) {
+      buffer.append("[" + this.filter + "]");
+    }
+    return buffer.toString();
   }
 
   @Override

diff --git a/lucene/core/src/test/org/apache/lucene/search/TestKnnByteVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestKnnByteVectorQuery.java
@@ -23,6 +23,7 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.index.Term;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.TestVectorUtil;
@@ -78,6 +79,11 @@ public void testToString() throws IOException {
       assertEquals("KnnByteVectorQuery:field[0,...][10]", query.toString("ignored"));
 
       assertDocScoreQueryToString(query.rewrite(newSearcher(reader)));
+
+      // test with filter
+      Query filter = new TermQuery(new Term("id", "text"));
+      query = getKnnVectorQuery("field", new float[] {0, 1}, 10, filter);
+      assertEquals("KnnByteVectorQuery:field[0,...][10][id:text]", query.toString("ignored"));
     }
   }
 

diff --git a/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java b/lucene/core/src/test/org/apache/lucene/search/TestKnnFloatVectorQuery.java
@@ -35,6 +35,7 @@
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.index.Term;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
@@ -77,6 +78,11 @@ public void testToString() throws IOException {
       assertEquals("KnnFloatVectorQuery:field[0.0,...][10]", query.toString("ignored"));
 
       assertDocScoreQueryToString(query.rewrite(newSearcher(reader)));
+
+      // test with filter
+      Query filter = new TermQuery(new Term("id", "text"));
+      query = getKnnVectorQuery("field", new float[] {0.0f, 1.0f}, 10, filter);
+      assertEquals("KnnFloatVectorQuery:field[0.0,...][10][id:text]", query.toString("ignored"));
     }
   }
 

diff --git a/...e/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java b/...e/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenByteKnnVectorQuery.java
@@ -154,7 +154,14 @@ protected TopDocs approximateSearch(
 
   @Override
   public String toString(String field) {
-    return getClass().getSimpleName() + ":" + this.field + "[" + query[0] + ",...][" + k + "]";
+    StringBuilder buffer = new StringBuilder();
+    buffer.append(getClass().getSimpleName() + ":");
+    buffer.append(this.field + "[" + query[0] + ",...]");
+    buffer.append("[" + k + "]");
+    if (this.filter != null) {
+      buffer.append("[" + this.filter + "]");
+    }
+    return buffer.toString();
   }
 
   @Override

diff --git a/.../join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java b/.../join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloatKnnVectorQuery.java
@@ -153,7 +153,14 @@ protected TopDocs approximateSearch(
 
   @Override
   public String toString(String field) {
-    return getClass().getSimpleName() + ":" + this.field + "[" + query[0] + ",...][" + k + "]";
+    StringBuilder buffer = new StringBuilder();
+    buffer.append(getClass().getSimpleName() + ":");
+    buffer.append(this.field + "[" + query[0] + ",...]");
+    buffer.append("[" + k + "]");
+    if (this.filter != null) {
+      buffer.append("[" + this.filter + "]");
+    }
+    return buffer.toString();
   }
 
   @Override

diff --git a/...ne/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinByteKnnVectorQuery.java b/...ne/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinByteKnnVectorQuery.java
@@ -29,9 +29,11 @@
 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.index.VectorSimilarityFunction;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.store.Directory;
 
 public class TestParentBlockJoinByteKnnVectorQuery extends ParentBlockJoinKnnVectorQueryTestCase {
@@ -81,6 +83,20 @@ d, new IndexWriterConfig().setMergePolicy(newMergePolicy(random(), false)))) {
     }
   }
 
+  public void testToString() {
+    // test without filter
+    Query query = getParentJoinKnnQuery("field", new float[] {0, 1}, null, 10, null);
+    assertEquals(
+        "DiversifyingChildrenByteKnnVectorQuery:field[0,...][10]", query.toString("ignored"));
+
+    // test with filter
+    Query filter = new TermQuery(new Term("id", "text"));
+    query = getParentJoinKnnQuery("field", new float[] {0, 1}, filter, 10, null);
+    assertEquals(
+        "DiversifyingChildrenByteKnnVectorQuery:field[0,...][10][id:text]",
+        query.toString("ignored"));
+  }
+
   private static byte[] fromFloat(float[] queryVector) {
     byte[] query = new byte[queryVector.length];
     for (int i = 0; i < queryVector.length; i++) {

diff --git a/...e/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinFloatKnnVectorQuery.java b/...e/join/src/test/org/apache/lucene/search/join/TestParentBlockJoinFloatKnnVectorQuery.java
@@ -29,9 +29,11 @@
 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.index.VectorSimilarityFunction;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.store.Directory;
 
 public class TestParentBlockJoinFloatKnnVectorQuery extends ParentBlockJoinKnnVectorQueryTestCase {
@@ -110,6 +112,20 @@ d, new IndexWriterConfig().setMergePolicy(newMergePolicy(random(), false)))) {
     }
   }
 
+  public void testToString() {
+    // test without filter
+    Query query = getParentJoinKnnQuery("field", new float[] {0, 1}, null, 10, null);
+    assertEquals(
+        "DiversifyingChildrenFloatKnnVectorQuery:field[0.0,...][10]", query.toString("ignored"));
+
+    // test with filter
+    Query filter = new TermQuery(new Term("id", "text"));
+    query = getParentJoinKnnQuery("field", new float[] {0.0f, 1.0f}, filter, 10, null);
+    assertEquals(
+        "DiversifyingChildrenFloatKnnVectorQuery:field[0.0,...][10][id:text]",
+        query.toString("ignored"));
+  }
+
   @Override
   Field getKnnVectorField(String name, float[] vector) {
     return new KnnFloatVectorField(name, vector);