ESQL: INLINESTATS

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -69,6 +69,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.EsqlAggregate;
 import org.elasticsearch.xpack.esql.plan.logical.EsqlUnresolvedRelation;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
+import org.elasticsearch.xpack.esql.plan.logical.InlineStats;
 import org.elasticsearch.xpack.esql.plan.logical.Keep;
 import org.elasticsearch.xpack.esql.plan.logical.Lookup;
 import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
@@ -333,6 +334,9 @@ private static class ResolveLookupTables extends ParameterizedAnalyzerRule<Looku
 
         @Override
         protected LogicalPlan rule(Lookup lookup, AnalyzerContext context) {
+            if (lookup.localRelation() != null) {
+                return lookup;
+            }
             // the parser passes the string wrapped in a literal
             Source source = lookup.source();
             Expression tableNameExpression = lookup.tableName();
@@ -394,6 +398,10 @@ protected LogicalPlan doRule(LogicalPlan plan) {
                 return resolveAggregate(agg, childrenOutput);
             }
 
+            if (plan instanceof InlineStats stats) {
+                return resolveInlineStats(stats, childrenOutput);
+            }
+
             if (plan instanceof Drop d) {
                 return resolveDrop(d, childrenOutput);
             }
@@ -478,6 +486,59 @@ private LogicalPlan resolveAggregate(Aggregate a, List<Attribute> childrenOutput
             return a;
         }
 
+        private LogicalPlan resolveInlineStats(InlineStats stats, List<Attribute> childrenOutput) {
+            // if the grouping is resolved but the aggs are not, use the former to resolve the latter
+            // e.g. STATS a ... GROUP BY a = x + 1
+            Holder<Boolean> changed = new Holder<>(false);
+            List<Expression> groupings = stats.groupings();
+            // first resolve groupings since the aggs might refer to them
+            // trying to globally resolve unresolved attributes will lead to some being marked as unresolvable
+            if (Resolvables.resolved(groupings) == false) {
+                List<Expression> newGroupings = new ArrayList<>(groupings.size());
+                for (Expression g : groupings) {
+                    Expression resolved = g.transformUp(UnresolvedAttribute.class, ua -> maybeResolveAttribute(ua, childrenOutput));
+                    if (resolved != g) {
+                        changed.set(true);
+                    }
+                    newGroupings.add(resolved);
+                }
+                groupings = newGroupings;
+                if (changed.get()) {
+                    stats = new InlineStats(stats.source(), stats.child(), newGroupings, stats.aggregates());
+                    changed.set(false);
+                }
+            }
+
+            if (stats.expressionsResolved() == false) {
+                AttributeMap<Expression> resolved = new AttributeMap<>();
+                for (Expression e : groupings) {
+                    Attribute attr = Expressions.attribute(e);
+                    if (attr != null && attr.resolved()) {
+                        resolved.put(attr, attr);
+                    }
+                }
+                List<Attribute> resolvedList = NamedExpressions.mergeOutputAttributes(new ArrayList<>(resolved.keySet()), childrenOutput);
+                List<NamedExpression> newAggregates = new ArrayList<>();
+
+                for (NamedExpression aggregate : stats.aggregates()) {
+                    var agg = (NamedExpression) aggregate.transformUp(UnresolvedAttribute.class, ua -> {
+                        Expression ne = ua;
+                        Attribute maybeResolved = maybeResolveAttribute(ua, resolvedList);
+                        if (maybeResolved != null) {
+                            changed.set(true);
+                            ne = maybeResolved;
+                        }
+                        return ne;
+                    });
+                    newAggregates.add(agg);
+                }
+
+                stats = changed.get() ? new InlineStats(stats.source(), stats.child(), groupings, newAggregates) : stats;
+            }
+
+            return stats;
+        }
+
         private LogicalPlan resolveMvExpand(MvExpand p, List<Attribute> childrenOutput) {
             if (p.target() instanceof UnresolvedAttribute ua) {
                 Attribute resolved = maybeResolveAttribute(ua, childrenOutput);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java

@@ -19,12 +19,15 @@
 import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.planner.Mapper;
+import org.elasticsearch.xpack.esql.plugin.ComputeService;
 import org.elasticsearch.xpack.esql.session.EsqlConfiguration;
 import org.elasticsearch.xpack.esql.session.EsqlIndexResolver;
 import org.elasticsearch.xpack.esql.session.EsqlSession;
 import org.elasticsearch.xpack.esql.stats.Metrics;
 import org.elasticsearch.xpack.esql.stats.QueryMetric;
 
+import java.util.function.BiConsumer;
+
 import static org.elasticsearch.action.ActionListener.wrap;
 
 public class PlanExecutor {
@@ -52,7 +55,8 @@ public void esql(
         String sessionId,
         EsqlConfiguration cfg,
         EnrichPolicyResolver enrichPolicyResolver,
-        ActionListener<PhysicalPlan> listener
+        BiConsumer<PhysicalPlan, ActionListener<EsqlSession.Result>> runPhase,
+        ActionListener<EsqlSession.Result> listener
     ) {
         final var session = new EsqlSession(
             sessionId,
@@ -68,7 +72,7 @@ public void esql(
         );
         QueryMetric clientId = QueryMetric.fromString("rest");
         metrics.total(clientId);
-        session.execute(request, wrap(listener::onResponse, ex -> {
+        session.execute(request, runPhase, wrap(listener::onResponse, ex -> {
             // TODO when we decide if we will differentiate Kibana from REST, this String value will likely come from the request
             metrics.failed(clientId);
             listener.onFailure(ex);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/InlineStats.java

@@ -7,23 +7,40 @@
 
 package org.elasticsearch.xpack.esql.plan.logical;
 
+import org.elasticsearch.common.collect.Iterators;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
+import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
+import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.core.plan.logical.UnaryPlan;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.core.type.EsField;
+import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
+import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier;
+import org.elasticsearch.xpack.esql.planner.PlannerUtils;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
+import java.util.stream.IntStream;
 
-public class InlineStats extends UnaryPlan {
+import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
+
+public class InlineStats extends UnaryPlan implements Phased {
 
     private final List<Expression> groupings;
     private final List<? extends NamedExpression> aggregates;
+    private List<Attribute> output;
 
     public InlineStats(Source source, LogicalPlan child, List<Expression> groupings, List<? extends NamedExpression> aggregates) {
         super(source, child);
@@ -56,7 +73,34 @@ public boolean expressionsResolved() {
 
     @Override
     public List<Attribute> output() {
-        return Expressions.asAttributes(aggregates);
+        // NOCOMMIT when is this one called?
+        if (this.output == null) {
+            this.output = mergeOutputAttributes(Expressions.asAttributes(aggregates), child().output());
+        }
+        return output;
+    }
+
+    @Override
+    public LogicalPlan firstPhase() {
+        return new Aggregate(source(), child(), groupings, aggregates);
+    }
+
+    @Override
+    public LogicalPlan nextPhase(List<Attribute> layout, List<Page> firstPhaseResult) {
+        // NOCOMMIT memory tracking
+        if (firstPhaseResult.size() > 1) {
+            throw new UnsupportedOperationException();
+        }
+        List<NamedExpression> namedGroupings = groupings.stream().map(Expressions::wrapAsNamed).toList();
+        Page page = firstPhaseResult.get(0);
+        Block[] blocks = IntStream.range(0, page.getBlockCount()).mapToObj(b -> {
+            Block block = page.getBlock(b);
+            Block.Builder builder = block.elementType().newBlockBuilder(block.getPositionCount(), PlannerUtils.NON_BREAKING_BLOCK_FACTORY);
+            builder.copyFrom(block, 0, block.getPositionCount());
+            return builder.build();
+        }).toArray(Block[]::new);
+        LocalRelation local = new LocalRelation(source(), layout, LocalSupplier.of(blocks));
+        return new Lookup(source(), child(), new ReferenceAttribute(source(), "unused", DataType.KEYWORD), namedGroupings, local);
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Phased.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.plan.logical;
+
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.core.util.Holder;
+
+import java.util.List;
+
+public interface Phased {
+    LogicalPlan firstPhase();
+    LogicalPlan nextPhase(List<Attribute> layout, List<Page> firstPhaseResult);
+
+    static LogicalPlan extractNextPhase(LogicalPlan plan) {
+        var firstPhase = new Holder<LogicalPlan>();
+        plan.forEachUp(t -> {
+            if (t instanceof Phased phased) {
+                firstPhase.set(phased.firstPhase());
+            }
+        });
+        return firstPhase.get();
+    }
+
+    static LogicalPlan applyResultsFromNextPhase(LogicalPlan plan, List<Attribute> layout, List<Page> result) {
+        return plan.transformUp(logicalPlan -> {
+            // NOCOMMIT make sure this stops after the first one.
+            if (logicalPlan instanceof Phased phased) {
+                return phased.nextPhase(layout, result);
+            }
+            return logicalPlan;
+        });
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/Mapper.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.Grok;
+import org.elasticsearch.xpack.esql.plan.logical.InlineStats;
 import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.Row;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java

@@ -12,11 +12,13 @@
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.ElementType;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
@@ -36,6 +38,7 @@
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
+import org.elasticsearch.xpack.esql.plan.logical.Phased;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsSourceExec;

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

@@ -72,6 +72,7 @@
 import org.elasticsearch.xpack.esql.planner.LocalExecutionPlanner;
 import org.elasticsearch.xpack.esql.planner.PlannerUtils;
 import org.elasticsearch.xpack.esql.session.EsqlConfiguration;
+import org.elasticsearch.xpack.esql.session.EsqlSession;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -89,8 +90,6 @@
  * Computes the result of a {@link PhysicalPlan}.
  */
 public class ComputeService {
-    public record Result(List<Page> pages, List<DriverProfile> profiles) {}
-
     private static final Logger LOGGER = LogManager.getLogger(ComputeService.class);
     private final SearchService searchService;
     private final BigArrays bigArrays;
@@ -136,7 +135,7 @@ public void execute(
         CancellableTask rootTask,
         PhysicalPlan physicalPlan,
         EsqlConfiguration configuration,
-        ActionListener<Result> listener
+        ActionListener<EsqlSession.Result> listener
     ) {
         Tuple<PhysicalPlan, PhysicalPlan> coordinatorAndDataNodePlan = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(
             physicalPlan,
@@ -176,7 +175,7 @@ public void execute(
                 rootTask,
                 computeContext,
                 coordinatorPlan,
-                listener.map(driverProfiles -> new Result(collectedPages, driverProfiles))
+                listener.map(driverProfiles -> new EsqlSession.Result(collectedPages, physicalPlan.output(), driverProfiles))
             );
             return;
         } else {
@@ -201,7 +200,9 @@ public void execute(
         );
         try (
             Releasable ignored = exchangeSource.addEmptySink();
-            RefCountingListener refs = new RefCountingListener(listener.map(unused -> new Result(collectedPages, collectedProfiles)))
+            RefCountingListener refs = new RefCountingListener(
+                listener.map(unused -> new EsqlSession.Result(collectedPages, physicalPlan.output(), collectedProfiles))
+            )
         ) {
             // run compute on the coordinator
             exchangeSource.addCompletionListener(refs.acquire());

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

@@ -158,35 +158,21 @@ private void innerExecute(Task task, EsqlQueryRequest request, ActionListener<Es
             request.tables()
         );
         String sessionId = sessionID(task);
-        planExecutor.esql(
-            request,
-            sessionId,
-            configuration,
-            enrichPolicyResolver,
-            listener.delegateFailureAndWrap(
-                (delegate, physicalPlan) -> computeService.execute(
-                    sessionId,
-                    (CancellableTask) task,
-                    physicalPlan,
-                    configuration,
-                    delegate.map(result -> {
-                        List<ColumnInfo> columns = physicalPlan.output()
-                            .stream()
-                            .map(c -> new ColumnInfo(c.qualifiedName(), EsqlDataTypes.outputType(c.dataType())))
-                            .toList();
-                        EsqlQueryResponse.Profile profile = configuration.profile()
-                            ? new EsqlQueryResponse.Profile(result.profiles())
-                            : null;
-                        if (task instanceof EsqlQueryTask asyncTask && request.keepOnCompletion()) {
-                            String id = asyncTask.getExecutionId().getEncoded();
-                            return new EsqlQueryResponse(columns, result.pages(), profile, request.columnar(), id, false, request.async());
-                        } else {
-                            return new EsqlQueryResponse(columns, result.pages(), profile, request.columnar(), request.async());
-                        }
-                    })
-                )
-            )
-        );
+        planExecutor.esql(request, sessionId, configuration, enrichPolicyResolver, (physicalPlan, resultListener) -> {
+            computeService.execute(sessionId, (CancellableTask) task, physicalPlan, configuration, resultListener);
+        }, listener.map(result -> {
+            List<ColumnInfo> columns = result.layout()
+                .stream()
+                .map(c -> new ColumnInfo(c.qualifiedName(), EsqlDataTypes.outputType(c.dataType())))
+                .toList();
+            EsqlQueryResponse.Profile profile = configuration.profile() ? new EsqlQueryResponse.Profile(result.profiles()) : null;
+            if (task instanceof EsqlQueryTask asyncTask && request.keepOnCompletion()) {
+                String id = asyncTask.getExecutionId().getEncoded();
+                return new EsqlQueryResponse(columns, result.pages(), profile, request.columnar(), id, false, request.async());
+            } else {
+                return new EsqlQueryResponse(columns, result.pages(), profile, request.columnar(), request.async());
+            }
+        }));
     }
 
     /**