Skip to content

Commit 6481b60

Browse files
Fixed failed test after knn pushed multiple inner hits feature (opensearch-project#1026) (opensearch-project#1027)
Signed-off-by: Martin Gaievski <gaievski@amazon.com> (cherry picked from commit 3d72cc3) Co-authored-by: Martin Gaievski <gaievski@amazon.com>
1 parent d52c3e4 commit 6481b60

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java

+5-4
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
import org.opensearch.knn.index.mapper.KNNVectorFieldType;
6868
import org.opensearch.knn.index.query.KNNQuery;
6969
import org.opensearch.knn.index.query.KNNQueryBuilder;
70+
import org.opensearch.knn.index.query.nativelib.NativeEngineKnnVectorQuery;
7071
import org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils;
7172
import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil;
7273

@@ -141,8 +142,8 @@ public void testDoToQuery_whenOneSubquery_thenBuildSuccessfully() {
141142
assertNotNull(queryOnlyNeural);
142143
assertTrue(queryOnlyNeural instanceof HybridQuery);
143144
assertEquals(1, ((HybridQuery) queryOnlyNeural).getSubQueries().size());
144-
assertTrue(((HybridQuery) queryOnlyNeural).getSubQueries().iterator().next() instanceof KNNQuery);
145-
KNNQuery knnQuery = (KNNQuery) ((HybridQuery) queryOnlyNeural).getSubQueries().iterator().next();
145+
assertTrue(((HybridQuery) queryOnlyNeural).getSubQueries().iterator().next() instanceof NativeEngineKnnVectorQuery);
146+
KNNQuery knnQuery = ((NativeEngineKnnVectorQuery) ((HybridQuery) queryOnlyNeural).getSubQueries().iterator().next()).getKnnQuery();
146147
assertEquals(VECTOR_FIELD_NAME, knnQuery.getField());
147148
assertEquals(K, knnQuery.getK());
148149
assertNotNull(knnQuery.getQueryVector());
@@ -183,8 +184,8 @@ public void testDoToQuery_whenMultipleSubqueries_thenBuildSuccessfully() {
183184
// verify knn vector query
184185
Iterator<Query> queryIterator = ((HybridQuery) queryTwoSubQueries).getSubQueries().iterator();
185186
Query firstQuery = queryIterator.next();
186-
assertTrue(firstQuery instanceof KNNQuery);
187-
KNNQuery knnQuery = (KNNQuery) firstQuery;
187+
assertTrue(firstQuery instanceof NativeEngineKnnVectorQuery);
188+
KNNQuery knnQuery = ((NativeEngineKnnVectorQuery) firstQuery).getKnnQuery();
188189
assertEquals(VECTOR_FIELD_NAME, knnQuery.getField());
189190
assertEquals(K, knnQuery.getK());
190191
assertNotNull(knnQuery.getQueryVector());

0 commit comments

Comments
 (0)