Changed unit tests and used builder instead of constructor

Signed-off-by: Sahil Buddharaju <[email protected]>
opensearch-project · Jan 2, 2025 · 17d7a24 · 17d7a24
1 parent 84e9c74
commit 17d7a24
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 13 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
@@ -666,19 +666,19 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws I
         if (Objects.nonNull(filter)) {
             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;
+                KNNQueryBuilder rewrittenQueryBuilder = KNNQueryBuilder.builder()
+                        .fieldName(this.fieldName)
+                        .vector(this.vector)
+                        .k(this.k)
+                        .maxDistance(this.maxDistance)
+                        .minScore(this.minScore)
+                        .methodParameters(this.methodParameters)
+                        .filter(rewrittenFilter)
+                        .ignoreUnmapped(this.ignoreUnmapped)
+                        .rescoreContext(this.rescoreContext)
+                        .expandNested(this.expandNested)
+                        .build();
+                return rewrittenQueryBuilder;
             }
         }
         return super.doRewrite(queryShardContext);

diff --git a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java
@@ -1066,11 +1066,18 @@ public void testDoRewrite_whenFilterSet_thenSuccessful() {
             .filter(rewrittenFilter)
             .k(K)
             .build();
+
+        QueryBuilder filterBefore = filter;
+
         // When
         KNNQueryBuilder knnQueryBuilder = KNNQueryBuilder.builder().fieldName(FIELD_NAME).vector(QUERY_VECTOR).filter(filter).k(K).build();
 
         QueryBuilder actual = knnQueryBuilder.rewrite(context);
 
+        QueryBuilder filterAfter = knnQueryBuilder.getFilter();
+
+        assertEquals(filterBefore, filterAfter);
+
         // Then
         assertEquals(expected, actual);
     }