Skip to content

Commit 5dd108f

Browse files
mengweiericmengweieric
committed
Add comprehensive test coverage for vector search hardening
- Create VectorSearchIndexTest: 7 tests covering buildKnnQueryJson() for top-k, max_distance, min_score, nested fields, multi-element and single-element vectors, numeric option rendering - Add edge case tests to VectorSearchTableFunctionImplementationTest: NaN vector component, empty option key/value, negative k, NaN for max_distance and min_score (6 new tests) - Add VectorSearchQueryBuilderTest: min_score radial mode LIMIT, pushDownSort delegation to parent (2 new tests) - Extract buildKnnQueryJson() as package-private for direct testing Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
1 parent 354635c commit 5dd108f

4 files changed

Lines changed: 230 additions & 5 deletions

File tree

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/VectorSearchIndex.java‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,11 @@ public TableScanBuilder createScanBuilder() {
6868
}
6969

7070
private QueryBuilder buildKnnQuery() {
71+
return new WrapperQueryBuilder(buildKnnQueryJson());
72+
}
73+
74+
// Package-private for testing
75+
String buildKnnQueryJson() {
7176
StringBuilder vectorJson = new StringBuilder("[");
7277
for (int i = 0; i < vector.length; i++) {
7378
if (i > 0) vectorJson.append(",");
@@ -88,11 +93,9 @@ private QueryBuilder buildKnnQuery() {
8893
}
8994
}
9095

91-
String knnQueryJson =
92-
String.format(
93-
"{\"knn\":{\"%s\":{\"vector\":%s%s}}}",
94-
field, vectorJson.toString(), optionsJson.toString());
95-
return new WrapperQueryBuilder(knnQueryJson);
96+
return String.format(
97+
"{\"knn\":{\"%s\":{\"vector\":%s%s}}}",
98+
field, vectorJson.toString(), optionsJson.toString());
9699
}
97100

98101
private static boolean isNumeric(String str) {

‎opensearch/src/test/java/org/opensearch/sql/opensearch/storage/VectorSearchIndexTest.java‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.storage;
7+
8+
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertTrue;
10+
11+
import java.util.LinkedHashMap;
12+
import java.util.Map;
13+
import org.junit.jupiter.api.Test;
14+
import org.junit.jupiter.api.extension.ExtendWith;
15+
import org.mockito.Mock;
16+
import org.mockito.junit.jupiter.MockitoExtension;
17+
import org.opensearch.sql.common.setting.Settings;
18+
import org.opensearch.sql.opensearch.client.OpenSearchClient;
19+
20+
@ExtendWith(MockitoExtension.class)
21+
class VectorSearchIndexTest {
22+
23+
@Mock private OpenSearchClient client;
24+
25+
@Mock private Settings settings;
26+
27+
@Test
28+
void buildKnnQueryJsonTopK() {
29+
VectorSearchIndex index =
30+
new VectorSearchIndex(
31+
client,
32+
settings,
33+
"test-index",
34+
"embedding",
35+
new float[] {1.0f, 2.0f, 3.0f},
36+
Map.of("k", "5"));
37+
38+
String json = index.buildKnnQueryJson();
39+
assertEquals("{\"knn\":{\"embedding\":{\"vector\":[1.0,2.0,3.0],\"k\":5}}}", json);
40+
}
41+
42+
@Test
43+
void buildKnnQueryJsonRadialMaxDistance() {
44+
VectorSearchIndex index =
45+
new VectorSearchIndex(
46+
client,
47+
settings,
48+
"test-index",
49+
"embedding",
50+
new float[] {1.0f, 2.0f},
51+
Map.of("max_distance", "10.5"));
52+
53+
String json = index.buildKnnQueryJson();
54+
assertEquals("{\"knn\":{\"embedding\":{\"vector\":[1.0,2.0],\"max_distance\":10.5}}}", json);
55+
}
56+
57+
@Test
58+
void buildKnnQueryJsonRadialMinScore() {
59+
VectorSearchIndex index =
60+
new VectorSearchIndex(
61+
client,
62+
settings,
63+
"test-index",
64+
"embedding",
65+
new float[] {0.5f},
66+
Map.of("min_score", "0.8"));
67+
68+
String json = index.buildKnnQueryJson();
69+
assertEquals("{\"knn\":{\"embedding\":{\"vector\":[0.5],\"min_score\":0.8}}}", json);
70+
}
71+
72+
@Test
73+
void buildKnnQueryJsonNestedFieldName() {
74+
VectorSearchIndex index =
75+
new VectorSearchIndex(
76+
client,
77+
settings,
78+
"test-index",
79+
"doc.embedding",
80+
new float[] {1.0f, 2.0f},
81+
Map.of("k", "10"));
82+
83+
String json = index.buildKnnQueryJson();
84+
assertTrue(json.contains("\"doc.embedding\""), "Should contain nested field name with dot");
85+
}
86+
87+
@Test
88+
void buildKnnQueryJsonMultiElementVector() {
89+
VectorSearchIndex index =
90+
new VectorSearchIndex(
91+
client,
92+
settings,
93+
"test-index",
94+
"embedding",
95+
new float[] {1.0f, -2.5f, 0.0f, 3.14f, 100.0f},
96+
Map.of("k", "3"));
97+
98+
String json = index.buildKnnQueryJson();
99+
assertTrue(
100+
json.contains("[1.0,-2.5,0.0,3.14,100.0]"),
101+
"Should contain all vector components with correct comma separation");
102+
}
103+
104+
@Test
105+
void buildKnnQueryJsonSingleElementVector() {
106+
VectorSearchIndex index =
107+
new VectorSearchIndex(
108+
client, settings, "test-index", "embedding", new float[] {42.0f}, Map.of("k", "1"));
109+
110+
String json = index.buildKnnQueryJson();
111+
assertTrue(json.contains("[42.0]"), "Should contain single-element vector");
112+
}
113+
114+
@Test
115+
void buildKnnQueryJsonNumericOptionRenderedUnquoted() {
116+
LinkedHashMap<String, String> options = new LinkedHashMap<>();
117+
options.put("k", "5");
118+
119+
VectorSearchIndex index =
120+
new VectorSearchIndex(
121+
client, settings, "test-index", "embedding", new float[] {1.0f}, options);
122+
123+
String json = index.buildKnnQueryJson();
124+
assertTrue(json.contains("\"k\":5"), "Numeric option should be unquoted");
125+
}
126+
127+
@Test
128+
void isInstanceOfOpenSearchIndex() {
129+
VectorSearchIndex index =
130+
new VectorSearchIndex(
131+
client, settings, "test-index", "embedding", new float[] {1.0f}, Map.of("k", "5"));
132+
assertTrue(index instanceof OpenSearchIndex);
133+
}
134+
}

‎opensearch/src/test/java/org/opensearch/sql/opensearch/storage/VectorSearchTableFunctionImplementationTest.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,60 @@ void testNullArgNameThrows() {
319319
assertTrue(ex.getMessage().contains("requires named arguments"));
320320
}
321321

322+
@Test
323+
void testNaNVectorComponentThrows() {
324+
VectorSearchTableFunctionImplementation impl =
325+
createImplWithArgs("my-index", "embedding", "[1.0, NaN, 3.0]", "k=5");
326+
ExpressionEvaluationException ex =
327+
assertThrows(ExpressionEvaluationException.class, () -> impl.applyArguments());
328+
assertTrue(ex.getMessage().contains("must be a finite number"));
329+
}
330+
331+
@Test
332+
void testEmptyOptionKeyThrows() {
333+
ExpressionEvaluationException ex =
334+
assertThrows(
335+
ExpressionEvaluationException.class,
336+
() -> VectorSearchTableFunctionImplementation.parseOptions("=value"));
337+
assertTrue(ex.getMessage().contains("Malformed option segment"));
338+
}
339+
340+
@Test
341+
void testEmptyOptionValueThrows() {
342+
ExpressionEvaluationException ex =
343+
assertThrows(
344+
ExpressionEvaluationException.class,
345+
() -> VectorSearchTableFunctionImplementation.parseOptions("key="));
346+
assertTrue(ex.getMessage().contains("Malformed option segment"));
347+
}
348+
349+
@Test
350+
void testNegativeKThrows() {
351+
VectorSearchTableFunctionImplementation impl =
352+
createImplWithArgs("my-index", "embedding", "[1.0, 2.0]", "k=-1");
353+
ExpressionEvaluationException ex =
354+
assertThrows(ExpressionEvaluationException.class, () -> impl.applyArguments());
355+
assertTrue(ex.getMessage().contains("k must be between 1 and 10000"));
356+
}
357+
358+
@Test
359+
void testNaNMaxDistanceThrows() {
360+
VectorSearchTableFunctionImplementation impl =
361+
createImplWithArgs("my-index", "embedding", "[1.0, 2.0]", "max_distance=NaN");
362+
ExpressionEvaluationException ex =
363+
assertThrows(ExpressionEvaluationException.class, () -> impl.applyArguments());
364+
assertTrue(ex.getMessage().contains("must be a finite number"));
365+
}
366+
367+
@Test
368+
void testNaNMinScoreThrows() {
369+
VectorSearchTableFunctionImplementation impl =
370+
createImplWithArgs("my-index", "embedding", "[1.0, 2.0]", "min_score=NaN");
371+
ExpressionEvaluationException ex =
372+
assertThrows(ExpressionEvaluationException.class, () -> impl.applyArguments());
373+
assertTrue(ex.getMessage().contains("must be a finite number"));
374+
}
375+
322376
private VectorSearchTableFunctionImplementation createImpl() {
323377
return createImplWithArgs("my-index", "embedding", "[1.0, 2.0, 3.0]", "k=5");
324378
}

‎opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/VectorSearchQueryBuilderTest.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
1313

1414
import java.util.Collections;
15+
import java.util.List;
1516
import java.util.Map;
1617
import org.junit.jupiter.api.Test;
1718
import org.opensearch.index.query.BoolQueryBuilder;
@@ -120,6 +121,39 @@ void pushDownLimitRadialModeNoRestriction() {
120121
assertTrue(pushed, "Radial mode should not restrict LIMIT");
121122
}
122123

124+
@Test
125+
void pushDownLimitMinScoreModeNoRestriction() {
126+
var requestBuilder = createRequestBuilder();
127+
var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
128+
var builder =
129+
new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("min_score", "0.5"));
130+
131+
var dummyChild = new LogicalValues(Collections.emptyList());
132+
var limit = new LogicalLimit(dummyChild, 100, 0);
133+
134+
boolean pushed = builder.pushDownLimit(limit);
135+
assertTrue(pushed, "min_score mode should not restrict LIMIT");
136+
}
137+
138+
@Test
139+
void pushDownSortDelegatesToParent() {
140+
var requestBuilder = createRequestBuilder();
141+
var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
142+
var builder = new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("k", "5"));
143+
144+
var dummyChild = new LogicalValues(Collections.emptyList());
145+
var sort =
146+
new org.opensearch.sql.planner.logical.LogicalSort(
147+
dummyChild,
148+
List.of(
149+
org.apache.commons.lang3.tuple.ImmutablePair.of(
150+
org.opensearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC,
151+
new ReferenceExpression("name", STRING))));
152+
153+
boolean pushed = builder.pushDownSort(sort);
154+
assertTrue(pushed, "pushDownSort should delegate to parent and succeed");
155+
}
156+
123157
private OpenSearchRequestBuilder createRequestBuilder() {
124158
return new OpenSearchRequestBuilder(
125159
mock(OpenSearchExprValueFactory.class), 10000, mock(Settings.class));

0 commit comments

Comments
 (0)