From 3948edb28dcc57f826fab4d56c30769668d78ff4 Mon Sep 17 00:00:00 2001
From: Clint Wylie <cwylie@apache.org>
Date: Fri, 8 Dec 2023 14:04:18 -0800
Subject: [PATCH] lift restriction of array_to_mv to only support direct column
 access

---
 .../org/apache/druid/math/expr/Function.java  |  8 ---
 .../apache/druid/math/expr/FunctionTest.java  |  7 +++
 ...yToMultiValueStringOperatorConversion.java |  2 -
 .../calcite/CalciteNestedDataQueryTest.java   | 52 +++++++++++++++++++
 4 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/math/expr/Function.java b/processing/src/main/java/org/apache/druid/math/expr/Function.java
index 61bce0b2b20a..3f90f7c39996 100644
--- a/processing/src/main/java/org/apache/druid/math/expr/Function.java
+++ b/processing/src/main/java/org/apache/druid/math/expr/Function.java
@@ -3231,14 +3231,6 @@ public ExprEval apply(List<Expr> args, Expr.ObjectBinding bindings)
     public void validateArguments(List<Expr> args)
     {
       validationHelperCheckArgumentCount(args, 1);
-      IdentifierExpr expr = args.get(0).getIdentifierExprIfIdentifierExpr();
-
-      if (expr == null) {
-        throw validationFailed(
-            "argument %s should be an identifier expression. Use array() instead",
-            args.get(0).toString()
-        );
-      }
     }
 
     @Nullable
diff --git a/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java b/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java
index b359ec27c278..e38b2ea51452 100644
--- a/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java
+++ b/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java
@@ -1163,6 +1163,13 @@ public void testArrayToMultiValueStringWithValidInputs()
     assertArrayExpr("array_to_mv(a)", new String[]{"foo", "bar", "baz", "foobar"});
     assertArrayExpr("array_to_mv(b)", new String[]{"1", "2", "3", "4", "5"});
     assertArrayExpr("array_to_mv(c)", new String[]{"3.1", "4.2", "5.3"});
+    assertArrayExpr("array_to_mv(array(y,z))", new String[]{"2", "3"});
+    // array type is determined by the first array type
+    assertArrayExpr("array_to_mv(array_concat(b,c))", new String[]{"1", "2", "3", "4", "5", "3", "4", "5"});
+    assertArrayExpr(
+        "array_to_mv(array_concat(c,b))",
+        new String[]{"3.1", "4.2", "5.3", "1.0", "2.0", "3.0", "4.0", "5.0"}
+    );
   }
 
   @Test
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/ArrayToMultiValueStringOperatorConversion.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/ArrayToMultiValueStringOperatorConversion.java
index 551df41b97a8..eb1ea63541c9 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/ArrayToMultiValueStringOperatorConversion.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/ArrayToMultiValueStringOperatorConversion.java
@@ -45,6 +45,4 @@ public ArrayToMultiValueStringOperatorConversion()
   {
     super(SQL_FUNCTION, "array_to_mv");
   }
-
-
 }
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
index 8261694df5b5..90b3932d1e6f 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
@@ -5344,6 +5344,58 @@ public void testGroupByRootSingleTypeArrayLongNullsAsMvd()
         .run();
   }
 
+  @Test
+  public void testGroupByRootSingleTypeArrayLongNullsAsMvdWithExpression()
+  {
+    cannotVectorize();
+    testBuilder()
+        .sql(
+            "SELECT "
+            + "ARRAY_TO_MV(ARRAY_CONCAT(arrayLongNulls, arrayLong)), "
+            + "SUM(cnt) "
+            + "FROM druid.arrays GROUP BY 1"
+        )
+        .queryContext(QUERY_CONTEXT_NO_STRINGIFY_ARRAY)
+        .expectedQueries(
+            ImmutableList.of(
+                GroupByQuery.builder()
+                            .setDataSource(TableDataSource.create(DATA_SOURCE_ARRAYS))
+                            .setInterval(querySegmentSpec(Filtration.eternity()))
+                            .setGranularity(Granularities.ALL)
+                            .setDimensions(
+                                dimensions(
+                                    new DefaultDimensionSpec("v0", "d0", ColumnType.STRING)
+                                )
+                            )
+                            .setVirtualColumns(expressionVirtualColumn(
+                                "v0",
+                                "array_to_mv(array_concat(\"arrayLongNulls\",\"arrayLong\"))",
+                                ColumnType.STRING
+                            ))
+                            .setAggregatorSpecs(aggregators(new LongSumAggregatorFactory("a0", "cnt")))
+                            .setContext(QUERY_CONTEXT_NO_STRINGIFY_ARRAY)
+                            .build()
+            )
+        )
+        .expectedResults(
+            // 9 isn't present in result because arrayLong rows are null in rows of arrayLongNulls that have value 9
+            ImmutableList.of(
+                new Object[]{NullHandling.defaultStringValue(), 10L},
+                new Object[]{"1", 12L},
+                new Object[]{"2", 7L},
+                new Object[]{"3", 9L},
+                new Object[]{"4", 4L}
+            )
+        )
+        .expectedSignature(
+            RowSignature.builder()
+                        .add("EXPR$0", ColumnType.STRING)
+                        .add("EXPR$1", ColumnType.LONG)
+                        .build()
+        )
+        .run();
+  }
+
   /**
    * MVD version of {@link #testGroupByRootSingleTypeArrayLongNullsFiltered()}
    * - implicit unnest since it is an mvd instead of array grouping