Changed filter logic

buddharajusahil · Dec 30, 2024 · 008919a · 008919a
1 parent d57fdea
commit 008919a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java b/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java
@@ -23,6 +23,7 @@
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.index.query.RangeQueryBuilder;
 import org.opensearch.knn.index.engine.KNNMethodConfigContext;
 import org.opensearch.knn.index.engine.model.QueryContext;
 import org.opensearch.knn.index.engine.qframe.QuantizationConfig;
@@ -660,15 +661,22 @@ public String getWriteableName() {
         return NAME;
     }
 
+
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws IOException {
-        // rewrite filter query if it exists to avoid runtime errors in next steps of query phase
+        QueryBuilder rewrittenFilter;
         if (Objects.nonNull(filter)) {
-            filter = filter.rewrite(queryShardContext);
+            rewrittenFilter = filter.rewrite(queryShardContext);
+            if (rewrittenFilter != filter) {
+                KNNQueryBuilder newKNNQuery = new KNNQueryBuilder(this.fieldName, this.vector, this.k, this.maxDistance, this.minScore,
+                        this.methodParameters, rewrittenFilter, this.ignoreUnmapped, this.rescoreContext, this.expandNested);
+                return newKNNQuery;
+            }
         }
         return super.doRewrite(queryShardContext);
     }
 
+
     @Getter
     @AllArgsConstructor
     private static class QueryConfigFromMapping {

diff --git a/src/test/java/org/opensearch/knn/index/FaissIT.java b/src/test/java/org/opensearch/knn/index/FaissIT.java
@@ -1956,7 +1956,8 @@ public void testQueryWithFilter_whenNonExistingFieldUsedInFilter_thenSuccessful(
         Map<String, Object> mappingMap = xContentBuilderToMap(builder);
         String mapping = builder.toString();
 
-        createKnnIndex(INDEX_NAME, mapping);
+        createIndex(INDEX_NAME, Settings.builder().put("number_of_shards", 2).put("number_of_replicas", 1).put("index.knn", true).build());
+        putMappingRequest(INDEX_NAME, mapping);
 
         Float[] vector = new Float[] { 2.0f, 4.5f, 6.5f };
 

diff --git a/src/test/java/org/opensearch/knn/index/LuceneEngineIT.java b/src/test/java/org/opensearch/knn/index/LuceneEngineIT.java
@@ -15,6 +15,7 @@
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseException;
 import org.opensearch.common.Nullable;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.index.query.QueryBuilder;
@@ -304,7 +305,8 @@ public void testQueryWithFilter_whenNonExistingFieldUsedInFilter_thenSuccessful(
         Map<String, Object> mappingMap = xContentBuilderToMap(builder);
         String mapping = builder.toString();
 
-        createKnnIndex(INDEX_NAME, mapping);
+        createIndex(INDEX_NAME, Settings.builder().put("number_of_shards", 2).put("number_of_replicas", 1).put("index.knn", true).build());
+        putMappingRequest(INDEX_NAME, mapping);
 
         Float[] vector = new Float[] { 2.0f, 4.5f, 6.5f };