Skip to content

Commit 187d810

Browse files
authored
[8.19] Fix KQL usage in in STATS .. BY (elastic#128371) (elastic#128547)
* Fix KQL usage in in STATS .. BY (elastic#128371) (cherry picked from commit 13f3864) # Conflicts: # x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java # x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java * Fix test
1 parent bad9277 commit 187d810

File tree

7 files changed

+135
-2
lines changed

7 files changed

+135
-2
lines changed

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluator.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,10 @@ private class ShardState {
183183
private final List<SegmentState> perSegmentState;
184184

185185
ShardState(ShardConfig config) throws IOException {
186-
weight = config.searcher.createWeight(config.query, scoreMode(), 1.0f);
186+
// At this point, only the QueryBuilder has been rewritten into the query, but not the query itself.
187+
// The query needs to be rewritten before creating the Weight so it can be transformed into the final Query to execute.
188+
Query rewritten = config.searcher.rewrite(config.query);
189+
weight = config.searcher.createWeight(rewritten, scoreMode(), 1.0f);
187190
searcher = config.searcher;
188191
perSegmentState = new ArrayList<>(Collections.nCopies(searcher.getLeafContexts().size(), null));
189192
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,3 +270,22 @@ r:double | author: text
270270
4.670000076293945 | Walter Scheps
271271
4.559999942779541 | J.R.R. Tolkien
272272
;
273+
274+
testKqlInStatsWithGroupingBy
275+
required_capability: kql_function
276+
required_capability: lucene_query_evaluator_query_rewrite
277+
FROM airports
278+
| STATS c = COUNT(*) where kql("country: United States") BY scalerank
279+
| SORT scalerank desc
280+
;
281+
282+
c: long | scalerank: long
283+
0 | 9
284+
44 | 8
285+
10 | 7
286+
28 | 6
287+
10 | 5
288+
12 | 4
289+
10 | 3
290+
15 | 2
291+
;

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -846,3 +846,22 @@ r:double | author: text
846846
4.670000076293945 | Walter Scheps
847847
4.559999942779541 | J.R.R. Tolkien
848848
;
849+
850+
testMatchInStatsWithGroupingBy
851+
required_capability: match_function
852+
required_capability: full_text_functions_in_stats_where
853+
FROM airports
854+
| STATS c = COUNT(*) where match(country, "United States") BY scalerank
855+
| SORT scalerank desc
856+
;
857+
858+
c: long | scalerank: long
859+
0 | 9
860+
44 | 8
861+
10 | 7
862+
28 | 6
863+
10 | 5
864+
12 | 4
865+
10 | 3
866+
15 | 2
867+
;

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,3 +300,22 @@ r:double | author: text
300300
4.670000076293945 | Walter Scheps
301301
4.559999942779541 | J.R.R. Tolkien
302302
;
303+
304+
testQstrInStatsWithGroupingBy
305+
required_capability: qstr_function
306+
required_capability: full_text_functions_in_stats_where
307+
FROM airports
308+
| STATS c = COUNT(*) where qstr("country: \"United States\"") BY scalerank
309+
| SORT scalerank desc
310+
;
311+
312+
c: long | scalerank: long
313+
0 | 9
314+
44 | 8
315+
10 | 7
316+
28 | 6
317+
10 | 5
318+
12 | 4
319+
10 | 3
320+
15 | 2
321+
;

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
package org.elasticsearch.xpack.esql.plugin;
99

10+
import org.elasticsearch.ElasticsearchException;
1011
import org.elasticsearch.action.index.IndexRequest;
1112
import org.elasticsearch.action.support.WriteRequest;
1213
import org.elasticsearch.common.settings.Settings;
@@ -16,12 +17,14 @@
1617
import org.elasticsearch.xpack.esql.VerificationException;
1718
import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase;
1819
import org.elasticsearch.xpack.kql.KqlPlugin;
20+
import org.hamcrest.Matchers;
1921
import org.junit.Before;
2022

2123
import java.util.Collection;
2224
import java.util.List;
2325

2426
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
27+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
2528
import static org.hamcrest.CoreMatchers.containsString;
2629

2730
public class KqlFunctionIT extends AbstractEsqlIntegTestCase {
@@ -91,6 +94,42 @@ public void testInvalidKqlQueryLexicalError() {
9194
assertThat(error.getRootCause().getMessage(), containsString("line 1:1: extraneous input ':' "));
9295
}
9396

97+
public void testKqlhWithStats() {
98+
var errorQuery = """
99+
FROM test
100+
| STATS c = count(*) BY kql("content: fox")
101+
""";
102+
103+
var error = expectThrows(ElasticsearchException.class, () -> run(errorQuery));
104+
assertThat(error.getMessage(), containsString("[KQL] function is only supported in WHERE and STATS commands"));
105+
106+
var query = """
107+
FROM test
108+
| STATS c = count(*) WHERE kql("content: fox"), d = count(*) WHERE kql("content: dog")
109+
""";
110+
111+
try (var resp = run(query)) {
112+
assertColumnNames(resp.columns(), List.of("c", "d"));
113+
assertColumnTypes(resp.columns(), List.of("long", "long"));
114+
assertValues(resp.values(), List.of(List.of(4L, 4L)));
115+
}
116+
117+
query = """
118+
FROM test METADATA _score
119+
| WHERE kql("content: fox")
120+
| STATS m = max(_score), n = min(_score)
121+
""";
122+
123+
try (var resp = run(query)) {
124+
assertColumnNames(resp.columns(), List.of("m", "n"));
125+
assertColumnTypes(resp.columns(), List.of("double", "double"));
126+
List<List<Object>> valuesList = getValuesList(resp.values());
127+
assertEquals(1, valuesList.size());
128+
assertThat((double) valuesList.get(0).get(0), Matchers.lessThan(1.0));
129+
assertThat((double) valuesList.get(0).get(1), Matchers.greaterThan(0.0));
130+
}
131+
}
132+
94133
private void createAndPopulateIndex() {
95134
var indexName = "test";
96135
var client = client().admin().indices();

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -882,7 +882,13 @@ public enum Cap {
882882
/**
883883
* Dense vector field type support
884884
*/
885-
DENSE_VECTOR_FIELD_TYPE(EsqlCorePlugin.DENSE_VECTOR_FEATURE_FLAG);
885+
DENSE_VECTOR_FIELD_TYPE(EsqlCorePlugin.DENSE_VECTOR_FEATURE_FLAG),
886+
887+
/**
888+
* {@link org.elasticsearch.compute.lucene.LuceneQueryEvaluator} rewrites the query before executing it in Lucene. This
889+
* provides support for KQL in a STATS ... BY command that uses a KQL query for filter, for example.
890+
*/
891+
LUCENE_QUERY_EVALUATOR_QUERY_REWRITE;
886892

887893
private final boolean enabled;
888894

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
1111

1212
import org.apache.lucene.search.IndexSearcher;
13+
import org.apache.lucene.util.BytesRef;
1314
import org.elasticsearch.common.network.NetworkAddress;
1415
import org.elasticsearch.common.settings.Settings;
1516
import org.elasticsearch.common.unit.Fuzziness;
@@ -1917,6 +1918,33 @@ public void testMatchFunctionStatisWithNonPushableCondition() {
19171918
assertNull(esQuery.query());
19181919
}
19191920

1921+
public void testMatchFunctionWithStatsBy() {
1922+
String query = """
1923+
from test
1924+
| stats count(*) where match(job_positions, "Data Scientist") by gender
1925+
""";
1926+
var analyzer = makeAnalyzer("mapping-default.json", new EnrichResolution());
1927+
var plannerOptimizer = new TestPlannerOptimizer(config, analyzer);
1928+
var plan = plannerOptimizer.plan(query);
1929+
1930+
var limit = as(plan, LimitExec.class);
1931+
var agg = as(limit.child(), AggregateExec.class);
1932+
var grouping = as(agg.groupings().get(0), FieldAttribute.class);
1933+
assertEquals("gender", grouping.name());
1934+
var aggregateAlias = as(agg.aggregates().get(0), Alias.class);
1935+
assertEquals("count(*) where match(job_positions, \"Data Scientist\")", aggregateAlias.name());
1936+
var count = as(aggregateAlias.child(), Count.class);
1937+
var countFilter = as(count.filter(), Match.class);
1938+
assertEquals("Data Scientist", ((BytesRef) ((Literal) countFilter.query()).value()).utf8ToString());
1939+
var aggregateFieldAttr = as(agg.aggregates().get(1), FieldAttribute.class);
1940+
assertEquals("gender", aggregateFieldAttr.name());
1941+
var exchange = as(agg.child(), ExchangeExec.class);
1942+
var aggExec = as(exchange.child(), AggregateExec.class);
1943+
var fieldExtract = as(aggExec.child(), FieldExtractExec.class);
1944+
var esQuery = as(fieldExtract.child(), EsQueryExec.class);
1945+
assertNull(esQuery.query());
1946+
}
1947+
19201948
private QueryBuilder wrapWithSingleQuery(String query, QueryBuilder inner, String fieldName, Source source) {
19211949
return FilterTests.singleValueQuery(query, inner, fieldName, source);
19221950
}

0 commit comments

Comments
 (0)