From 0f489653d4aca59d056ef5c95fa3e9099ec0191c Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Wed, 11 Dec 2024 11:04:13 -0800 Subject: [PATCH] Minor code fixes which includes: (#2206) (#2322) * Passing correct score mode in NativeEngineKNNVectorQuery * Ensuring visitor is called in KnnQuery Signed-off-by: Navneet Verma (cherry picked from commit f810493cf9d255da7532e0bb1d59d4bcee927d87) Co-authored-by: Navneet Verma --- .../org/opensearch/knn/index/query/KNNQuery.java | 2 +- .../nativelib/NativeEngineKnnVectorQuery.java | 2 +- .../nativelib/NativeEngineKNNVectorQueryTests.java | 14 ++++++++------ 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/query/KNNQuery.java b/src/main/java/org/opensearch/knn/index/query/KNNQuery.java index f5c4d3131..f0974f7e9 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNQuery.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNQuery.java @@ -188,7 +188,7 @@ private Weight getFilterWeight(IndexSearcher searcher) throws IOException { @Override public void visit(QueryVisitor visitor) { - + visitor.visitLeaf(this); } @Override diff --git a/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java b/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java index bf12e63c0..f782b0180 100644 --- a/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java +++ b/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java @@ -54,7 +54,7 @@ public class NativeEngineKnnVectorQuery extends Query { @Override public Weight createWeight(IndexSearcher indexSearcher, ScoreMode scoreMode, float boost) throws IOException { final IndexReader reader = indexSearcher.getIndexReader(); - final KNNWeight knnWeight = (KNNWeight) knnQuery.createWeight(indexSearcher, ScoreMode.COMPLETE, 1); + final KNNWeight knnWeight = (KNNWeight) knnQuery.createWeight(indexSearcher, scoreMode, 1); List leafReaderContexts = reader.leaves(); List> perLeafResults; RescoreContext rescoreContext = knnQuery.getRescoreContext(); diff --git a/src/test/java/org/opensearch/knn/index/query/nativelib/NativeEngineKNNVectorQueryTests.java b/src/test/java/org/opensearch/knn/index/query/nativelib/NativeEngineKNNVectorQueryTests.java index fab67189d..01e3fa6f9 100644 --- a/src/test/java/org/opensearch/knn/index/query/nativelib/NativeEngineKNNVectorQueryTests.java +++ b/src/test/java/org/opensearch/knn/index/query/nativelib/NativeEngineKNNVectorQueryTests.java @@ -76,6 +76,8 @@ public class NativeEngineKNNVectorQueryTests extends OpenSearchTestCase { @InjectMocks private NativeEngineKnnVectorQuery objectUnderTest; + private static ScoreMode scoreMode = ScoreMode.TOP_SCORES; + @Override public void setUp() throws Exception { super.setUp(); @@ -85,7 +87,7 @@ public void setUp() throws Exception { when(leaf2.reader()).thenReturn(leafReader2); when(searcher.getIndexReader()).thenReturn(reader); - when(knnQuery.createWeight(searcher, ScoreMode.COMPLETE, 1)).thenReturn(knnWeight); + when(knnQuery.createWeight(searcher, scoreMode, 1)).thenReturn(knnWeight); when(searcher.getTaskExecutor()).thenReturn(taskExecutor); when(taskExecutor.invokeAll(any())).thenAnswer(invocationOnMock -> { @@ -135,7 +137,7 @@ public void testMultiLeaf() { Query expected = new DocAndScoreQuery(4, expectedDocs, expectedScores, findSegments, 1); // When - Weight actual = objectUnderTest.createWeight(searcher, ScoreMode.COMPLETE, 1); + Weight actual = objectUnderTest.createWeight(searcher, scoreMode, 1); // Then assertEquals(expected, actual.getQuery()); @@ -176,7 +178,7 @@ public void testRescoreWhenShardLevelRescoringEnabled() { mockedResultUtil.when(() -> ResultUtil.reduceToTopK(any(), anyInt())).thenCallRealMethod(); // When - Weight actual = objectUnderTest.createWeight(searcher, ScoreMode.COMPLETE, 1); + Weight actual = objectUnderTest.createWeight(searcher, scoreMode, 1); // Then mockedResultUtil.verify(() -> ResultUtil.reduceToTopK(any(), anyInt()), times(2)); @@ -199,7 +201,7 @@ public void testSingleLeaf() { Query expected = new DocAndScoreQuery(4, expectedDocs, expectedScores, findSegments, 1); // When - Weight actual = objectUnderTest.createWeight(searcher, ScoreMode.COMPLETE, 1); + Weight actual = objectUnderTest.createWeight(searcher, scoreMode, 1); // Then assertEquals(expected, actual.getQuery()); @@ -214,7 +216,7 @@ public void testNoMatch() { when(knnQuery.getK()).thenReturn(4); // When - Weight actual = objectUnderTest.createWeight(searcher, ScoreMode.COMPLETE, 1); + Weight actual = objectUnderTest.createWeight(searcher, scoreMode, 1); // Then assertEquals(new MatchNoDocsQuery(), actual.getQuery()); @@ -263,7 +265,7 @@ public void testRescore() { try (MockedStatic mockedStaticNativeKnnVectorQuery = mockStatic(NativeEngineKnnVectorQuery.class)) { mockedStaticNativeKnnVectorQuery.when(() -> NativeEngineKnnVectorQuery.findSegmentStarts(any(), any())) .thenReturn(new int[] { 0, 4, 2 }); - Weight actual = objectUnderTest.createWeight(searcher, ScoreMode.COMPLETE, 1); + Weight actual = objectUnderTest.createWeight(searcher, scoreMode, 1); assertEquals(expected, actual.getQuery()); } }