[CALCITE-6349] CoreRules.PROJECT_REDUCE_EXPRESSIONS crashes on expres…

…sion with ARRAY_REPEAT Signed-off-by: Mihai Budiu <[email protected]>
apache · Apr 9, 2024 · 3cd0f8e · 3cd0f8e
1 parent 283c1c1
commit 3cd0f8e
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 3 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
@@ -1399,7 +1399,7 @@ private static RelDataType arrayInsertReturnType(SqlOperatorBinding opBinding) {
   @LibraryOperator(libraries = {SPARK})
   public static final SqlFunction ARRAY_REPEAT =
       SqlBasicFunction.create(SqlKind.ARRAY_REPEAT,
-          ReturnTypes.TO_ARRAY,
+          ReturnTypes.TO_ARRAY.andThen(SqlTypeTransforms.TO_NULLABLE),
           OperandTypes.sequence(
               "ARRAY_REPEAT(ANY, INTEGER)",
               OperandTypes.ANY, OperandTypes.typeName(SqlTypeName.INTEGER)));

diff --git a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
@@ -2002,6 +2002,20 @@ private void checkSemiOrAntiJoinProjectTranspose(JoinRelType type) {
         .checkUnchanged();
   }
 
+  /** Test case for <a href="https://issues.apache.org/jira/projects/CALCITE/issues/CALCITE-6349">
+   * [CALCITE-6349] CoreRules.PROJECT_REDUCE_EXPRESSIONS crashes on expression
+   * with ARRAY_REPEAT</a>. */
+  @Test void testArrayRepeat() {
+    final String sql = "select array_repeat(1, null)";
+    sql(sql)
+        .withFactory(
+            t -> t.withOperatorTable(
+                opTab -> SqlLibraryOperatorTableFactory.INSTANCE.getOperatorTable(
+                    SqlLibrary.STANDARD, SqlLibrary.SPARK)))
+        .withRule(CoreRules.PROJECT_REDUCE_EXPRESSIONS)
+        .check();
+  }
+
   @Test void testDistinctCountMixed() {
     final String sql = "select deptno, count(distinct deptno, job) as cddj,\n"
         + "  sum(sal) as s\n"

diff --git a/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml b/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
@@ -1297,9 +1297,26 @@ LogicalProject(EXPR$0=[ARRAY_CONCAT(ARRAY(1, 2), ARRAY(3, 4))])
 ]]>
     </Resource>
     <Resource name="planAfter">
-      <![CDATA[
+<![CDATA[
 LogicalProject(EXPR$0=[CAST(ARRAY(1, 2, 3, 4)):INTEGER NOT NULL ARRAY NOT NULL])
   LogicalValues(tuples=[[{ 0 }]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testArrayRepeat">
+    <Resource name="sql">
+      <![CDATA[select array_repeat(1, null)]]>
+    </Resource>
+    <Resource name="planBefore">
+      <![CDATA[
+LogicalProject(EXPR$0=[ARRAY_REPEAT(1, null:DECIMAL(19, 9))])
+  LogicalValues(tuples=[[{ 0 }]])
+]]>
+    </Resource>
+    <Resource name="planAfter">
+      <![CDATA[
+LogicalProject(EXPR$0=[null:INTEGER NOT NULL ARRAY])
+  LogicalValues(tuples=[[{ 0 }]])
 ]]>
     </Resource>
   </TestCase>

diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
@@ -6820,7 +6820,7 @@ void checkRegexpExtract(SqlOperatorFixture f0, FunctionAlias functionAlias) {
     f.checkScalar("array_repeat(map[1, 'a', 2, 'b'], 2)", "[{1=a, 2=b}, {1=a, 2=b}]",
         "(INTEGER NOT NULL, CHAR(1) NOT NULL) MAP NOT NULL ARRAY NOT NULL");
     f.checkScalar("array_repeat(cast(null as integer), 2)", "[null, null]",
-        "INTEGER ARRAY NOT NULL");
+        "INTEGER ARRAY");
     // elements cast
     f.checkScalar("array_repeat(cast(1 as tinyint), 2)", "[1, 1]",
         "TINYINT NOT NULL ARRAY NOT NULL");