From e179b3ce43a6acfaf1accc5ef370fcc7f25488f3 Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Mon, 30 Dec 2024 14:19:53 -0800 Subject: [PATCH 1/8] Changed filter logic Signed-off-by: Sahil Buddharaju --- .../opensearch/knn/index/query/KNNQueryBuilder.java | 12 ++++++++++-- src/test/java/org/opensearch/knn/index/FaissIT.java | 3 ++- .../org/opensearch/knn/index/LuceneEngineIT.java | 4 +++- 3 files changed, 15 insertions(+), 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 index d2b169b2a..1d706e098 100644 --- 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 index f6aef8cb1..1618ec4d8 100644 --- 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 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 index 688d22e74..25dc6ccd4 100644 --- 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 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 }; From aa00c5c0a426873b6ef821110ce2dfb04724e12f Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Mon, 30 Dec 2024 14:25:46 -0800 Subject: [PATCH 2/8] spotless Signed-off-by: Sahil Buddharaju --- .../knn/index/query/KNNQueryBuilder.java | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 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 index 1d706e098..62400a509 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java @@ -23,7 +23,6 @@ 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; @@ -661,22 +660,30 @@ public String getWriteableName() { return NAME; } - @Override protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws IOException { QueryBuilder rewrittenFilter; 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); + 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 { From bdd753d219bca28c25785b6349e0995c8e9c008b Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Mon, 30 Dec 2024 15:54:18 -0800 Subject: [PATCH 3/8] Changelog Signed-off-by: Sahil Buddharaju --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5fda2a2..a767a63b7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Bug Fixes * Fixing the bug when a segment has no vector field present for disk based vector search (#2282)[https://github.com/opensearch-project/k-NN/pull/2282] * Allow validation for non knn index only after 2.17.0 (#2315)[https://github.com/opensearch-project/k-NN/pull/2315] +* Release query vector memory after execution (#2346)[https://github.com/opensearch-project/k-NN/pull/2346] +* Fix shard level rescoring disabled setting flag (#2352)[https://github.com/opensearch-project/k-NN/pull/2352] +* Fix filter rewrite logic getting inconsistent / incorrect results (#2359)[https://github.com/opensearch-project/k-NN/pull/2359] ### Infrastructure * Updated C++ version in JNI from c++11 to c++17 [#2259](https://github.com/opensearch-project/k-NN/pull/2259) * Upgrade bytebuddy and objenesis version to match OpenSearch core and, update github ci runner for macos [#2279](https://github.com/opensearch-project/k-NN/pull/2279) From 84e9c74c382a0d1b66082f0d22f06a5d10ccf0cf Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Tue, 31 Dec 2024 12:57:32 -0800 Subject: [PATCH 4/8] Added unit tests to test multi shard filters Signed-off-by: Sahil Buddharaju --- .../org/opensearch/knn/index/FaissIT.java | 48 ++++++++++++++++++- .../opensearch/knn/index/LuceneEngineIT.java | 48 ++++++++++++++++++- 2 files changed, 92 insertions(+), 4 deletions(-) diff --git a/src/test/java/org/opensearch/knn/index/FaissIT.java b/src/test/java/org/opensearch/knn/index/FaissIT.java index 1618ec4d8..c2e75ecb2 100644 --- a/src/test/java/org/opensearch/knn/index/FaissIT.java +++ b/src/test/java/org/opensearch/knn/index/FaissIT.java @@ -412,6 +412,51 @@ public void testRadialQuery_withFilter_thenSuccess() { deleteKNNIndex(INDEX_NAME); } + @SneakyThrows + public void testQueryWithFilterMultipleShards() { + XContentBuilder builder = XContentFactory.jsonBuilder() + .startObject() + .startObject(PROPERTIES_FIELD_NAME) + .startObject(FIELD_NAME) + .field(TYPE_FIELD_NAME, KNN_VECTOR_TYPE) + .field(DIMENSION_FIELD_NAME, "3") + .startObject(KNNConstants.KNN_METHOD) + .field(KNNConstants.NAME, METHOD_HNSW) + .field(KNNConstants.METHOD_PARAMETER_SPACE_TYPE, SpaceType.L2.getValue()) + .field(KNNConstants.KNN_ENGINE, KNNEngine.FAISS.getName()) + .endObject() + .endObject() + .startObject(INTEGER_FIELD_NAME) + .field(TYPE_FIELD_NAME, FILED_TYPE_INTEGER) + .endObject() + .endObject() + .endObject(); + String mapping = builder.toString(); + + createIndex(INDEX_NAME, Settings.builder().put("number_of_shards", 10).put("number_of_replicas", 0).put("index.knn", true).build()); + putMappingRequest(INDEX_NAME, mapping); + + addKnnDocWithAttributes("doc1", new float[] { 7.0f, 7.0f, 3.0f }, ImmutableMap.of("dateReceived", "2024-10-01")); + + refreshIndex(INDEX_NAME); + + final float[] searchVector = { 6.0f, 7.0f, 3.0f }; + final Response response = searchKNNIndex( + INDEX_NAME, + new KNNQueryBuilder( + FIELD_NAME, + searchVector, + 1, + QueryBuilders.boolQuery().must(QueryBuilders.rangeQuery("dateReceived").gte("2023-11-01")) + ), + 10 + ); + final String responseBody = EntityUtils.toString(response.getEntity()); + final List knnResults = parseSearchResponse(responseBody, FIELD_NAME); + + assertEquals(1, knnResults.size()); + } + @SneakyThrows public void testEndToEnd_whenMethodIsHNSWPQ_thenSucceed() { String indexName = "test-index"; @@ -1956,8 +2001,7 @@ public void testQueryWithFilter_whenNonExistingFieldUsedInFilter_thenSuccessful( Map mappingMap = xContentBuilderToMap(builder); String mapping = builder.toString(); - createIndex(INDEX_NAME, Settings.builder().put("number_of_shards", 2).put("number_of_replicas", 1).put("index.knn", true).build()); - putMappingRequest(INDEX_NAME, mapping); + createKnnIndex(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 index 25dc6ccd4..b478415a0 100644 --- a/src/test/java/org/opensearch/knn/index/LuceneEngineIT.java +++ b/src/test/java/org/opensearch/knn/index/LuceneEngineIT.java @@ -283,6 +283,51 @@ public void testQueryWithFilterUsingByteVectorDataType() { validateQueryResultsWithFilters(searchVector, 5, 1, expectedDocIdsKGreaterThanFilterResult, expectedDocIdsKLimitsFilterResult); } + @SneakyThrows + public void testQueryWithFilterMultipleShards() { + XContentBuilder builder = XContentFactory.jsonBuilder() + .startObject() + .startObject(PROPERTIES_FIELD_NAME) + .startObject(FIELD_NAME) + .field(TYPE_FIELD_NAME, KNN_VECTOR_TYPE) + .field(DIMENSION_FIELD_NAME, DIMENSION) + .startObject(KNNConstants.KNN_METHOD) + .field(KNNConstants.NAME, METHOD_HNSW) + .field(KNNConstants.METHOD_PARAMETER_SPACE_TYPE, SpaceType.L2.getValue()) + .field(KNNConstants.KNN_ENGINE, KNNEngine.LUCENE.getName()) + .endObject() + .endObject() + .startObject(INTEGER_FIELD_NAME) + .field(TYPE_FIELD_NAME, FILED_TYPE_INTEGER) + .endObject() + .endObject() + .endObject(); + String mapping = builder.toString(); + + createIndex(INDEX_NAME, Settings.builder().put("number_of_shards", 10).put("number_of_replicas", 0).put("index.knn", true).build()); + putMappingRequest(INDEX_NAME, mapping); + + addKnnDocWithAttributes("doc1", new float[] { 7.0f, 7.0f, 3.0f }, ImmutableMap.of("dateReceived", "2024-10-01")); + + refreshIndex(INDEX_NAME); + + final float[] searchVector = { 6.0f, 7.0f, 3.0f }; + final Response response = searchKNNIndex( + INDEX_NAME, + new KNNQueryBuilder( + FIELD_NAME, + searchVector, + 1, + QueryBuilders.boolQuery().must(QueryBuilders.rangeQuery("dateReceived").gte("2023-11-01")) + ), + 10 + ); + final String responseBody = EntityUtils.toString(response.getEntity()); + final List knnResults = parseSearchResponse(responseBody, FIELD_NAME); + + assertEquals(1, knnResults.size()); + } + @SneakyThrows public void testQueryWithFilter_whenNonExistingFieldUsedInFilter_thenSuccessful() { XContentBuilder builder = XContentFactory.jsonBuilder() @@ -305,8 +350,7 @@ public void testQueryWithFilter_whenNonExistingFieldUsedInFilter_thenSuccessful( Map mappingMap = xContentBuilderToMap(builder); String mapping = builder.toString(); - createIndex(INDEX_NAME, Settings.builder().put("number_of_shards", 2).put("number_of_replicas", 1).put("index.knn", true).build()); - putMappingRequest(INDEX_NAME, mapping); + createKnnIndex(INDEX_NAME, mapping); Float[] vector = new Float[] { 2.0f, 4.5f, 6.5f }; From 17d7a2435887929eb1a9c4588ed72eca2a3c3375 Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Thu, 2 Jan 2025 13:51:42 -0800 Subject: [PATCH 5/8] Changed unit tests and used builder instead of constructor Signed-off-by: Sahil Buddharaju --- .../knn/index/query/KNNQueryBuilder.java | 26 +++++++++---------- .../knn/index/query/KNNQueryBuilderTests.java | 7 +++++ 2 files changed, 20 insertions(+), 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 index 62400a509..b03a40693 100644 --- 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 index 30c8007e6..e82e19453 100644 --- 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); } From d6a2ebb982f8bc22c7643811446a0af84cc379da Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Thu, 2 Jan 2025 14:12:14 -0800 Subject: [PATCH 6/8] Spotless apply Signed-off-by: Sahil Buddharaju --- .../knn/index/query/KNNQueryBuilder.java | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 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 index b03a40693..ee18394f6 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java @@ -667,17 +667,17 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws I rewrittenFilter = filter.rewrite(queryShardContext); if (rewrittenFilter != filter) { 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(); + .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; } } From 9a03e20f3fe86435f807b2b694df2d7fc97c2e42 Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Fri, 3 Jan 2025 09:20:02 -0800 Subject: [PATCH 7/8] slight unit test adjustment Signed-off-by: Sahil Buddharaju --- .../opensearch/knn/index/query/KNNQueryBuilderTests.java | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java index e82e19453..1333d616e 100644 --- a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java +++ b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java @@ -1067,16 +1067,12 @@ public void testDoRewrite_whenFilterSet_thenSuccessful() { .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); + assertEquals(knnQueryBuilder, KNNQueryBuilder.builder().fieldName(FIELD_NAME).vector(QUERY_VECTOR).filter(filter).k(K).build()); // Then assertEquals(expected, actual); From 887af64a4b18bd9602965dd4f6e2fda9de2509bd Mon Sep 17 00:00:00 2001 From: Sahil Buddharaju Date: Fri, 3 Jan 2025 09:58:39 -0800 Subject: [PATCH 8/8] changelog Signed-off-by: Sahil Buddharaju --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a767a63b7..ab1d87b6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,7 +28,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), * Allow validation for non knn index only after 2.17.0 (#2315)[https://github.com/opensearch-project/k-NN/pull/2315] * Release query vector memory after execution (#2346)[https://github.com/opensearch-project/k-NN/pull/2346] * Fix shard level rescoring disabled setting flag (#2352)[https://github.com/opensearch-project/k-NN/pull/2352] -* Fix filter rewrite logic getting inconsistent / incorrect results (#2359)[https://github.com/opensearch-project/k-NN/pull/2359] +* Fix filter rewrite logic which was resulting in getting inconsistent / incorrect results for cases where filter was getting rewritten for shards (#2359)[https://github.com/opensearch-project/k-NN/pull/2359] ### Infrastructure * Updated C++ version in JNI from c++11 to c++17 [#2259](https://github.com/opensearch-project/k-NN/pull/2259) * Upgrade bytebuddy and objenesis version to match OpenSearch core and, update github ci runner for macos [#2279](https://github.com/opensearch-project/k-NN/pull/2279)