Skip to content

Commit 6931643

Browse files
carlosdelestcarlosdelest
committed
Fix tests
1 parent 54e6407 commit 6931643

File tree

9 files changed

+47
-43
lines changed

9 files changed

+47
-43
lines changed

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/knn-function.csv-spec‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
# top-n query at the shard level
44

55
knnSearch
6-
required_capability: knn_function
6+
required_capability: knn_function_v2
77

88
// tag::knn-function[]
99
from colors metadata _score
@@ -30,7 +30,7 @@ chartreuse | [127.0, 255.0, 0.0]
3030
;
3131

3232
knnSearchWithKOption
33-
required_capability: knn_function
33+
required_capability: knn_function_v2
3434

3535
// tag::knn-function-options[]
3636
from colors metadata _score
@@ -49,7 +49,7 @@ teal | [0.0, 128.0, 128.0]
4949
;
5050

5151
knnSearchWithSimilarityOption
52-
required_capability: knn_function
52+
required_capability: knn_function_v2
5353

5454
from colors metadata _score
5555
| where knn(rgb_vector, [255,192,203], {"similarity": 40})
@@ -66,7 +66,7 @@ wheat | [245.0, 222.0, 179.0]
6666
;
6767

6868
knnHybridSearch
69-
required_capability: knn_function
69+
required_capability: knn_function_v2
7070

7171
from colors metadata _score
7272
| where match(color, "blue") or knn(rgb_vector, [65,105,225])
@@ -89,7 +89,7 @@ aqua marine | [127.0, 255.0, 212.0]
8989
;
9090

9191
knnWithMultipleFunctions
92-
required_capability: knn_function
92+
required_capability: knn_function_v2
9393

9494
from colors metadata _score
9595
| where knn(rgb_vector, [128,128,0]) and match(color, "olive")
@@ -102,7 +102,7 @@ olive | [128.0, 128.0, 0.0]
102102
;
103103

104104
knnAfterKeep
105-
required_capability: knn_function
105+
required_capability: knn_function_v2
106106

107107
from colors metadata _score
108108
| keep rgb_vector, color, _score
@@ -121,7 +121,7 @@ rgb_vector:dense_vector
121121
;
122122

123123
knnAfterDrop
124-
required_capability: knn_function
124+
required_capability: knn_function_v2
125125

126126
from colors metadata _score
127127
| drop primary
@@ -140,7 +140,7 @@ lime | [0.0, 255.0, 0.0]
140140
;
141141

142142
knnAfterEval
143-
required_capability: knn_function
143+
required_capability: knn_function_v2
144144

145145
from colors metadata _score
146146
| eval composed_name = locate(color, " ") > 0
@@ -159,7 +159,7 @@ golden rod | true
159159
;
160160

161161
knnWithConjunction
162-
required_capability: knn_function
162+
required_capability: knn_function_v2
163163

164164
# TODO We need kNN prefiltering here so we get more candidates that pass the filter
165165
from colors metadata _score
@@ -180,14 +180,13 @@ yellow | #FFFF00 | [255.0, 255.0, 0.0]
180180
;
181181

182182
knnWithDisjunctionAndFiltersConjunction
183-
required_capability: knn_function
183+
required_capability: knn_function_v2
184184

185185
# TODO We need kNN prefiltering here so we get more candidates that pass the filter
186186
from colors metadata _score
187187
| where (knn(rgb_vector, [0,255,255]) or knn(rgb_vector, [128, 0, 255])) and primary == true
188188
| sort _score desc, color asc
189189
| drop _score, primary, hex_code
190-
| limit 140
191190
;
192191

193192
color:text | rgb_vector:dense_vector
@@ -203,7 +202,7 @@ yellow | [255.0, 255.0, 0.0]
203202
;
204203

205204
knnWithNonPushableConjunction
206-
required_capability: knn_function
205+
required_capability: knn_function_v2
207206

208207
from colors metadata _score
209208
| eval composed_name = locate(color, " ") > 0
@@ -226,7 +225,7 @@ maroon | false
226225
;
227226

228227
testKnnWithNonPushableDisjunctions
229-
required_capability: knn_function
228+
required_capability: knn_function_v2
230229

231230
from colors metadata _score
232231
| where knn(rgb_vector, [128,128,0], {"similarity": 30}) or length(color) > 10
@@ -242,7 +241,7 @@ papaya whip
242241
;
243242

244243
testKnnWithNonPushableDisjunctionsOnComplexExpressions
245-
required_capability: knn_function
244+
required_capability: knn_function_v2
246245

247246
from colors metadata _score
248247
| where (knn(rgb_vector, [128,128,0], {"similarity": 70}) and length(color) < 10) or (knn(rgb_vector, [128,0,128], {"similarity": 60}) and primary == false)

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KnnFunctionIT.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ public void testKnnDefaults() {
4949
assertColumnTypes(resp.columns(), List.of("integer", "double", "double", "dense_vector"));
5050

5151
List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
52-
assertEquals(Math.min(indexedVectors.size(), 10), valuesList.size());
52+
assertEquals(indexedVectors.size(), valuesList.size());
5353
for (int i = 0; i < valuesList.size(); i++) {
5454
List<Object> row = valuesList.get(i);
5555
// Vectors should be in order of ID, as they're less similar than the query vector as the ID increases
@@ -111,7 +111,7 @@ public void testKnnNonPushedDown() {
111111

112112
@Before
113113
public void setup() throws IOException {
114-
assumeTrue("Needs KNN support", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
114+
assumeTrue("Needs KNN support", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
115115

116116
var indexName = "test";
117117
var client = client().admin().indices();

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1195,7 +1195,7 @@ public enum Cap {
11951195
/**
11961196
* Support knn function
11971197
*/
1198-
KNN_FUNCTION(Build.current().isSnapshot());
1198+
KNN_FUNCTION_V2(Build.current().isSnapshot());
11991199

12001200
private final boolean enabled;
12011201

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/ExpressionWritables.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ private static List<NamedWriteableRegistry.Entry> fullText() {
257257
}
258258

259259
private static List<NamedWriteableRegistry.Entry> vector() {
260-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
260+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
261261
return List.of(Knn.ENTRY);
262262
}
263263
return List.of();

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -298,7 +298,7 @@ public final void test() throws Throwable {
298298
);
299299
assumeFalse(
300300
"can't use KNN function in csv tests",
301-
testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.KNN_FUNCTION.capabilityName())
301+
testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.KNN_FUNCTION_V2.capabilityName())
302302
);
303303
assumeFalse(
304304
"lookup join disabled for csv tests",

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/VerifierTests.java‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1235,7 +1235,7 @@ public void testFieldBasedFullTextFunctions() throws Exception {
12351235
checkFieldBasedWithNonIndexedColumn("Term", "term(text, \"cat\")", "function");
12361236
checkFieldBasedFunctionNotAllowedAfterCommands("Term", "function", "term(title, \"Meditation\")");
12371237
}
1238-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
1238+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
12391239
checkFieldBasedFunctionNotAllowedAfterCommands("KNN", "function", "knn(vector, [1, 2, 3])");
12401240
}
12411241
}
@@ -1369,7 +1369,7 @@ public void testFullTextFunctionsOnlyAllowedInWhere() throws Exception {
13691369
if (EsqlCapabilities.Cap.MULTI_MATCH_FUNCTION.isEnabled()) {
13701370
checkFullTextFunctionsOnlyAllowedInWhere("MultiMatch", "multi_match(\"Meditation\", title, body)", "function", defaultErrorMsg);
13711371
}
1372-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
1372+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
13731373
checkFullTextFunctionsOnlyAllowedInWhere("KNN", "knn(vector, [0, 1, 2])", "function", "is only supported in WHERE commands");
13741374
}
13751375
}
@@ -1412,7 +1412,7 @@ public void testFullTextFunctionsDisjunctions() {
14121412
if (EsqlCapabilities.Cap.TERM_FUNCTION.isEnabled()) {
14131413
checkWithFullTextFunctionsDisjunctions("term(title, \"Meditation\")");
14141414
}
1415-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
1415+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
14161416
checkWithFullTextFunctionsDisjunctions("knn(vector, [1, 2, 3])");
14171417
}
14181418
}
@@ -1477,7 +1477,7 @@ public void testFullTextFunctionsWithNonBooleanFunctions() {
14771477
if (EsqlCapabilities.Cap.TERM_FUNCTION.isEnabled()) {
14781478
checkFullTextFunctionsWithNonBooleanFunctions("Term", "term(title, \"Meditation\")", "function");
14791479
}
1480-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
1480+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
14811481
checkFullTextFunctionsWithNonBooleanFunctions("KNN", "knn(vector, [1, 2, 3])", "function");
14821482
}
14831483
}
@@ -1548,7 +1548,7 @@ public void testFullTextFunctionsTargetsExistingField() throws Exception {
15481548
if (EsqlCapabilities.Cap.TERM_FUNCTION.isEnabled()) {
15491549
testFullTextFunctionTargetsExistingField("term(fist_name, \"Meditation\")");
15501550
}
1551-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
1551+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
15521552
testFullTextFunctionTargetsExistingField("knn(vector, [0, 1, 2])");
15531553
}
15541554
}
@@ -2076,7 +2076,7 @@ public void testFullTextFunctionOptions() {
20762076
if (EsqlCapabilities.Cap.MULTI_MATCH_FUNCTION.isEnabled()) {
20772077
checkOptionDataTypes(MultiMatch.OPTIONS, "FROM test | WHERE MULTI_MATCH(\"Jean\", title, body, {\"%s\": %s})");
20782078
}
2079-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
2079+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
20802080
checkOptionDataTypes(Knn.ALLOWED_OPTIONS, "FROM test | WHERE KNN(vector, [0.1, 0.2, 0.3], {\"%s\": %s})");
20812081
}
20822082
}
@@ -2164,7 +2164,7 @@ public void testFullTextFunctionsNullArgs() throws Exception {
21642164
checkFullTextFunctionNullArgs("term(null, \"query\")", "first");
21652165
checkFullTextFunctionNullArgs("term(title, null)", "second");
21662166
}
2167-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
2167+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
21682168
checkFullTextFunctionNullArgs("knn(null, [0, 1, 2])", "first");
21692169
checkFullTextFunctionNullArgs("knn(vector, null)", "second");
21702170
}
@@ -2189,7 +2189,7 @@ public void testFullTextFunctionsConstantQuery() throws Exception {
21892189
if (EsqlCapabilities.Cap.TERM_FUNCTION.isEnabled()) {
21902190
checkFullTextFunctionsConstantQuery("term(title, tags)", "second");
21912191
}
2192-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
2192+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
21932193
checkFullTextFunctionsConstantQuery("knn(vector, vector)", "second");
21942194
}
21952195
}
@@ -2219,7 +2219,7 @@ public void testFullTextFunctionsInStats() {
22192219
if (EsqlCapabilities.Cap.MULTI_MATCH_FUNCTION.isEnabled()) {
22202220
checkFullTextFunctionsInStats("multi_match(\"Meditation\", title, body)");
22212221
}
2222-
if (EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled()) {
2222+
if (EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled()) {
22232223
checkFullTextFunctionsInStatsError("knn(vector, [0, 1, 2])");
22242224
}
22252225
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/fulltext/KnnTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ public static Iterable<Object[]> parameters() {
5050

5151
@Before
5252
public void checkCapability() {
53-
assumeTrue("KNN is not enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
53+
assumeTrue("KNN is not enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
5454
}
5555

5656
private static List<TestCaseSupplier> testCaseSuppliers() {

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1365,7 +1365,7 @@ public void testMultiMatchOptionsPushDown() {
13651365
}
13661366

13671367
public void testKnnOptionsPushDown() {
1368-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
1368+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
13691369

13701370
String query = """
13711371
from test
@@ -2061,7 +2061,7 @@ public void testToDateNanosPushDown() {
20612061
}
20622062

20632063
public void testKnnWithoutExplicitLimit() {
2064-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2064+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
20652065

20662066
var query = """
20672067
from test
@@ -2081,7 +2081,7 @@ public void testKnnWithoutExplicitLimit() {
20812081
}
20822082

20832083
public void testKnnWithoutExplicitLimitSorted() {
2084-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2084+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
20852085

20862086
var query = """
20872087
from test metadata _score
@@ -2102,7 +2102,7 @@ public void testKnnWithoutExplicitLimitSorted() {
21022102
}
21032103

21042104
public void testKnnWithExplicitLimitSorted() {
2105-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2105+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
21062106

21072107
var query = """
21082108
from test metadata _score
@@ -2124,7 +2124,7 @@ public void testKnnWithExplicitLimitSorted() {
21242124
}
21252125

21262126
public void testKnnWithStatsAndLimit() {
2127-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2127+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
21282128

21292129
var query = """
21302130
from test metadata _score
@@ -2150,7 +2150,7 @@ public void testKnnWithStatsAndLimit() {
21502150
}
21512151

21522152
public void testKnnWithoutExplicitLimitSortedAndCommandsInBetween() {
2153-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2153+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
21542154

21552155
var query = """
21562156
from test metadata _score
@@ -2173,11 +2173,11 @@ public void testKnnWithoutExplicitLimitSortedAndCommandsInBetween() {
21732173
}
21742174

21752175
public void testKnnWithExplicitLimit() {
2176-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2176+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
21772177

21782178
var query = """
21792179
from test metadata _score
2180-
| where knn(dense_vector, [0, 1, 2]) or match(text, "blue")
2180+
| where (knn(dense_vector, [0, 1, 2]) or knn(dense_vector, [1, 2, 3])) and match(text, "blue")
21812181
| sort _score desc, text asc
21822182
| keep text, dense_vector
21832183
| limit 140
@@ -2195,13 +2195,18 @@ public void testKnnWithExplicitLimit() {
21952195
var fieldExtractExec2 = as(topNExec2.child(), FieldExtractExec.class);
21962196
var queryExec = as(fieldExtractExec2.child(), EsQueryExec.class);
21972197
assertNull(queryExec.limit());
2198-
var expectedQuery = boolQuery().should(new KnnVectorQueryBuilder("dense_vector", new float[] { 0f, 1f, 2f }, 140, null, null, null))
2199-
.should(matchQuery("text", "blue").lenient(true));
2198+
var expectedQuery = boolQuery()
2199+
.must(
2200+
boolQuery()
2201+
.should(new KnnVectorQueryBuilder("dense_vector", new float[] { 0f, 1f, 2f }, 140, null, null, null))
2202+
.should(new KnnVectorQueryBuilder("dense_vector", new float[] { 1f, 2f, 3f }, 140, null, null, null))
2203+
)
2204+
.must(matchQuery("text", "blue").lenient(true));
22002205
assertEquals(expectedQuery.toString(), queryExec.query().toString());
22012206
}
22022207

22032208
public void testKnnWithExplicitLimitAndExistingTopK() {
2204-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2209+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
22052210

22062211
var query = """
22072212
from test
@@ -2222,7 +2227,7 @@ public void testKnnWithExplicitLimitAndExistingTopK() {
22222227
}
22232228

22242229
public void testMultipleKnnQueriesLimit() {
2225-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2230+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
22262231

22272232
var query = """
22282233
from test
@@ -2269,7 +2274,7 @@ public void testMultipleKnnQueriesLimit() {
22692274
}
22702275

22712276
public void testKnnWithCombinedLimits() {
2272-
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
2277+
assumeTrue("KNN should be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
22732278

22742279
var query = """
22752280
from test

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7997,7 +7997,7 @@ public void testSampleNoPushDownChangePoint() {
79977997
}
79987998

79997999
public void testKnnFunctionHasLimitOrKSpecified() {
8000-
assumeTrue("KNN must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION.isEnabled());
8000+
assumeTrue("KNN must be enabled", EsqlCapabilities.Cap.KNN_FUNCTION_V2.isEnabled());
80018001

80028002
var query = """
80038003
FROM test

0 commit comments

Comments
 (0)