Update with() methods

Author: ivancea

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/SpatialDocValuesExtraction.java

@@ -20,14 +20,11 @@
 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;
 
@@ -79,7 +76,7 @@ protected PhysicalPlan rule(AggregateExec aggregate, LocalPhysicalOptimizerConte
         var foundAttributes = new HashSet<FieldAttribute>();
 
         PhysicalPlan plan = aggregate.transformDown(UnaryExec.class, exec -> {
-            if (exec instanceof AbstractAggregateExec agg) {
+            if (exec instanceof AggregateExec agg) {
                 var orderedAggregates = new ArrayList<NamedExpression>();
                 var changedAggregates = false;
                 for (NamedExpression aggExpr : agg.aggregates()) {
@@ -100,39 +97,7 @@ && allowedForDocValues(fieldAttribute, ctx.searchStats(), agg, foundAttributes))
                     }
                 }
                 if (changedAggregates) {
-                    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());
-                    };
+                    exec = agg.withAggregates(orderedAggregates);
                 }
             }
             if (exec instanceof EvalExec evalExec) {
@@ -206,7 +171,7 @@ private boolean foundField(Expression expression, Set<FieldAttribute> foundAttri
     private boolean allowedForDocValues(
         FieldAttribute fieldAttribute,
         SearchStats stats,
-        AbstractAggregateExec agg,
+        AggregateExec agg,
         Set<FieldAttribute> foundAttributes
     ) {
         if (stats.hasDocValues(fieldAttribute.fieldName()) == false) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/AbstractAggregateExec.java

@@ -113,10 +113,12 @@ public PhysicalPlan estimateRowSize(State state) {
         return Objects.equals(this.estimatedRowSize, size) ? this : withEstimatedSize(size);
     }
 
-    protected abstract AbstractAggregateExec withEstimatedSize(int estimatedRowSize);
+    public abstract AbstractAggregateExec withAggregates(List<? extends NamedExpression> aggregates);
 
     public abstract AbstractAggregateExec withMode(AggregatorMode newMode);
 
+    protected abstract AbstractAggregateExec withEstimatedSize(int estimatedRowSize);
+
     public AggregatorMode getMode() {
         return mode;
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/AggregateExec.java

@@ -64,6 +64,11 @@ public AggregateExec replaceChild(PhysicalPlan newChild) {
         return new AggregateExec(source(), newChild, groupings, aggregates, mode, intermediateAttributes, estimatedRowSize);
     }
 
+    @Override
+    public AggregateExec withAggregates(List<? extends NamedExpression> aggregates) {
+        return new AggregateExec(source(), child(), groupings, aggregates, mode, intermediateAttributes, estimatedRowSize);
+    }
+
     @Override
     public AggregateExec withMode(AggregatorMode newMode) {
         return new AggregateExec(source(), child(), groupings, aggregates, newMode, intermediateAttributes, estimatedRowSize);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/TimeSeriesAggregateExec.java

@@ -97,6 +97,21 @@ public TimeSeriesAggregateExec replaceChild(PhysicalPlan newChild) {
         );
     }
 
+    @Override
+    public TimeSeriesAggregateExec withAggregates(List<? extends NamedExpression> aggregates) {
+        return new TimeSeriesAggregateExec(
+            source(),
+            child(),
+            groupings,
+            aggregates,
+            mode,
+            intermediateAttributes,
+            estimatedRowSize,
+            timeBucket
+        );
+    }
+
+    @Override
     public TimeSeriesAggregateExec withMode(AggregatorMode newMode) {
         return new TimeSeriesAggregateExec(
             source(),
@@ -111,7 +126,7 @@ public TimeSeriesAggregateExec withMode(AggregatorMode newMode) {
     }
 
     @Override
-    protected AggregateExec withEstimatedSize(int estimatedRowSize) {
+    protected TimeSeriesAggregateExec withEstimatedSize(int estimatedRowSize) {
         return new TimeSeriesAggregateExec(
             source(),
             child(),

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/TopNAggregateExec.java

@@ -108,7 +108,7 @@ public Expression limit() {
     }
 
     @Override
-    protected TopNAggregateExec withEstimatedSize(int estimatedRowSize) {
+    public TopNAggregateExec withAggregates(List<? extends NamedExpression> aggregates) {
         return new TopNAggregateExec(
             source(),
             child(),
@@ -137,6 +137,21 @@ public TopNAggregateExec withMode(AggregatorMode newMode) {
         );
     }
 
+    @Override
+    protected TopNAggregateExec withEstimatedSize(int estimatedRowSize) {
+        return new TopNAggregateExec(
+            source(),
+            child(),
+            groupings,
+            aggregates,
+            mode,
+            intermediateAttributes,
+            estimatedRowSize,
+            order,
+            limit
+        );
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(groupings, aggregates, mode, intermediateAttributes, estimatedRowSize, order, limit, child());