Rename fields with project (#421)

Fixes #418

Renames in `project` require a small conceptual change in the local
execution planning: So far, the layout always contained one entry per
block and there was a 1:1 mapping from attributes to blocks/channels.
This assumption no longer holds as the same block might be mapped to
multiple attribute ids. E.g. the query `from test | project x = count, y
= count` creates a project operator that drops all blocks except `count`
and is associated with the mapping `Id(x) -> 0, Id(y) -> 0`.

Author: Lukas Wegmann

x-pack/plugin/esql/qa/server/src/main/resources/project.csv-spec

@@ -0,0 +1,27 @@
+projectRename
+row a = 1, b = 2 | project c = a;
+
+c:integer
+1
+;
+
+projectRenameDuplicate
+row a = 1, b = 2 | project c = a, d = a;
+
+c:integer | d:integer
+1 | 1
+;
+
+projectRenameEval
+row a = 1, b = 2 | project c = a, d = a | eval e = c + d;
+
+c:integer | d:integer | e:integer
+1 | 1 | 2
+;
+
+projectRenameEvalProject
+row a = 1, b = 2 | project c = a, d = a | eval e = c + d | project e, c, d;
+
+e:integer | c:integer | d:integer
+2 | 1 | 1
+;

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java

@@ -49,6 +49,7 @@
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.is;
 
 @Experimental
 @ESIntegTestCase.ClusterScope(scope = SUITE, numDataNodes = 1, numClientNodes = 0, supportsDedicatedMasters = false)
@@ -392,6 +393,55 @@ public void testEvalOverride() {
         }
     }
 
+    public void testProjectRename() {
+        EsqlQueryResponse results = run("from test | project x = count, y = count");
+        logger.info(results);
+        Assert.assertEquals(40, results.values().size());
+        assertThat(results.columns(), contains(new ColumnInfo("x", "long"), new ColumnInfo("y", "long")));
+        for (List<Object> values : results.values()) {
+            assertThat((Long) values.get(0), greaterThanOrEqualTo(40L));
+            assertThat(values.get(1), is(values.get(0)));
+        }
+    }
+
+    public void testProjectRenameEval() {
+        EsqlQueryResponse results = run("from test | project x = count, y = count | eval x2 = x + 1 | eval y2 = y + 2");
+        logger.info(results);
+        Assert.assertEquals(40, results.values().size());
+        assertThat(
+            results.columns(),
+            contains(new ColumnInfo("x", "long"), new ColumnInfo("y", "long"), new ColumnInfo("x2", "long"), new ColumnInfo("y2", "long"))
+        );
+        for (List<Object> values : results.values()) {
+            assertThat((Long) values.get(0), greaterThanOrEqualTo(40L));
+            assertThat(values.get(1), is(values.get(0)));
+            assertThat(values.get(2), is(((Long) values.get(0)) + 1));
+            assertThat(values.get(3), is(((Long) values.get(0)) + 2));
+        }
+    }
+
+    public void testProjectRenameEvalProject() {
+        EsqlQueryResponse results = run("from test | project x = count, y = count | eval z = x + y | project x, y, z");
+        logger.info(results);
+        Assert.assertEquals(40, results.values().size());
+        assertThat(results.columns(), contains(new ColumnInfo("x", "long"), new ColumnInfo("y", "long"), new ColumnInfo("z", "long")));
+        for (List<Object> values : results.values()) {
+            assertThat((Long) values.get(0), greaterThanOrEqualTo(40L));
+            assertThat(values.get(1), is(values.get(0)));
+            assertThat(values.get(2), is((Long) values.get(0) * 2));
+        }
+    }
+
+    public void testProjectOverride() {
+        EsqlQueryResponse results = run("from test | project count, data = count");
+        logger.info(results);
+        Assert.assertEquals(40, results.values().size());
+        assertThat(results.columns(), contains(new ColumnInfo("count", "long"), new ColumnInfo("data", "long")));
+        for (List<Object> values : results.values()) {
+            assertThat(values.get(1), is(values.get(0)));
+        }
+    }
+
     public void testRefreshSearchIdleShards() throws Exception {
         String indexName = "test_refresh";
         ElasticsearchAssertions.assertAcked(

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

@@ -74,6 +74,7 @@
 import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.Literal;
+import org.elasticsearch.xpack.ql.expression.NameId;
 import org.elasticsearch.xpack.ql.expression.NamedExpression;
 import org.elasticsearch.xpack.ql.expression.Order;
 import org.elasticsearch.xpack.ql.expression.function.aggregate.AggregateFunction;
@@ -84,6 +85,7 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -278,7 +280,7 @@ public PhysicalOperation plan(PhysicalPlan node, LocalExecutionPlanContext conte
                 }
             }
             Function<Page, Page> mapper = transformRequired ? p -> {
-                var blocks = new Block[p.getBlockCount()];
+                var blocks = new Block[mappedPosition.length];
                 for (int i = 0; i < blocks.length; i++) {
                     blocks[i] = p.getBlock(mappedPosition[i]);
                 }
@@ -344,9 +346,8 @@ public PhysicalOperation plan(PhysicalPlan node, LocalExecutionPlanContext conte
             } else {
                 throw new UnsupportedOperationException();
             }
-            Map<Object, Integer> layout = new HashMap<>();
-            layout.putAll(source.layout);
-            layout.put(namedExpression.toAttribute().id(), layout.size());
+            Map<Object, Integer> layout = new HashMap<>(source.layout);
+            layout.put(namedExpression.toAttribute().id(), nextFreeChannel(layout));
             return new PhysicalOperation(
                 new EvalOperatorFactory(evaluator, namedExpression.dataType().isRational() ? Double.TYPE : Long.TYPE),
                 layout,
@@ -368,22 +369,46 @@ public PhysicalOperation plan(PhysicalPlan node, LocalExecutionPlanContext conte
             return new PhysicalOperation(new RowOperatorFactory(obj), layout);
         } else if (node instanceof ProjectExec project) {
             var source = plan(project.child(), context);
-            Map<Object, Integer> layout = new HashMap<>();
 
-            var outputSet = project.outputSet();
-            var input = project.child().output();
-            var mask = new BitSet(input.size());
-            int layoutPos = 0;
-            for (Attribute element : input) {
-                var id = element.id();
-                var maskPosition = source.layout.get(id);
-                var keepColumn = outputSet.contains(element);
-                mask.set(maskPosition, keepColumn);
-                if (keepColumn) {
-                    layout.put(id, layoutPos++);
+            Map<Integer, Set<NameId>> inputChannelToInputIds = new HashMap<>(source.layout.size());
+            for (Map.Entry<Object, Integer> entry : source.layout.entrySet()) {
+                inputChannelToInputIds.computeIfAbsent(entry.getValue(), ignore -> new HashSet<>()).add((NameId) entry.getKey());
+            }
+
+            Map<Integer, Set<NameId>> inputChannelToOutputIds = new HashMap<>(inputChannelToInputIds.size());
+            for (NamedExpression ne : project.projections()) {
+                NameId inputId;
+                if (ne instanceof Alias a) {
+                    inputId = ((NamedExpression) a.child()).id();
+                } else {
+                    inputId = ne.id();
+                }
+                int inputChannel = source.layout.get(inputId);
+                inputChannelToOutputIds.computeIfAbsent(inputChannel, ignore -> new HashSet<>()).add(ne.id());
+            }
+
+            BitSet mask = new BitSet(inputChannelToInputIds.size());
+            Map<Object, Integer> layout = new HashMap<>(project.projections().size());
+            int outChannel = 0;
+
+            for (int inChannel = 0; inChannel < inputChannelToInputIds.size(); inChannel++) {
+                Set<NameId> outputIds = inputChannelToOutputIds.get(inChannel);
+
+                if (outputIds != null) {
+                    mask.set(inChannel);
+                    for (NameId outId : outputIds) {
+                        layout.put(outId, outChannel);
+                    }
+                    outChannel++;
                 }
             }
-            return new PhysicalOperation(new ProjectOperatorFactory(mask), layout, source);
+
+            if (mask.cardinality() == inputChannelToInputIds.size()) {
+                // all columns are retained, project operator is not needed but the layout needs to be updated
+                return new PhysicalOperation(source.operatorFactories, layout);
+            } else {
+                return new PhysicalOperation(new ProjectOperatorFactory(mask), layout, source);
+            }
         } else if (node instanceof FilterExec filter) {
             PhysicalOperation source = plan(filter.child(), context);
             return new PhysicalOperation(new FilterOperatorFactory(toEvaluator(filter.condition(), source.layout)), source.layout, source);
@@ -421,7 +446,7 @@ private PhysicalOperation planFieldExtractNode(LocalExecutionPlanContext context
         PhysicalOperation op = source;
         for (Attribute attr : fieldExtractExec.attributesToExtract()) {
             layout = new HashMap<>(layout);
-            layout.put(attr.id(), layout.size());
+            layout.put(attr.id(), nextFreeChannel(layout));
             Map<Object, Integer> previousLayout = op.layout;
 
             // Create ValuesSource object for the field to extract its values
@@ -519,6 +544,11 @@ public static class PhysicalOperation implements Describable {
             this.layout = layout;
         }
 
+        PhysicalOperation(List<OperatorFactory> operatorFactories, Map<Object, Integer> layout) {
+            this.operatorFactories.addAll(operatorFactories);
+            this.layout = layout;
+        }
+
         PhysicalOperation(OperatorFactory operatorFactory, Map<Object, Integer> layout, PhysicalOperation source) {
             this.operatorFactories.addAll(source.operatorFactories);
             this.operatorFactories.add(operatorFactory);
@@ -535,6 +565,16 @@ public String describe() {
         }
     }
 
+    private static int nextFreeChannel(Map<Object, Integer> layout) {
+        int nextChannel = 0;
+        for (int channel : layout.values()) {
+            if (channel >= nextChannel) {
+                nextChannel = channel + 1;
+            }
+        }
+        return nextChannel;
+    }
+
     /**
      * The count and type of driver parallelism.
      */