Skip to content

Commit dc7aeb6

Browse files
committed
Add VectorSearchQueryBuilder unit tests
Verifies knn query is placed in scoring (must) context, not wrapped in bool.filter when no WHERE clause is present. Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
1 parent e05fef7 commit dc7aeb6

1 file changed

Lines changed: 50 additions & 0 deletions

File tree

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.storage.scan;
7+
8+
import static org.junit.jupiter.api.Assertions.assertTrue;
9+
import static org.mockito.Mockito.mock;
10+
11+
import org.junit.jupiter.api.Test;
12+
import org.opensearch.index.query.QueryBuilder;
13+
import org.opensearch.index.query.WrapperQueryBuilder;
14+
import org.opensearch.sql.common.setting.Settings;
15+
import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
16+
import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
17+
18+
class VectorSearchQueryBuilderTest {
19+
20+
@Test
21+
void knnQuerySetAsScoringQuery() {
22+
var requestBuilder = createRequestBuilder();
23+
var knnQuery = new WrapperQueryBuilder("eyJrbm4iOnt9fQ==");
24+
25+
new VectorSearchQueryBuilder(requestBuilder, knnQuery);
26+
27+
QueryBuilder query = requestBuilder.getSourceBuilder().query();
28+
assertTrue(
29+
query instanceof WrapperQueryBuilder,
30+
"knn query should be set directly as top-level query (scoring context)");
31+
}
32+
33+
@Test
34+
void knnQueryNotWrappedInFilterWhenNoWhere() {
35+
var requestBuilder = createRequestBuilder();
36+
var knnQuery = new WrapperQueryBuilder("eyJrbm4iOnt9fQ==");
37+
38+
new VectorSearchQueryBuilder(requestBuilder, knnQuery);
39+
40+
QueryBuilder query = requestBuilder.getSourceBuilder().query();
41+
assertTrue(
42+
query instanceof WrapperQueryBuilder,
43+
"Without WHERE clause, knn query should NOT be wrapped in bool.filter");
44+
}
45+
46+
private OpenSearchRequestBuilder createRequestBuilder() {
47+
return new OpenSearchRequestBuilder(
48+
mock(OpenSearchExprValueFactory.class), 10000, mock(Settings.class));
49+
}
50+
}

0 commit comments

Comments
 (0)