Move attributeSet reuse inside UnaryPlan(Exec)

Author: costin

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

@@ -10,15 +10,12 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
-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.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 
 import java.io.IOException;
-import java.util.List;
 import java.util.Objects;
 
 public class Limit extends UnaryPlan implements TelemetryAware {
@@ -78,16 +75,6 @@ public String getWriteableName() {
         return ENTRY.name;
     }
 
-    @Override
-    public List<Attribute> output() {
-        return child().output();
-    }
-
-    @Override
-    public AttributeSet outputSet() {
-        return child().outputSet();
-    }
-
     @Override
     protected NodeInfo<Limit> info() {
         return NodeInfo.create(this, Limit::new, limit, child(), duplicated);

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

@@ -14,8 +14,6 @@
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
-import org.elasticsearch.xpack.esql.core.expression.Attribute;
-import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
@@ -63,16 +61,6 @@ public String getWriteableName() {
         return ENTRY.name;
     }
 
-    @Override
-    public List<Attribute> output() {
-        return child().output();
-    }
-
-    @Override
-    public AttributeSet outputSet() {
-        return child().outputSet();
-    }
-
     @Override
     protected NodeInfo<OrderBy> info() {
         return NodeInfo.create(this, OrderBy::new, child(), order);

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

@@ -11,8 +11,6 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
-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.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -57,16 +55,6 @@ public String getWriteableName() {
         return ENTRY.name;
     }
 
-    @Override
-    public List<Attribute> output() {
-        return child().output();
-    }
-
-    @Override
-    public AttributeSet outputSet() {
-        return child().outputSet();
-    }
-
     @Override
     public boolean expressionsResolved() {
         return limit.resolved() && Resolvables.resolved(order);

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

@@ -21,6 +21,7 @@
 public abstract class UnaryPlan extends LogicalPlan {
 
     private final LogicalPlan child;
+    private AttributeSet lazyOutputSet;
 
     protected UnaryPlan(Source source, LogicalPlan child) {
         super(source, Collections.singletonList(child));
@@ -43,9 +44,12 @@ public List<Attribute> output() {
         return child.output();
     }
 
-    @Override
     public AttributeSet outputSet() {
-        return child().outputSet();
+        if (lazyOutputSet == null) {
+            List<Attribute> output = output();
+            lazyOutputSet = (output == child.output() ? child.outputSet() : new AttributeSet(output));
+        }
+        return lazyOutputSet;
     }
 
     @Override

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

@@ -10,7 +10,6 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 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.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -69,11 +68,6 @@ public List<Attribute> output() {
         return child().output();
     }
 
-    @Override
-    public AttributeSet outputSet() {
-        return child().outputSet();
-    }
-
     @Override
     public int hashCode() {
         return Objects.hash(condition, child());

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

@@ -10,15 +10,12 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-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.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 
 import java.io.IOException;
-import java.util.List;
 import java.util.Objects;
 
 public class LimitExec extends UnaryExec {
@@ -51,16 +48,6 @@ public String getWriteableName() {
         return ENTRY.name;
     }
 
-    @Override
-    public List<Attribute> output() {
-        return child().output();
-    }
-
-    @Override
-    public AttributeSet outputSet() {
-        return child().outputSet();
-    }
-
     @Override
     protected NodeInfo<? extends LimitExec> info() {
         return NodeInfo.create(this, LimitExec::new, child(), limit);

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

@@ -11,7 +11,6 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 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.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -70,16 +69,6 @@ public String getWriteableName() {
         return ENTRY.name;
     }
 
-    @Override
-    public List<Attribute> output() {
-        return child().output();
-    }
-
-    @Override
-    public AttributeSet outputSet() {
-        return child().outputSet();
-    }
-
     @Override
     protected NodeInfo<TopNExec> info() {
         return NodeInfo.create(this, TopNExec::new, child(), order, limit, estimatedRowSize);

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

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.plan.physical;
 
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 
 import java.util.Collections;
@@ -17,6 +18,7 @@
 public abstract class UnaryExec extends PhysicalPlan {
 
     private final PhysicalPlan child;
+    private AttributeSet lazyOutputSet;
 
     protected UnaryExec(Source source, PhysicalPlan child) {
         super(source, Collections.singletonList(child));
@@ -39,6 +41,16 @@ public List<Attribute> output() {
         return child.output();
     }
 
+    @Override
+    public AttributeSet outputSet() {
+        if (lazyOutputSet == null) {
+            List<Attribute> output = output();
+            lazyOutputSet = (output == child.output() ? child.outputSet() : new AttributeSet(output));
+            return lazyOutputSet;
+        }
+        return lazyOutputSet;
+    }
+
     @Override
     public int hashCode() {
         return Objects.hashCode(child());