Skip to content

Commit 204efda

Browse files
carlosdelestcarlosdelest
committed
Initial CSV tests
1 parent 7ae9909 commit 204efda

File tree

6 files changed

+220
-2
lines changed

6 files changed

+220
-2
lines changed

‎x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestsDataLoader.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ public class CsvTestsDataLoader {
144144
private static final TestDataset LOGS = new TestDataset("logs");
145145
private static final TestDataset MV_TEXT = new TestDataset("mv_text");
146146
private static final TestDataset DENSE_VECTOR = new TestDataset("dense_vector");
147+
private static final TestDataset COLORS = new TestDataset("colors");
147148

148149
public static final Map<String, TestDataset> CSV_DATASET_MAP = Map.ofEntries(
149150
Map.entry(EMPLOYEES.indexName, EMPLOYEES),
@@ -204,7 +205,8 @@ public class CsvTestsDataLoader {
204205
Map.entry(SEMANTIC_TEXT.indexName, SEMANTIC_TEXT),
205206
Map.entry(LOGS.indexName, LOGS),
206207
Map.entry(MV_TEXT.indexName, MV_TEXT),
207-
Map.entry(DENSE_VECTOR.indexName, DENSE_VECTOR)
208+
Map.entry(DENSE_VECTOR.indexName, DENSE_VECTOR),
209+
Map.entry(COLORS.indexName, COLORS)
208210
);
209211

210212
private static final EnrichConfig LANGUAGES_ENRICH = new EnrichConfig("languages_policy", "enrich-policy-languages.json");

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/data/colors.csv‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
color:keyword,hex_code:keyword,rgb_vector:dense_vector
2+
maroon, #800000, [128,0,0]
3+
dark red, #8B0000, [139,0,0]
4+
brown, #A52A2A, [165,42,42]
5+
firebrick, #B22222, [178,34,34]
6+
crimson, #DC143C, [220,20,60]
7+
red, #FF0000, [255,0,0]
8+
tomato, #FF6347, [255,99,71]
9+
coral, #FF7F50, [255,127,80]
10+
indian red, #CD5C5C, [205,92,92]
11+
light coral, #F08080, [240,128,128]
12+
dark salmon, #E9967A, [233,150,122]
13+
salmon, #FA8072, [250,128,114]
14+
light salmon, #FFA07A, [255,160,122]
15+
orange red, #FF4500, [255,69,0]
16+
dark orange, #FF8C00, [255,140,0]
17+
orange, #FFA500, [255,165,0]
18+
gold, #FFD700, [255,215,0]
19+
dark golden rod, #B8860B, [184,134,11]
20+
golden rod, #DAA520, [218,165,32]
21+
pale golden rod, #EEE8AA, [238,232,170]
22+
dark khaki, #BDB76B, [189,183,107]
23+
khaki, #F0E68C, [240,230,140]
24+
olive, #808000, [128,128,0]
25+
yellow, #FFFF00, [255,255,0]
26+
yellow green, #9ACD32, [154,205,50]
27+
dark olive green, #556B2F, [85,107,47]
28+
olive drab, #6B8E23, [107,142,35]
29+
lawn green, #7CFC00, [124,252,0]
30+
chartreuse, #7FFF00, [127,255,0]
31+
green yellow, #ADFF2F, [173,255,47]
32+
dark green, #006400, [0,100,0]
33+
green, #008000, [0,128,0]
34+
forest green, #228B22, [34,139,34]
35+
lime, #00FF00, [0,255,0]
36+
lime green, #32CD32, [50,205,50]
37+
light green, #90EE90, [144,238,144]
38+
pale green, #98FB98, [152,251,152]
39+
dark sea green, #8FBC8F, [143,188,143]
40+
medium spring green, #00FA9A, [0,250,154]
41+
spring green, #00FF7F, [0,255,127]
42+
sea green, #2E8B57, [46,139,87]
43+
medium aqua marine, #66CDAA, [102,205,170]
44+
medium sea green, #3CB371, [60,179,113]
45+
light sea green, #20B2AA, [32,178,170]
46+
dark slate gray, #2F4F4F, [47,79,79]
47+
teal, #008080, [0,128,128]
48+
dark cyan, #008B8B, [0,139,139]
49+
aqua, #00FFFF, [0,255,255]
50+
cyan, #00FFFF, [0,255,255]
51+
light cyan, #E0FFFF, [224,255,255]
52+
dark turquoise, #00CED1, [0,206,209]
53+
turquoise, #40E0D0, [64,224,208]
54+
medium turquoise, #48D1CC, [72,209,204]
55+
pale turquoise, #AFEEEE, [175,238,238]
56+
aqua marine, #7FFFD4, [127,255,212]
57+
powder blue, #B0E0E6, [176,224,230]
58+
cadet blue, #5F9EA0, [95,158,160]
59+
steel blue, #4682B4, [70,130,180]
60+
corn flower blue, #6495ED, [100,149,237]
61+
deep sky blue, #00BFFF, [0,191,255]
62+
dodger blue, #1E90FF, [30,144,255]
63+
light blue, #ADD8E6, [173,216,230]
64+
sky blue, #87CEEB, [135,206,235]
65+
light sky blue, #87CEFA, [135,206,250]
66+
midnight blue, #191970, [25,25,112]
67+
navy, #000080, [0,0,128]
68+
dark blue, #00008B, [0,0,139]
69+
medium blue, #0000CD, [0,0,205]
70+
blue, #0000FF, [0,0,255]
71+
royal blue, #4169E1, [65,105,225]
72+
blue violet, #8A2BE2, [138,43,226]
73+
indigo, #4B0082, [75,0,130]
74+
dark slate blue, #483D8B, [72,61,139]
75+
slate blue, #6A5ACD, [106,90,205]
76+
medium slate blue, #7B68EE, [123,104,238]
77+
medium purple, #9370DB, [147,112,219]
78+
dark magenta, #8B008B, [139,0,139]
79+
dark violet, #9400D3, [148,0,211]
80+
dark orchid, #9932CC, [153,50,204]
81+
medium orchid, #BA55D3, [186,85,211]
82+
purple, #800080, [128,0,128]
83+
thistle, #D8BFD8, [216,191,216]
84+
plum, #DDA0DD, [221,160,221]
85+
violet, #EE82EE, [238,130,238]
86+
magenta / fuchsia, #FF00FF, [255,0,255]
87+
orchid, #DA70D6, [218,112,214]
88+
medium violet red, #C71585, [199,21,133]
89+
pale violet red, #DB7093, [219,112,147]
90+
deep pink, #FF1493, [255,20,147]
91+
hot pink, #FF69B4, [255,105,180]
92+
light pink, #FFB6C1, [255,182,193]
93+
pink, #FFC0CB, [255,192,203]
94+
antique white, #FAEBD7, [250,235,215]
95+
beige, #F5F5DC, [245,245,220]
96+
bisque, #FFE4C4, [255,228,196]
97+
blanched almond, #FFEBCD, [255,235,205]
98+
wheat, #F5DEB3, [245,222,179]
99+
corn silk, #FFF8DC, [255,248,220]
100+
lemon chiffon, #FFFACD, [255,250,205]
101+
light golden rod yellow, #FAFAD2, [250,250,210]
102+
light yellow, #FFFFE0, [255,255,224]
103+
saddle brown, #8B4513, [139,69,19]
104+
sienna, #A0522D, [160,82,45]
105+
chocolate, #D2691E, [210,105,30]
106+
peru, #CD853F, [205,133,63]
107+
sandy brown, #F4A460, [244,164,96]
108+
burly wood, #DEB887, [222,184,135]
109+
tan, #D2B48C, [210,180,140]
110+
rosy brown, #BC8F8F, [188,143,143]
111+
moccasin, #FFE4B5, [255,228,181]
112+
navajo white, #FFDEAD, [255,222,173]
113+
peach puff, #FFDAB9, [255,218,185]
114+
misty rose, #FFE4E1, [255,228,225]
115+
lavender blush, #FFF0F5, [255,240,245]
116+
linen, #FAF0E6, [250,240,230]
117+
old lace, #FDF5E6, [253,245,230]
118+
papaya whip, #FFEFD5, [255,239,213]
119+
sea shell, #FFF5EE, [255,245,238]
120+
mint cream, #F5FFFA, [245,255,250]
121+
slate gray, #708090, [112,128,144]
122+
light slate gray, #778899, [119,136,153]
123+
light steel blue, #B0C4DE, [176,196,222]
124+
lavender, #E6E6FA, [230,230,250]
125+
floral white, #FFFAF0, [255,250,240]
126+
alice blue, #F0F8FF, [240,248,255]
127+
ghost white, #F8F8FF, [248,248,255]
128+
honeydew, #F0FFF0, [240,255,240]
129+
ivory, #FFFFF0, [255,255,240]
130+
azure, #F0FFFF, [240,255,255]
131+
snow, #FFFAFA, [255,250,250]
132+
black, #000000, [0,0,0]
133+
dim gray / dim grey, #696969, [105,105,105]
134+
gray / grey, #808080, [128,128,128]
135+
dark gray / dark grey, #A9A9A9, [169,169,169]
136+
silver, #C0C0C0, [192,192,192]
137+
light gray / light grey, #D3D3D3, [211,211,211]
138+
gainsboro, #DCDCDC, [220,220,220]
139+
white smoke, #F5F5F5, [245,245,245]
140+
white, #FFFFFF, [255,255,255]

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
knnSearch
2+
required_capability: knn_function
3+
4+
from colors metadata _score
5+
| where knn(rgb_vector, [0, 120, 0])
6+
| sort _score desc
7+
| keep color, rgb_vector, _score
8+
;
9+
10+
color:keyword | rgb_vector:dense_vector | _score:double
11+
green | [0.0, 128.0, 0.0] | 0.017461242154240608
12+
dark green | [0.0, 100.0, 0.0] | 0.002344448585063219
13+
forest green | [34.0, 139.0, 34.0] | 8.307271637022495E-4
14+
dark olive green | [85.0, 107.0, 47.0] | 1.564669219078496E-4
15+
sea green | [46.0, 139.0, 87.0] | 1.4884951815474778E-4
16+
dark slate gray | [47.0, 79.0, 79.0] | 1.3984512770548463E-4
17+
olive drab | [107.0, 142.0, 35.0] | 1.0368906077928841E-4
18+
lime green | [50.0, 205.0, 50.0] | 1.0350035154260695E-4
19+
black | [0.0, 0.0, 0.0] | 8.751000859774649E-5
20+
olive | [128.0, 128.0, 0.0] | 7.608329178765416E-5
21+
;
22+
23+
knnSearchWithKOption
24+
required_capability: knn_function
25+
26+
from colors metadata _score
27+
| where knn(rgb_vector, [0,255,255], {"k": 5})
28+
| sort _score desc
29+
| keep color, rgb_vector, _score
30+
;
31+
32+
color:keyword | rgb_vector:dense_vector | _score:double
33+
cyan | [0.0, 255.0, 255.0] | 1.0
34+
aqua | [0.0, 255.0, 255.0] | 1.0
35+
deep sky blue | [0.0, 191.0, 255.0] | 2.6103827985934913E-4
36+
dark turquoise | [0.0, 206.0, 209.0] | 2.463386917952448E-4
37+
turquoise | [64.0, 224.0, 208.0] | 1.8581181939225644E-4
38+
;
39+
40+
knnSearchWithSimilarityOption
41+
required_capability: knn_function
42+
43+
from colors metadata _score
44+
| where knn(rgb_vector, [255,192,203], {"similarity": 40})
45+
| sort _score desc
46+
| keep color, rgb_vector, _score
47+
;
48+
49+
color:keyword | rgb_vector:dense_vector | _score:double
50+
pink | [255.0, 192.0, 203.0] | 1.0
51+
light pink | [255.0, 182.0, 193.0] | 0.004639764316380024
52+
peach puff | [255.0, 218.0, 185.0] | 9.59702767431736E-4
53+
bisque | [255.0, 228.0, 196.0] | 7.536620832979679E-4
54+
thistle | [216.0, 191.0, 216.0] | 6.913584074936807E-4
55+
wheat | [245.0, 222.0, 179.0] | 6.44357583951205E-4
56+
;
57+
58+

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-colors.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
{
2+
"properties": {
3+
"color": {
4+
"type": "keyword"
5+
},
6+
"hex_code": {
7+
"type": "keyword"
8+
},
9+
"rgb_vector": {
10+
"type": "dense_vector",
11+
"similarity": "l2_norm"
12+
}
13+
}
14+
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Knn.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ protected Query translate(TranslatorHandler handler) {
110110
Check.notNull(fieldAttribute, "Match must have a field attribute as the first argument");
111111
String fieldName = getNameFromFieldAttribute(fieldAttribute);
112112
@SuppressWarnings("unchecked")
113-
List<Double> queryFolded = (List<Double>) query().fold(FoldContext.small() /* TODO remove me */);
113+
List<Number> queryFolded = (List<Number>) query().fold(FoldContext.small() /* TODO remove me */);
114114
float[] queryAsFloats = new float[queryFolded.size()];
115115
for (int i = 0; i < queryFolded.size(); i++) {
116116
queryAsFloats[i] = queryFolded.get(i).floatValue();

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,10 @@ public final void test() throws Throwable {
288288
"can't use KQL function in csv tests",
289289
testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.KQL_FUNCTION.capabilityName())
290290
);
291+
assumeFalse(
292+
"can't use KNN function in csv tests",
293+
testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.KNN_FUNCTION.capabilityName())
294+
);
291295
assumeFalse(
292296
"lookup join disabled for csv tests",
293297
testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.JOIN_LOOKUP_V12.capabilityName())

0 commit comments

Comments
 (0)