Simplify EsRelation transformation (#138100)

Author: idegtiarenko

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

@@ -1942,13 +1942,7 @@ private static LogicalPlan addGeneratedFieldsToEsRelations(LogicalPlan plan, Lis
                 }
 
                 if (missing.isEmpty() == false) {
-                    return new EsRelation(
-                        esr.source(),
-                        esr.indexPattern(),
-                        esr.indexMode(),
-                        esr.indexNameWithModes(),
-                        CollectionUtils.combine(esr.output(), missing)
-                    );
+                    return esr.withAttributes(CollectionUtils.combine(esr.output(), missing));
                 }
                 return esr;
             });

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java

@@ -193,7 +193,7 @@ private static LogicalPlan pruneColumnsInEsRelation(EsRelation esr, AttributeSet
             // it works differently as we extract all fields (other than the join key) that the EsRelation has.
             var remaining = pruneUnusedAndAddReferences(esr.output(), used);
             if (remaining != null) {
-                p = new EsRelation(esr.source(), esr.indexPattern(), esr.indexMode(), esr.indexNameWithModes(), remaining);
+                p = esr.withAttributes(remaining);
             }
         }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneUnusedIndexMode.java

@@ -26,7 +26,7 @@ public PruneUnusedIndexMode() {
     protected LogicalPlan rule(EsRelation r) {
         if (r.indexMode() == IndexMode.TIME_SERIES) {
             if (Expressions.anyMatch(r.output(), a -> MetadataAttribute.TSID_FIELD.equals(((Attribute) a).name())) == false) {
-                return new EsRelation(r.source(), r.indexPattern(), IndexMode.STANDARD, r.indexNameWithModes(), r.output());
+                return r.withIndexMode(IndexMode.STANDARD);
             }
         }
         return r;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/TranslateTimeSeriesAggregate.java

@@ -320,15 +320,9 @@ protected LogicalPlan rule(TimeSeriesAggregate aggregate, LogicalOptimizerContex
         LogicalPlan newChild = aggregate.child().transformUp(EsRelation.class, r -> {
             IndexMode indexMode = requiredTimeSeriesSource.get() ? r.indexMode() : IndexMode.STANDARD;
             if (r.output().contains(tsid.get()) == false) {
-                return new EsRelation(
-                    r.source(),
-                    r.indexPattern(),
-                    indexMode,
-                    r.indexNameWithModes(),
-                    CollectionUtils.combine(r.output(), tsid.get())
-                );
+                return r.withIndexMode(indexMode).withAttributes(CollectionUtils.combine(r.output(), tsid.get()));
             } else {
-                return new EsRelation(r.source(), r.indexPattern(), indexMode, r.indexNameWithModes(), r.output());
+                return r.withIndexMode(indexMode);
             }
         });
         final var firstPhase = new TimeSeriesAggregate(

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

@@ -196,4 +196,8 @@ public String nodeString() {
     public EsRelation withAttributes(List<Attribute> newAttributes) {
         return new EsRelation(source(), indexPattern, indexMode, indexNameWithModes, newAttributes);
     }
+
+    public EsRelation withIndexMode(IndexMode indexMode) {
+        return new EsRelation(source(), indexPattern, indexMode, indexNameWithModes, attrs);
+    }
 }

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

@@ -106,7 +106,7 @@ public static Optional<ReductionPlan> planReduceDriverTopN(
                 return r;
             }
             List<Attribute> attributes = CollectionUtils.prependToCopy(doc, r.output());
-            return new EsRelation(r.source(), r.indexPattern(), r.indexMode(), r.indexNameWithModes(), attributes);
+            return r.withAttributes(attributes);
         });
         if (withAddedDocToRelation.output().stream().noneMatch(EsQueryExec::isDocAttribute)) {
             // Defensive check: if any intermediate projects (or possibly another operator) removed the doc field, just abort this