upd

Author: idegtiarenko

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/AttributeSet.java

@@ -6,11 +6,13 @@
  */
 package org.elasticsearch.xpack.esql.core.expression;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Set;
 import java.util.Spliterator;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
@@ -184,6 +186,24 @@ public static AttributeSet of(Collection<? extends Attribute> c) {
         return new AttributeSet(mapBuilder.build());
     }
 
+    public static <T> AttributeSet of(Collection<T> sources, Function<T, Collection<Attribute>> mapper) {
+        if (sources.isEmpty()) {
+            return AttributeSet.EMPTY;
+        }
+        var size = 0;
+        var attributeSets = new ArrayList<Collection<Attribute>>(sources.size());
+        for (T source : sources) {
+            var attrs = mapper.apply(source);
+            size += attrs.size();
+            attributeSets.add(attrs);
+        }
+        var builder = AttributeSet.builder(size);
+        for (Collection<Attribute> attributeSet : attributeSets) {
+            builder.addAll(attributeSet);
+        }
+        return builder.build();
+    }
+
     public static Builder builder() {
         return new Builder(AttributeMap.builder());
     }

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/Expressions.java

@@ -105,23 +105,7 @@ public static List<Object> fold(FoldContext ctx, List<? extends Expression> exps
     }
 
     public static AttributeSet references(List<? extends Expression> exps) {
-        if (exps.isEmpty()) {
-            return AttributeSet.EMPTY;
-        }
-
-        var size = 0;
-        var references = new ArrayList<AttributeSet>(exps.size());
-        for (Expression exp : exps) {
-            var refs = exp.references();
-            size += refs.size();
-            references.add(refs);
-        }
-
-        var builder = AttributeSet.builder(size);
-        for (AttributeSet refs : references) {
-            builder.addAll(refs);
-        }
-        return builder.build();
+        return AttributeSet.of(exps, Expression::references);
     }
 
     public static String name(Expression e) {

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

@@ -48,11 +48,7 @@ public AttributeSet outputSet() {
 
     public AttributeSet inputSet() {
         if (lazyInputSet == null) {
-            List<Attribute> attrs = new ArrayList<>();
-            for (PlanType child : children()) {
-                attrs.addAll(child.output());
-            }
-            lazyInputSet = AttributeSet.of(attrs);
+            lazyInputSet = AttributeSet.of(children(), QueryPlan::output);
         }
         return lazyInputSet;
     }