Skip to content

Commit 494f95e

Browse files
tteofilitteofili
committed
first chicken_score PoC
1 parent 457aebe commit 494f95e

File tree

2 files changed

+407
-0
lines changed

2 files changed

+407
-0
lines changed

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

Lines changed: 294 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,294 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.plugin;
9+
10+
import org.elasticsearch.ElasticsearchException;
11+
import org.elasticsearch.action.index.IndexRequest;
12+
import org.elasticsearch.action.support.WriteRequest;
13+
import org.elasticsearch.common.settings.Settings;
14+
import org.elasticsearch.xpack.esql.VerificationException;
15+
import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase;
16+
import org.junit.Before;
17+
18+
import java.util.List;
19+
20+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
21+
import static org.hamcrest.CoreMatchers.containsString;
22+
23+
//@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE,org.elasticsearch.compute:TRACE", reason = "debug")
24+
public class ChickenScoreFunctionIT extends AbstractEsqlIntegTestCase {
25+
26+
@Before
27+
public void setupIndex() {
28+
createAndPopulateIndex();
29+
}
30+
31+
public void testSimpleChickenScoreWithWhereMatch() {
32+
var query = """
33+
FROM test METADATA _score
34+
| EVAL first_score = chicken_score(match(content, "brown"))
35+
| WHERE match(content, "fox")
36+
| KEEP id, _score, first_score
37+
| SORT id
38+
""";
39+
40+
try (var resp = run(query)) {
41+
assertColumnNames(resp.columns(), List.of("id", "_score", "first_score"));
42+
assertColumnTypes(resp.columns(), List.of("integer", "double", "double"));
43+
assertValues(resp.values(), List.of(List.of(1, 1.156558871269226, 0.2708943784236908), List.of(6, 0.9114001989364624, 0.21347221732139587)));
44+
}
45+
}
46+
47+
public void testSimpleEvalScoreWithWhereMatch() {
48+
var query = """
49+
FROM test METADATA _score
50+
| EVAL first_matching = match(content, "brown")
51+
| WHERE match(content, "fox")
52+
| KEEP id, _score, first_score
53+
| SORT id
54+
""";
55+
56+
try (var resp = run(query)) {
57+
assertColumnNames(resp.columns(), List.of("id", "_score", "first_matching"));
58+
assertColumnTypes(resp.columns(), List.of("integer", "double", "boolean"));
59+
assertValues(resp.values(), List.of(List.of(1, 1.156558871269226, true), List.of(6, 0.9114001989364624, true)));
60+
}
61+
}
62+
63+
/*public void testCombinedWhereMatch() {
64+
var query = """
65+
FROM test
66+
| WHERE match(content, "fox") AND id > 5
67+
| KEEP id
68+
| SORT id
69+
""";
70+
71+
try (var resp = run(query)) {
72+
assertColumnNames(resp.columns(), List.of("id"));
73+
assertColumnTypes(resp.columns(), List.of("integer"));
74+
assertValues(resp.values(), List.of(List.of(6)));
75+
}
76+
}
77+
78+
public void testMultipleMatch() {
79+
var query = """
80+
FROM test
81+
| WHERE match(content, "fox") AND match(content, "brown")
82+
| KEEP id
83+
| SORT id
84+
""";
85+
86+
try (var resp = run(query)) {
87+
assertColumnNames(resp.columns(), List.of("id"));
88+
assertColumnTypes(resp.columns(), List.of("integer"));
89+
assertValues(resp.values(), List.of(List.of(1), List.of(6)));
90+
}
91+
}
92+
93+
public void testMultipleWhereMatch() {
94+
var query = """
95+
FROM test
96+
| WHERE match(content, "fox") AND match(content, "brown")
97+
| EVAL summary = CONCAT("document with id: ", to_str(id), "and content: ", content)
98+
| SORT summary
99+
| LIMIT 4
100+
| WHERE match(content, "brown fox")
101+
| KEEP id
102+
""";
103+
104+
var error = expectThrows(ElasticsearchException.class, () -> run(query));
105+
assertThat(error.getMessage(), containsString("[MATCH] function cannot be used after LIMIT"));
106+
}
107+
108+
public void testNotWhereMatch() {
109+
var query = """
110+
FROM test
111+
| WHERE NOT match(content, "brown fox")
112+
| KEEP id
113+
| SORT id
114+
""";
115+
116+
try (var resp = run(query)) {
117+
assertColumnNames(resp.columns(), List.of("id"));
118+
assertColumnTypes(resp.columns(), List.of("integer"));
119+
assertValues(resp.values(), List.of(List.of(5)));
120+
}
121+
}
122+
123+
public void testWhereMatchWithScoring() {
124+
var query = """
125+
FROM test
126+
METADATA _score
127+
| WHERE match(content, "fox")
128+
| KEEP id, _score
129+
| SORT id ASC
130+
""";
131+
132+
try (var resp = run(query)) {
133+
assertColumnNames(resp.columns(), List.of("id", "_score"));
134+
assertColumnTypes(resp.columns(), List.of("integer", "double"));
135+
assertValues(resp.values(), List.of(List.of(1, 1.156558871269226), List.of(6, 0.9114001989364624)));
136+
}
137+
}
138+
139+
public void testWhereMatchWithScoringDifferentSort() {
140+
141+
var query = """
142+
FROM test
143+
METADATA _score
144+
| WHERE match(content, "fox")
145+
| KEEP id, _score
146+
| SORT id DESC
147+
""";
148+
149+
try (var resp = run(query)) {
150+
assertColumnNames(resp.columns(), List.of("id", "_score"));
151+
assertColumnTypes(resp.columns(), List.of("integer", "double"));
152+
assertValues(resp.values(), List.of(List.of(6, 0.9114001989364624), List.of(1, 1.156558871269226)));
153+
}
154+
}
155+
156+
public void testWhereMatchWithScoringSortScore() {
157+
var query = """
158+
FROM test
159+
METADATA _score
160+
| WHERE match(content, "fox")
161+
| KEEP id, _score
162+
| SORT _score DESC
163+
""";
164+
165+
try (var resp = run(query)) {
166+
assertColumnNames(resp.columns(), List.of("id", "_score"));
167+
assertColumnTypes(resp.columns(), List.of("integer", "double"));
168+
assertValues(resp.values(), List.of(List.of(1, 1.156558871269226), List.of(6, 0.9114001989364624)));
169+
}
170+
}
171+
172+
public void testWhereMatchWithScoringNoSort() {
173+
var query = """
174+
FROM test
175+
METADATA _score
176+
| WHERE match(content, "fox")
177+
| KEEP id, _score
178+
""";
179+
180+
try (var resp = run(query)) {
181+
assertColumnNames(resp.columns(), List.of("id", "_score"));
182+
assertColumnTypes(resp.columns(), List.of("integer", "double"));
183+
assertValuesInAnyOrder(resp.values(), List.of(List.of(1, 1.156558871269226), List.of(6, 0.9114001989364624)));
184+
}
185+
}
186+
187+
public void testNonExistingColumn() {
188+
var query = """
189+
FROM test
190+
| WHERE match(something, "fox")
191+
""";
192+
193+
var error = expectThrows(VerificationException.class, () -> run(query));
194+
assertThat(error.getMessage(), containsString("Unknown column [something]"));
195+
}
196+
197+
public void testWhereMatchEvalColumn() {
198+
var query = """
199+
FROM test
200+
| EVAL upper_content = to_upper(content)
201+
| WHERE match(upper_content, "FOX")
202+
| KEEP id
203+
""";
204+
205+
var error = expectThrows(VerificationException.class, () -> run(query));
206+
assertThat(
207+
error.getMessage(),
208+
containsString("[MATCH] function cannot operate on [upper_content], which is not a field from an index mapping")
209+
);
210+
}
211+
212+
public void testWhereMatchOverWrittenColumn() {
213+
var query = """
214+
FROM test
215+
| DROP content
216+
| EVAL content = CONCAT("document with ID ", to_str(id))
217+
| WHERE match(content, "document")
218+
""";
219+
220+
var error = expectThrows(VerificationException.class, () -> run(query));
221+
assertThat(
222+
error.getMessage(),
223+
containsString("[MATCH] function cannot operate on [content], which is not a field from an index mapping")
224+
);
225+
}
226+
227+
public void testWhereMatchAfterStats() {
228+
var query = """
229+
FROM test
230+
| STATS count(*)
231+
| WHERE match(content, "fox")
232+
""";
233+
234+
var error = expectThrows(VerificationException.class, () -> run(query));
235+
assertThat(error.getMessage(), containsString("Unknown column [content]"));
236+
}
237+
238+
public void testWhereMatchNotPushedDown() {
239+
var query = """
240+
FROM test
241+
| WHERE match(content, "fox") OR length(content) < 20
242+
| KEEP id
243+
| SORT id
244+
""";
245+
246+
try (var resp = run(query)) {
247+
assertColumnNames(resp.columns(), List.of("id"));
248+
assertColumnTypes(resp.columns(), List.of("integer"));
249+
assertValues(resp.values(), List.of(List.of(1), List.of(2), List.of(6)));
250+
}
251+
}
252+
253+
public void testWhereMatchWithRow() {
254+
var query = """
255+
ROW content = "a brown fox"
256+
| WHERE match(content, "fox")
257+
""";
258+
259+
var error = expectThrows(ElasticsearchException.class, () -> run(query));
260+
assertThat(
261+
error.getMessage(),
262+
containsString("line 2:15: [MATCH] function cannot operate on [content], which is not a field from an index mapping")
263+
);
264+
}
265+
266+
public void testMatchWithinEval() {
267+
var query = """
268+
FROM test
269+
| EVAL matches_query = match(content, "fox")
270+
""";
271+
272+
var error = expectThrows(VerificationException.class, () -> run(query));
273+
assertThat(error.getMessage(), containsString("[MATCH] function is only supported in WHERE commands"));
274+
}*/
275+
276+
private void createAndPopulateIndex() {
277+
var indexName = "test";
278+
var client = client().admin().indices();
279+
var CreateRequest = client.prepareCreate(indexName)
280+
.setSettings(Settings.builder().put("index.number_of_shards", 1))
281+
.setMapping("id", "type=integer", "content", "type=text");
282+
assertAcked(CreateRequest);
283+
client().prepareBulk()
284+
.add(new IndexRequest(indexName).id("1").source("id", 1, "content", "This is a brown fox"))
285+
.add(new IndexRequest(indexName).id("2").source("id", 2, "content", "This is a brown dog"))
286+
.add(new IndexRequest(indexName).id("3").source("id", 3, "content", "This dog is really brown"))
287+
.add(new IndexRequest(indexName).id("4").source("id", 4, "content", "The dog is brown but this document is very very long"))
288+
.add(new IndexRequest(indexName).id("5").source("id", 5, "content", "There is also a white cat"))
289+
.add(new IndexRequest(indexName).id("6").source("id", 6, "content", "The quick brown fox jumps over the lazy dog"))
290+
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
291+
.get();
292+
ensureYellow(indexName);
293+
}
294+
}

0 commit comments

Comments
 (0)