diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluator.java index b741534736fab..4644dd31f204f 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluator.java @@ -184,7 +184,10 @@ private class ShardState { private final List perSegmentState; ShardState(ShardConfig config) throws IOException { - weight = config.searcher.createWeight(config.query, scoreMode(), 1.0f); + // At this point, only the QueryBuilder has been rewritten into the query, but not the query itself. + // The query needs to be rewritten before creating the Weight so it can be transformed into the final Query to execute. + Query rewritten = config.searcher.rewrite(config.query); + weight = config.searcher.createWeight(rewritten, scoreMode(), 1.0f); searcher = config.searcher; perSegmentState = new ArrayList<>(Collections.nCopies(searcher.getLeafContexts().size(), null)); } diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/kql-function.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/kql-function.csv-spec index 410da34d42cbb..8f7d3446c899c 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/kql-function.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/kql-function.csv-spec @@ -268,3 +268,22 @@ r:double | author: text 4.670000076293945 | Walter Scheps 4.559999942779541 | J.R.R. Tolkien ; + +testKqlInStatsWithGroupingBy +required_capability: kql_function +required_capability: lucene_query_evaluator_query_rewrite +FROM airports +| STATS c = COUNT(*) where kql("country: United States") BY scalerank +| SORT scalerank desc +; + +c: long | scalerank: long +0 | 9 +44 | 8 +10 | 7 +28 | 6 +10 | 5 +12 | 4 +10 | 3 +15 | 2 +; diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/match-function.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/match-function.csv-spec index 04532d2a9693d..cd08bb55c0f17 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/match-function.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/match-function.csv-spec @@ -841,3 +841,22 @@ r:double | author: text 4.670000076293945 | Walter Scheps 4.559999942779541 | J.R.R. Tolkien ; + +testMatchInStatsWithGroupingBy +required_capability: match_function +required_capability: full_text_functions_in_stats_where +FROM airports +| STATS c = COUNT(*) where match(country, "United States") BY scalerank +| SORT scalerank desc +; + +c: long | scalerank: long +0 | 9 +44 | 8 +10 | 7 +28 | 6 +10 | 5 +12 | 4 +10 | 3 +15 | 2 +; diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/qstr-function.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/qstr-function.csv-spec index 9819662a8af13..90902137db230 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/qstr-function.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/qstr-function.csv-spec @@ -299,3 +299,22 @@ r:double | author: text 4.670000076293945 | Walter Scheps 4.559999942779541 | J.R.R. Tolkien ; + +testQstrInStatsWithGroupingBy +required_capability: qstr_function +required_capability: full_text_functions_in_stats_where +FROM airports +| STATS c = COUNT(*) where qstr("country: \"United States\"") BY scalerank +| SORT scalerank desc +; + +c: long | scalerank: long +0 | 9 +44 | 8 +10 | 7 +28 | 6 +10 | 5 +12 | 4 +10 | 3 +15 | 2 +; diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KqlFunctionIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KqlFunctionIT.java index 8e0a5ed2e8ce9..10a17bb05135b 100644 --- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KqlFunctionIT.java +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KqlFunctionIT.java @@ -7,6 +7,7 @@ package org.elasticsearch.xpack.esql.plugin; +import org.elasticsearch.ElasticsearchException; import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.WriteRequest; import org.elasticsearch.common.settings.Settings; @@ -16,12 +17,14 @@ import org.elasticsearch.xpack.esql.VerificationException; import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase; import org.elasticsearch.xpack.kql.KqlPlugin; +import org.hamcrest.Matchers; import org.junit.Before; import java.util.Collection; import java.util.List; import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; +import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList; import static org.hamcrest.CoreMatchers.containsString; public class KqlFunctionIT extends AbstractEsqlIntegTestCase { @@ -91,6 +94,42 @@ public void testInvalidKqlQueryLexicalError() { assertThat(error.getRootCause().getMessage(), containsString("line 1:1: extraneous input ':' ")); } + public void testKqlhWithStats() { + var errorQuery = """ + FROM test + | STATS c = count(*) BY kql("content: fox") + """; + + var error = expectThrows(ElasticsearchException.class, () -> run(errorQuery)); + assertThat(error.getMessage(), containsString("[KQL] function is only supported in WHERE and STATS commands")); + + var query = """ + FROM test + | STATS c = count(*) WHERE kql("content: fox"), d = count(*) WHERE kql("content: dog") + """; + + try (var resp = run(query)) { + assertColumnNames(resp.columns(), List.of("c", "d")); + assertColumnTypes(resp.columns(), List.of("long", "long")); + assertValues(resp.values(), List.of(List.of(4L, 4L))); + } + + query = """ + FROM test METADATA _score + | WHERE kql("content: fox") + | STATS m = max(_score), n = min(_score) + """; + + try (var resp = run(query)) { + assertColumnNames(resp.columns(), List.of("m", "n")); + assertColumnTypes(resp.columns(), List.of("double", "double")); + List> valuesList = getValuesList(resp.values()); + assertEquals(1, valuesList.size()); + assertThat((double) valuesList.get(0).get(0), Matchers.lessThan(1.0)); + assertThat((double) valuesList.get(0).get(1), Matchers.greaterThan(0.0)); + } + } + private void createAndPopulateIndex() { var indexName = "test"; var client = client().admin().indices(); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java index 4dae054c35360..046186b9635c9 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java @@ -1126,14 +1126,20 @@ public enum Cap { LOOKUP_JOIN_ON_MIXED_NUMERIC_FIELDS, /** - * Dense vector field type support + * {@link org.elasticsearch.compute.lucene.LuceneQueryEvaluator} rewrites the query before executing it in Lucene. This + * provides support for KQL in a STATS ... BY command that uses a KQL query for filter, for example. */ - DENSE_VECTOR_FIELD_TYPE(EsqlCorePlugin.DENSE_VECTOR_FEATURE_FLAG), + LUCENE_QUERY_EVALUATOR_QUERY_REWRITE, /** * Support parameters for LiMIT command. */ - PARAMETER_FOR_LIMIT; + PARAMETER_FOR_LIMIT, + + /** + * Dense vector field type support + */ + DENSE_VECTOR_FIELD_TYPE(EsqlCorePlugin.DENSE_VECTOR_FEATURE_FLAG); private final boolean enabled; diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java index 8187941c0b6f0..02590ff680b08 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java @@ -10,6 +10,7 @@ import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.util.BytesRef; import org.elasticsearch.Build; import org.elasticsearch.common.network.NetworkAddress; import org.elasticsearch.common.settings.Settings; @@ -2217,6 +2218,33 @@ public void testMatchFunctionStatisWithNonPushableCondition() { assertNull(esQuery.query()); } + public void testMatchFunctionWithStatsBy() { + String query = """ + from test + | stats count(*) where match(job_positions, "Data Scientist") by gender + """; + var analyzer = makeAnalyzer("mapping-default.json"); + var plannerOptimizer = new TestPlannerOptimizer(config, analyzer); + var plan = plannerOptimizer.plan(query); + + var limit = as(plan, LimitExec.class); + var agg = as(limit.child(), AggregateExec.class); + var grouping = as(agg.groupings().get(0), FieldAttribute.class); + assertEquals("gender", grouping.name()); + var aggregateAlias = as(agg.aggregates().get(0), Alias.class); + assertEquals("count(*) where match(job_positions, \"Data Scientist\")", aggregateAlias.name()); + var count = as(aggregateAlias.child(), Count.class); + var countFilter = as(count.filter(), Match.class); + assertEquals("Data Scientist", ((BytesRef) ((Literal) countFilter.query()).value()).utf8ToString()); + var aggregateFieldAttr = as(agg.aggregates().get(1), FieldAttribute.class); + assertEquals("gender", aggregateFieldAttr.name()); + var exchange = as(agg.child(), ExchangeExec.class); + var aggExec = as(exchange.child(), AggregateExec.class); + var fieldExtract = as(aggExec.child(), FieldExtractExec.class); + var esQuery = as(fieldExtract.child(), EsQueryExec.class); + assertNull(esQuery.query()); + } + private QueryBuilder wrapWithSingleQuery(String query, QueryBuilder inner, String fieldName, Source source) { return FilterTests.singleValueQuery(query, inner, fieldName, source); }