Fix KQL usage in in STATS .. BY (#128371)

Author: carlosdelest

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

@@ -184,7 +184,10 @@ private class ShardState {
         private final List<SegmentState> 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));
         }

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              
+;

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              
+;

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              
+;

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<List<Object>> 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();

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;
 

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);
     }