WIP

Author: dnhatn

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/OutputOperator.java

@@ -11,7 +11,6 @@
 
 import java.util.List;
 import java.util.function.Consumer;
-import java.util.function.Function;
 
 /**
  * Sink operator that calls a given listener for each page received. The listener receives both the page as well as schema information,
@@ -21,15 +20,12 @@ public class OutputOperator extends SinkOperator {
 
     private final List<String> columns;
     private final Consumer<Page> pageConsumer;
-    private final Function<Page, Page> mapper;
 
-    public record OutputOperatorFactory(List<String> columns, Function<Page, Page> mapper, Consumer<Page> pageConsumer)
-        implements
-            SinkOperatorFactory {
+    public record OutputOperatorFactory(List<String> columns, Consumer<Page> pageConsumer) implements SinkOperatorFactory {
 
         @Override
         public SinkOperator get(DriverContext driverContext) {
-            return new OutputOperator(columns, mapper, pageConsumer);
+            return new OutputOperator(columns, pageConsumer);
         }
 
         @Override
@@ -38,9 +34,8 @@ public String describe() {
         }
     }
 
-    public OutputOperator(List<String> columns, Function<Page, Page> mapper, Consumer<Page> pageConsumer) {
+    public OutputOperator(List<String> columns, Consumer<Page> pageConsumer) {
         this.columns = columns;
-        this.mapper = mapper;
         this.pageConsumer = pageConsumer;
     }
 
@@ -63,7 +58,7 @@ public boolean needsInput() {
 
     @Override
     protected void doAddInput(Page page) {
-        pageConsumer.accept(mapper.apply(page));
+        pageConsumer.accept(page);
     }
 
     @Override

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/OutputOperatorTests.java

@@ -18,7 +18,7 @@
 public class OutputOperatorTests extends AnyOperatorTestCase {
     @Override
     protected Operator.OperatorFactory simple() {
-        return new OutputOperator.OutputOperatorFactory(List.of("a"), p -> p, p -> {});
+        return new OutputOperator.OutputOperatorFactory(List.of("a"), p -> {});
     }
 
     @Override
@@ -32,7 +32,7 @@ protected Matcher<String> expectedToStringOfSimple() {
     }
 
     private Operator.OperatorFactory big() {
-        return new OutputOperator.OutputOperatorFactory(IntStream.range(0, 20).mapToObj(i -> "a" + i).toList(), p -> p, p -> {});
+        return new OutputOperator.OutputOperatorFactory(IntStream.range(0, 20).mapToObj(i -> "a" + i).toList(), p -> {});
     }
 
     private String expectedDescriptionOfBig() {

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

@@ -80,7 +80,6 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.Executor;
-import java.util.function.Function;
 import java.util.stream.IntStream;
 
 /**
@@ -328,7 +327,7 @@ private void doLookup(T request, CancellableTask task, ActionListener<List<Page>
              * because the list will never grow mega large.
              */
             List<Page> collectedPages = Collections.synchronizedList(new ArrayList<>());
-            OutputOperator outputOperator = new OutputOperator(List.of(), Function.identity(), collectedPages::add);
+            OutputOperator outputOperator = new OutputOperator(List.of(), collectedPages::add);
             releasables.add(outputOperator);
             Driver driver = new Driver(
                 "enrich-lookup:" + request.sessionId,

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

@@ -32,6 +32,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneEmptyPlans;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneFilters;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneLiteralsInOrderBy;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneOverridingEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantOrderBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantSortClauses;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineFilters;
@@ -148,8 +149,9 @@ protected static Batch<LogicalPlan> substitutions() {
             new ReplaceAliasingEvalWithProject(),
             new SkipQueryOnEmptyMappings(),
             new SubstituteSpatialSurrogates(),
-            new ReplaceOrderByExpressionWithEval()
+            new ReplaceOrderByExpressionWithEval(),
             // new NormalizeAggregate(), - waits on https://github.com/elastic/elasticsearch/issues/100634
+            new PruneOverridingEval()
         );
     }
 

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.optimizer.rules.logical;
+
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.plan.logical.Eval;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.Project;
+
+import java.util.List;
+
+public final class PruneOverridingEval extends OptimizerRules.OptimizerRule<Eval> {
+    public PruneOverridingEval() {
+        super(OptimizerRules.TransformDirection.UP);
+    }
+
+    @Override
+    protected LogicalPlan rule(Eval eval) {
+        final List<Attribute> output = eval.output();
+        final List<Attribute> childOutput = eval.child().output();
+        if (output.size() < childOutput.size() + eval.fields().size()) {
+            return new Project(eval.source(), eval, output);
+        } else {
+            return eval;
+        }
+    }
+}

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

@@ -343,15 +343,8 @@ private PhysicalOperation planFieldExtractNode(FieldExtractExec fieldExtractExec
     private PhysicalOperation planOutput(OutputExec outputExec, LocalExecutionPlannerContext context) {
         PhysicalOperation source = plan(outputExec.child(), context);
         var output = outputExec.output();
-
-        return source.withSink(
-            new OutputOperatorFactory(
-                Expressions.names(output),
-                alignPageToAttributes(output, source.layout),
-                outputExec.getPageConsumer()
-            ),
-            source.layout
-        );
+        alignPageToAttributes(output, source.layout);
+        return source.withSink(new OutputOperatorFactory(Expressions.names(output), outputExec.getPageConsumer()), source.layout);
     }
 
     private static Function<Page, Page> alignPageToAttributes(List<Attribute> attrs, Layout layout) {
@@ -365,6 +358,9 @@ private static Function<Page, Page> alignPageToAttributes(List<Attribute> attrs,
             mappedPosition[++index] = layout.get(attribute.id()).channel();
             transformRequired |= mappedPosition[index] != index;
         }
+        if (transformRequired) {
+            throw new AssertionError("transform required attributes are not supported");
+        }
         Function<Page, Page> transformer = transformRequired ? p -> {
             var blocks = new Block[mappedPosition.length];
             for (int i = 0; i < blocks.length; i++) {