ESQL: Reuse child outputSet inside the plan where possible (elastic#124611) (elastic#125275)

Avoid creating outputSet between nodes that passthrough their input

Relates elastic#124395

Author: costin

docs/changelog/124611.yaml

@@ -0,0 +1,5 @@
+pr: 124611
+summary: Reuse child `outputSet` inside the plan where possible
+area: ES|QL
+type: enhancement
+issues: []

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

@@ -7,6 +7,7 @@
 package org.elasticsearch.xpack.esql.plan.logical;
 
 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;
@@ -20,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));
@@ -42,6 +44,14 @@ public List<Attribute> output() {
         return child.output();
     }
 
+    public AttributeSet outputSet() {
+        if (lazyOutputSet == null) {
+            List<Attribute> output = output();
+            lazyOutputSet = (output == child.output() ? child.outputSet() : new AttributeSet(output));
+        }
+        return lazyOutputSet;
+    }
+
     @Override
     public int hashCode() {
         return Objects.hashCode(child());

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

@@ -9,14 +9,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.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 FilterExec extends UnaryExec {
@@ -63,11 +61,6 @@ public Expression condition() {
         return condition;
     }
 
-    @Override
-    public List<Attribute> output() {
-        return child().output();
-    }
-
     @Override
     public int hashCode() {
         return Objects.hash(condition, child());

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());