Updated Physicial rules and fixed Spatial rule for time series

ivancea · ivancea · commit d0677c0ea381 · 2025-06-11T17:57:10.000+02:00
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/InsertFieldExtraction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/InsertFieldExtraction.java
@@ -14,6 +14,7 @@
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext;
 import org.elasticsearch.xpack.esql.optimizer.PhysicalOptimizerRules;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.ProjectAwayColumns;
+import org.elasticsearch.xpack.esql.plan.physical.AbstractAggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec;
 import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
@@ -53,7 +54,7 @@ public PhysicalPlan rule(PhysicalPlan plan, LocalPhysicalOptimizerContext contex
              * it loads the field lazily. If we have more than one field we need to
              * make sure the fields are loaded for the standard hash aggregator.
              */
-            if (p instanceof AggregateExec agg) {
+            if (p instanceof AbstractAggregateExec agg) {
                 var ordinalAttributes = agg.ordinalAttributes();
                 missing.removeAll(Expressions.references(ordinalAttributes));
             }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/SpatialDocValuesExtraction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/SpatialDocValuesExtraction.java
@@ -20,11 +20,14 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.SpatialRelatesFunction;
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext;
 import org.elasticsearch.xpack.esql.optimizer.PhysicalOptimizerRules;
+import org.elasticsearch.xpack.esql.plan.physical.AbstractAggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EvalExec;
 import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
 import org.elasticsearch.xpack.esql.plan.physical.FilterExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
+import org.elasticsearch.xpack.esql.plan.physical.TimeSeriesAggregateExec;
+import org.elasticsearch.xpack.esql.plan.physical.TopNAggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.UnaryExec;
 import org.elasticsearch.xpack.esql.stats.SearchStats;
 
@@ -76,7 +79,7 @@ protected PhysicalPlan rule(AggregateExec aggregate, LocalPhysicalOptimizerConte
         var foundAttributes = new HashSet<FieldAttribute>();
 
         PhysicalPlan plan = aggregate.transformDown(UnaryExec.class, exec -> {
-            if (exec instanceof AggregateExec agg) {
+            if (exec instanceof AbstractAggregateExec agg) {
                 var orderedAggregates = new ArrayList<NamedExpression>();
                 var changedAggregates = false;
                 for (NamedExpression aggExpr : agg.aggregates()) {
@@ -97,15 +100,39 @@ && allowedForDocValues(fieldAttribute, ctx.searchStats(), agg, foundAttributes))
                     }
                 }
                 if (changedAggregates) {
-                    exec = new AggregateExec(
-                        agg.source(),
-                        agg.child(),
-                        agg.groupings(),
-                        orderedAggregates,
-                        agg.getMode(),
-                        agg.intermediateAttributes(),
-                        agg.estimatedRowSize()
-                    );
+                    exec = switch (agg) {
+                        case TimeSeriesAggregateExec tsAggExec -> new TimeSeriesAggregateExec(
+                            agg.source(),
+                            agg.child(),
+                            agg.groupings(),
+                            orderedAggregates,
+                            agg.getMode(),
+                            agg.intermediateAttributes(),
+                            agg.estimatedRowSize(),
+                            tsAggExec.timeBucket()
+                        );
+                        case AggregateExec aggExec -> new AggregateExec(
+                            agg.source(),
+                            agg.child(),
+                            agg.groupings(),
+                            orderedAggregates,
+                            agg.getMode(),
+                            agg.intermediateAttributes(),
+                            agg.estimatedRowSize()
+                        );
+                        case TopNAggregateExec topNAggExec -> new TopNAggregateExec(
+                            agg.source(),
+                            agg.child(),
+                            agg.groupings(),
+                            orderedAggregates,
+                            agg.getMode(),
+                            agg.intermediateAttributes(),
+                            agg.estimatedRowSize(),
+                            topNAggExec.order(),
+                            topNAggExec.limit()
+                        );
+                        default -> throw new IllegalStateException("Unexpected aggregate type: " + agg.getClass().getName());
+                    };
                 }
             }
             if (exec instanceof EvalExec evalExec) {
@@ -179,7 +206,7 @@ private boolean foundField(Expression expression, Set<FieldAttribute> foundAttri
     private boolean allowedForDocValues(
         FieldAttribute fieldAttribute,
         SearchStats stats,
-        AggregateExec agg,
+        AbstractAggregateExec agg,
         Set<FieldAttribute> foundAttributes
     ) {
         if (stats.hasDocValues(fieldAttribute.fieldName()) == false) {
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/AbstractAggregateExec.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/AbstractAggregateExec.java
@@ -115,6 +115,8 @@ public PhysicalPlan estimateRowSize(State state) {
 
     protected abstract AbstractAggregateExec withEstimatedSize(int estimatedRowSize);
 
+    public abstract AbstractAggregateExec withMode(AggregatorMode newMode);
+
     public AggregatorMode getMode() {
         return mode;
     }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/AggregateExec.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/AggregateExec.java
@@ -65,6 +65,7 @@ public AggregateExec replaceChild(PhysicalPlan newChild) {
         return new AggregateExec(source(), newChild, groupings, aggregates, mode, intermediateAttributes, estimatedRowSize);
     }
 
+    @Override
     public AggregateExec withMode(AggregatorMode newMode) {
         return new AggregateExec(source(), child(), groupings, aggregates, newMode, intermediateAttributes, estimatedRowSize);
     }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/TopNAggregateExec.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/TopNAggregateExec.java
@@ -103,6 +103,11 @@ protected TopNAggregateExec withEstimatedSize(int estimatedRowSize) {
         );
     }
 
+    @Override
+    public TopNAggregateExec withMode(AggregatorMode newMode) {
+        return new TopNAggregateExec(source(), child(), groupings, aggregates, newMode, intermediateAttributes, estimatedRowSize, order, limit);
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(groupings, aggregates, mode, intermediateAttributes, estimatedRowSize, order, limit, child());
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java
@@ -38,6 +38,7 @@
 import org.elasticsearch.xpack.esql.plan.QueryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
+import org.elasticsearch.xpack.esql.plan.physical.AbstractAggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.EstimatesRowSize;
@@ -123,7 +124,7 @@ public static PhysicalPlan reductionPlan(PhysicalPlan plan) {
         final var pipelineBreaker = pipelineBreakers.getFirst();
         final LocalMapper mapper = new LocalMapper();
         PhysicalPlan reducePlan = mapper.map(pipelineBreaker);
-        if (reducePlan instanceof AggregateExec agg) {
+        if (reducePlan instanceof AbstractAggregateExec agg) {
             reducePlan = agg.withMode(AggregatorMode.INTERMEDIATE);
         }
         return EstimatesRowSize.estimateRowSize(fragment.estimatedRowSize(), reducePlan);

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,8 @@ public PhysicalPlan estimateRowSize(State state) {`
`115`	`115`
`116`	`116`	`protected abstract AbstractAggregateExec withEstimatedSize(int estimatedRowSize);`
`117`	`117`
	`118`	`+ public abstract AbstractAggregateExec withMode(AggregatorMode newMode);`
	`119`	`+`
`118`	`120`	`public AggregatorMode getMode() {`
`119`	`121`	`return mode;`
`120`	`122`	`}`
Original file line number	Diff line number	Diff line change
`@@ -65,6 +65,7 @@ public AggregateExec replaceChild(PhysicalPlan newChild) {`
`65`	`65`	`return new AggregateExec(source(), newChild, groupings, aggregates, mode, intermediateAttributes, estimatedRowSize);`
`66`	`66`	`}`
`67`	`67`
	`68`	`+ @Override`
`68`	`69`	`public AggregateExec withMode(AggregatorMode newMode) {`
`69`	`70`	`return new AggregateExec(source(), child(), groupings, aggregates, newMode, intermediateAttributes, estimatedRowSize);`
`70`	`71`	`}`