|
67 | 67 | import org.opensearch.knn.index.mapper.KNNVectorFieldType;
|
68 | 68 | import org.opensearch.knn.index.query.KNNQuery;
|
69 | 69 | import org.opensearch.knn.index.query.KNNQueryBuilder;
|
| 70 | +import org.opensearch.knn.index.query.nativelib.NativeEngineKnnVectorQuery; |
70 | 71 | import org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils;
|
71 | 72 | import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil;
|
72 | 73 |
|
@@ -141,8 +142,8 @@ public void testDoToQuery_whenOneSubquery_thenBuildSuccessfully() {
|
141 | 142 | assertNotNull(queryOnlyNeural);
|
142 | 143 | assertTrue(queryOnlyNeural instanceof HybridQuery);
|
143 | 144 | 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(); |
146 | 147 | assertEquals(VECTOR_FIELD_NAME, knnQuery.getField());
|
147 | 148 | assertEquals(K, knnQuery.getK());
|
148 | 149 | assertNotNull(knnQuery.getQueryVector());
|
@@ -183,8 +184,8 @@ public void testDoToQuery_whenMultipleSubqueries_thenBuildSuccessfully() {
|
183 | 184 | // verify knn vector query
|
184 | 185 | Iterator<Query> queryIterator = ((HybridQuery) queryTwoSubQueries).getSubQueries().iterator();
|
185 | 186 | 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(); |
188 | 189 | assertEquals(VECTOR_FIELD_NAME, knnQuery.getField());
|
189 | 190 | assertEquals(K, knnQuery.getK());
|
190 | 191 | assertNotNull(knnQuery.getQueryVector());
|
|
0 commit comments