[CALCITE-6008] ARRAY_AGG should returns ARRAY NULL when there's no in…

…put rows
QAQ-Jun · Jun 18, 2024 · a419a12 · a419a12
1 parent 4a1da22
commit a419a12
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 8 deletions.
diff --git a/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java b/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
@@ -677,7 +677,7 @@ static RelDataType deriveTypeSplit(SqlOperatorBinding operatorBinding,
       SqlBasicAggFunction
           .create(SqlKind.ARRAY_AGG,
               ReturnTypes.andThen(ReturnTypes::stripOrderBy,
-                  ReturnTypes.TO_ARRAY), OperandTypes.ANY)
+                  ReturnTypes.TO_ARRAY_NULLABLE), OperandTypes.ANY)
           .withFunctionType(SqlFunctionCategory.SYSTEM)
           .withSyntax(SqlSyntax.ORDERED_FUNCTION)
           .withAllowsNullTreatment(true);

diff --git a/core/src/main/java/org/apache/calcite/sql/type/ReturnTypes.java b/core/src/main/java/org/apache/calcite/sql/type/ReturnTypes.java
@@ -670,6 +670,13 @@ public static SqlCall stripSeparator(SqlCall call) {
   public static final SqlReturnTypeInference TO_ARRAY =
       ARG0.andThen(SqlTypeTransforms.TO_ARRAY);
 
+  /**
+   * Type-inference strategy whereby the result type of a call is nullable
+   * <code>ARRAY</code>.
+   */
+  public static final SqlReturnTypeInference TO_ARRAY_NULLABLE =
+      TO_ARRAY.andThen(SqlTypeTransforms.TO_NULLABLE);
+
   /**
    * Returns a MAP type.
    *

diff --git a/testkit/src/main/java/org/apache/calcite/sql/test/SqlOperatorFixture.java b/testkit/src/main/java/org/apache/calcite/sql/test/SqlOperatorFixture.java
@@ -450,6 +450,24 @@ void checkAgg(
       String[] inputValues,
       ResultChecker checker);
 
+  /**
+   * Checks that an aggregate expression returns the expected result.
+   *
+   * <p>For example, <code>checkAgg("AVG(DISTINCT x)", new String[] {"2", "3",
+   * null, "3" }, "INTEGER", isSingle([2, 3]));</code>
+   *
+   * @param expr        Aggregate expression, e.g. <code>SUM(DISTINCT x)</code>
+   * @param inputValues Array of input values, e.g. <code>["1", null,
+   *                    "2"]</code>.
+   * @param type        Expected result type
+   * @param checker     Result checker
+   */
+  void checkAgg(
+      String expr,
+      String[] inputValues,
+      String type,
+      ResultChecker checker);
+
   /**
    * Checks that an aggregate expression with multiple args returns the expected
    * result.

diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorFixtureImpl.java b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorFixtureImpl.java
@@ -194,9 +194,21 @@ void forEachQueryValidateAndThen(StringAndPos expression,
 
   @Override public void checkAgg(String expr, String[] inputValues,
       SqlTester.ResultChecker checker) {
+    checkAgg(expr, inputValues, SqlTests.ANY_TYPE_CHECKER, checker);
+  }
+
+  @Override public void checkAgg(String expr, String[] inputValues,
+      String type, SqlTester.ResultChecker checker) {
+    final SqlTester.TypeChecker typeChecker =
+        new SqlTests.StringTypeChecker(type);
+    checkAgg(expr, inputValues, typeChecker, checker);
+  }
+
+  private void checkAgg(String expr, String[] inputValues,
+      SqlTester.TypeChecker typeChecker, SqlTester.ResultChecker resultChecker) {
     String query =
         SqlTests.generateAggQuery(expr, inputValues);
-    tester.check(factory, query, SqlTests.ANY_TYPE_CHECKER, checker);
+    tester.check(factory, query, typeChecker, resultChecker);
   }
 
   @Override public void checkAggWithMultipleArgs(
@@ -214,9 +226,11 @@ void forEachQueryValidateAndThen(StringAndPos expression,
       String windowSpec,
       String type,
       SqlTester.ResultChecker resultChecker) {
+    final SqlTester.TypeChecker typeChecker =
+        new SqlTests.StringTypeChecker(type);
     String query =
         SqlTests.generateWinAggQuery(expr, windowSpec, inputValues);
-    tester.check(factory, query, SqlTests.ANY_TYPE_CHECKER, resultChecker);
+    tester.check(factory, query, typeChecker, resultChecker);
   }
 
   @Override public void checkScalar(String expression,

diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
@@ -11089,11 +11089,12 @@ private static void checkGroupConcatFuncFails(SqlOperatorFixture t) {
   }
 
   private static void checkArrayAggFunc(SqlOperatorFixture f) {
-    f.setFor(SqlLibraryOperators.ARRAY_CONCAT_AGG, VM_FENNEL, VM_JAVA);
+    f.setFor(SqlLibraryOperators.ARRAY_AGG, VM_FENNEL, VM_JAVA);
     final String[] values = {"'x'", "null", "'yz'"};
-    f.checkAgg("array_agg(x)", values, isSingle("[x, yz]"));
-    f.checkAgg("array_agg(x ignore nulls)", values, isSingle("[x, yz]"));
-    f.checkAgg("array_agg(x respect nulls)", values, isSingle("[x, yz]"));
+    f.checkAggType("array_agg(x)", "INTEGER NOT NULL ARRAY NOT NULL");
+    f.checkAgg("array_agg(x)", values, "CHAR(2) ARRAY", isSingle("[x, yz]"));
+    f.checkAgg("array_agg(x ignore nulls)", values, "CHAR(2) ARRAY", isSingle("[x, yz]"));
+    f.checkAgg("array_agg(x respect nulls)", values, "CHAR(2) ARRAY", isSingle("[x, yz]"));
     final String expectedError = "Invalid number of arguments "
         + "to function 'ARRAY_AGG'. Was expecting 1 arguments";
     f.checkAggFails("^array_agg(x,':')^", values, expectedError, false);
@@ -11104,7 +11105,7 @@ private static void checkArrayAggFunc(SqlOperatorFixture f) {
   }
 
   private static void checkArrayAggFuncFails(SqlOperatorFixture t) {
-    t.setFor(SqlLibraryOperators.ARRAY_CONCAT_AGG, VM_FENNEL, VM_JAVA);
+    t.setFor(SqlLibraryOperators.ARRAY_AGG, VM_FENNEL, VM_JAVA);
     final String[] values = {"'x'", "'y'"};
     final String expectedError = "No match found for function signature "
         + "ARRAY_AGG\\(<CHARACTER>\\)";