Replace KnnQueryVector by KnnFloatVectorQuery (#767) (#768)

Signed-off-by: Martin Gaievski <gaievski@amazon.com> (cherry picked from commit 71d7232) Co-authored-by: Martin Gaievski <gaievski@amazon.com>
opensearch-project · Feb 17, 2023 · eb88235 · eb88235
1 parent 26c2f3c
commit eb88235
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java b/src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java
@@ -11,7 +11,7 @@
 import lombok.NonNull;
 import lombok.Setter;
 import lombok.extern.log4j.Log4j2;
-import org.apache.lucene.search.KnnVectorQuery;
+import org.apache.lucene.search.KnnFloatVectorQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryShardContext;
@@ -73,13 +73,13 @@ public static Query create(CreateQueryRequest createQueryRequest) {
             );
             try {
                 final Query filterQuery = createQueryRequest.getFilter().get().toQuery(queryShardContext);
-                return new KnnVectorQuery(fieldName, vector, k, filterQuery);
+                return new KnnFloatVectorQuery(fieldName, vector, k, filterQuery);
             } catch (IOException e) {
                 throw new RuntimeException("Cannot create knn query with filter", e);
             }
         }
         log.debug(String.format("Creating Lucene k-NN query for index: %s \"\", field: %s \"\", k: %d", indexName, fieldName, k));
-        return new KnnVectorQuery(fieldName, vector, k);
+        return new KnnFloatVectorQuery(fieldName, vector, k);
     }
 
     /**

diff --git a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java
@@ -6,7 +6,7 @@
 package org.opensearch.knn.index.query;
 
 import com.google.common.collect.ImmutableMap;
-import org.apache.lucene.search.KnnVectorQuery;
+import org.apache.lucene.search.KnnFloatVectorQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.Version;
 import org.opensearch.cluster.ClusterModule;
@@ -196,7 +196,7 @@ public void testDoToQuery_KnnQueryWithFilter() throws Exception {
         when(mockQueryShardContext.fieldMapper(anyString())).thenReturn(mockKNNVectorField);
         Query query = knnQueryBuilder.doToQuery(mockQueryShardContext);
         assertNotNull(query);
-        assertTrue(query instanceof KnnVectorQuery);
+        assertTrue(query.getClass().isAssignableFrom(KnnFloatVectorQuery.class));
     }
 
     public void testDoToQuery_FromModel() {

diff --git a/src/test/java/org/opensearch/knn/index/query/KNNQueryFactoryTests.java b/src/test/java/org/opensearch/knn/index/query/KNNQueryFactoryTests.java
@@ -5,7 +5,7 @@
 
 package org.opensearch.knn.index.query;
 
-import org.apache.lucene.search.KnnVectorQuery;
+import org.apache.lucene.search.KnnFloatVectorQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.query.QueryBuilder;
@@ -47,7 +47,7 @@ public void testCreateLuceneDefaultQuery() {
             .collect(Collectors.toList());
         for (KNNEngine knnEngine : luceneDefaultQueryEngineList) {
             Query query = KNNQueryFactory.create(knnEngine, testIndexName, testFieldName, testQueryVector, testK);
-            assertTrue(query instanceof KnnVectorQuery);
+            assertTrue(query.getClass().isAssignableFrom(KnnFloatVectorQuery.class));
         }
     }
 
@@ -70,7 +70,7 @@ public void testCreateLuceneQueryWithFilter() {
                 .filter(filter)
                 .build();
             Query query = KNNQueryFactory.create(createQueryRequest);
-            assertTrue(query instanceof KnnVectorQuery);
+            assertTrue(query.getClass().isAssignableFrom(KnnFloatVectorQuery.class));
         }
     }
 }