SQL: Fix issue with duplicate columns in SELECT (#42122)

Previously, if a column (field, scalar, alias) appeared more than once in the
SELECT list, the value was returned only once (1st appearance) in each row.

Fixes: #41811

(cherry picked from commit 097ea36)

Author: emasab

x-pack/plugin/sql/qa/src/main/resources/agg-ordering.csv-spec

@@ -0,0 +1,26 @@
+aggSumWithColumnRepeatedWithOrderAsc
+SELECT gender AS g, gender, SUM(salary) AS s3, SUM(salary), SUM(salary) AS s5 FROM test_emp GROUP BY gender ORDER BY SUM(salary);
+
+g:s  | gender:s  | s3:i  | SUM(salary):i | s5:i
+null |null       |487605 |487605         |487605
+F    |F          |1666196|1666196        |1666196
+M    |M          |2671054|2671054        |2671054
+;
+
+aggSumWithAliasWithColumnRepeatedWithOrderDesc
+SELECT gender AS g, gender, SUM(salary) AS s3, SUM(salary), SUM(salary) AS s5 FROM test_emp GROUP BY g ORDER BY s5 DESC;
+
+g:s  | gender:s  | s3:i  | SUM(salary):i | s5:i
+M    |M          |2671054|2671054        |2671054
+F    |F          |1666196|1666196        |1666196
+null |null       |487605 |487605         |487605
+;
+
+aggSumWithNumericRefWithColumnRepeatedWithOrderDesc
+SELECT gender AS g, gender, SUM(salary) AS s3, SUM(salary), SUM(salary) AS s5 FROM test_emp GROUP BY 2 ORDER BY 3 DESC;
+
+g:s  | gender:s  | s3:i  | SUM(salary):i | s5:i
+M    |M          |2671054|2671054        |2671054
+F    |F          |1666196|1666196        |1666196
+null |null       |487605 |487605         |487605
+;

x-pack/plugin/sql/qa/src/main/resources/agg.csv-spec

@@ -146,6 +146,33 @@ null           |null           |null           |null          |null           |n
 4              |4              |72             |4             |4.0            |0.0            |NaN            |NaN   
 ;
 
+aggSumWithColumnRepeated
+SELECT gender AS g, gender, SUM(salary) AS s3, SUM(salary), SUM(salary) AS s5 FROM test_emp GROUP BY gender;
+
+g:s  | gender:s  | s3:i  | SUM(salary):i | s5:i
+null |null       |487605 |487605         |487605
+F    |F          |1666196|1666196        |1666196
+M    |M          |2671054|2671054        |2671054
+;
+
+aggSumWithAliasWithColumnRepeated
+SELECT gender AS g, gender, SUM(salary) AS s3, SUM(salary), SUM(salary) AS s5 FROM test_emp GROUP BY g;
+
+g:s  | gender:s  | s3:i  | SUM(salary):i | s5:i
+null |null       |487605 |487605         |487605
+F    |F          |1666196|1666196        |1666196
+M    |M          |2671054|2671054        |2671054
+;
+
+aggSumWithNumericRefWithColumnRepeated
+SELECT gender AS g, gender, SUM(salary) AS s3, SUM(salary), SUM(salary) AS s5 FROM test_emp GROUP BY 2;
+
+g:s  | gender:s  | s3:i  | SUM(salary):i | s5:i
+null |null       |487605 |487605         |487605
+F    |F          |1666196|1666196        |1666196
+M    |M          |2671054|2671054        |2671054
+;
+
 aggByComplexCastedValue
 SELECT CONVERT(CONCAT(LTRIM(CONVERT("emp_no", SQL_VARCHAR)), LTRIM(CONVERT("languages", SQL_VARCHAR))), SQL_BIGINT) AS "TEMP"
 FROM "test_emp" GROUP BY "TEMP" ORDER BY "TEMP" LIMIT 20;

x-pack/plugin/sql/qa/src/main/resources/select.sql-spec

@@ -33,6 +33,12 @@ multipleColumnsNoAliasWithLimit
 SELECT first_name, last_name FROM "test_emp" ORDER BY emp_no LIMIT 5;
 multipleColumnWithAliasWithAndWithoutAsWithLimit
 SELECT first_name f, last_name AS l FROM "test_emp" ORDER BY emp_no LIMIT 5;
+multipleColumnNoAliasWithColumnRepeatedWithLimit
+SELECT salary, first_name, salary FROM test_emp ORDER BY salary LIMIT 3;
+multipleColumnWithAliasWithAsWithColumnRepeatedWithLimit
+SELECT salary, first_name, salary AS x FROM test_emp ORDER BY x LIMIT 3;
+multipleColumnWithAliasWithAndWithoutAsWithColumnRepeatedWithLimit
+SELECT salary, first_name, salary AS x, salary y FROM test_emp ORDER BY y LIMIT 3;
 
 //
 // SELECT constant literals with FROM

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainer.java

@@ -182,7 +182,9 @@ public BitSet columnMask(List<Attribute> columns) {
                     .innerId() : alias.id()) : null;
             for (int i = 0; i < fields.size(); i++) {
                 Tuple<FieldExtraction, ExpressionId> tuple = fields.get(i);
-                if (tuple.v2().equals(id) || (aliasId != null && tuple.v2().equals(aliasId))) {
+                // if the index is already set there is a collision,
+                // so continue searching for the other tuple with the same id
+                if (mask.get(i)==false && (tuple.v2().equals(id) || (aliasId != null && tuple.v2().equals(aliasId)))) {
                     index = i;
                     break;
                 }
@@ -479,4 +481,4 @@ public String toString() {
             throw new RuntimeException("error rendering", e);
         }
     }
-}
+}

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainerTests.java

@@ -6,15 +6,25 @@
 package org.elasticsearch.xpack.sql.querydsl.container;
 
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.sql.expression.Alias;
+import org.elasticsearch.xpack.sql.expression.Attribute;
+import org.elasticsearch.xpack.sql.expression.AttributeMap;
+import org.elasticsearch.xpack.sql.expression.FieldAttribute;
 import org.elasticsearch.xpack.sql.querydsl.query.BoolQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.MatchAll;
 import org.elasticsearch.xpack.sql.querydsl.query.NestedQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.Query;
 import org.elasticsearch.xpack.sql.querydsl.query.RangeQuery;
 import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.tree.SourceTests;
+import org.elasticsearch.xpack.sql.type.DataType;
+import org.elasticsearch.xpack.sql.type.EsField;
 
 import java.util.AbstractMap.SimpleImmutableEntry;
+import java.util.Arrays;
+import java.util.BitSet;
+import java.util.LinkedHashMap;
+import java.util.Map;
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
@@ -60,4 +70,48 @@ public void testRewriteToContainsNestedFieldWhenDoesNotContainNestedFieldAndCant
                     new MatchAll(source)));
         assertEquals(expected, QueryContainer.rewriteToContainNestedField(original, source, path, name, format, hasDocValues));
     }
+
+    public void testColumnMaskShouldDuplicateSameAttributes() {
+
+        EsField esField = new EsField("str", DataType.TEXT, emptyMap(), true);
+
+        Attribute first = new FieldAttribute(Source.EMPTY, "first", esField);
+        Attribute second = new FieldAttribute(Source.EMPTY, "second", esField);
+        Attribute third = new FieldAttribute(Source.EMPTY, "third", esField);
+        Attribute fourth = new FieldAttribute(Source.EMPTY, "fourth", esField);
+        Alias firstAliased = new Alias(Source.EMPTY, "firstAliased", first);
+
+        Map<Attribute,Attribute> aliasesMap = new LinkedHashMap<>();
+        aliasesMap.put(firstAliased.toAttribute(), first);
+
+        QueryContainer queryContainer = new QueryContainer()
+            .withAliases(new AttributeMap<>(aliasesMap))
+            .addColumn(third)
+            .addColumn(first)
+            .addColumn(fourth)
+            .addColumn(firstAliased.toAttribute())
+            .addColumn(second)
+            .addColumn(first)
+            .addColumn(fourth);
+
+        BitSet result = queryContainer.columnMask(Arrays.asList(
+            first,
+            first,
+            second,
+            third,
+            firstAliased.toAttribute()
+        ));
+
+        BitSet expected = new BitSet();
+        expected.set(0, true);
+        expected.set(1, true);
+        expected.set(2, false);
+        expected.set(3, true);
+        expected.set(4, true);
+        expected.set(5, true);
+        expected.set(6, false);
+
+
+        assertEquals(expected, result);
+    }
 }