[CALCITE-5961] Type inference of ARRAY_COMPACT is incorrect

liyubin117 · Sep 19, 2023 · abb3e5b · abb3e5b
1 parent 008c553
commit abb3e5b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 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
@@ -1093,11 +1093,26 @@ private static RelDataType arrayAppendPrependReturnType(SqlOperatorBinding opBin
           SqlLibraryOperators::arrayAppendPrependReturnType,
           OperandTypes.ARRAY_ELEMENT);
 
+  @SuppressWarnings("argument.type.incompatible")
+  private static RelDataType arrayCompactReturnType(SqlOperatorBinding opBinding) {
+    final RelDataType arrayType = opBinding.collectOperandTypes().get(0);
+    if (arrayType.getSqlTypeName() == SqlTypeName.NULL) {
+      return arrayType;
+    }
+    RelDataType type = arrayType.getComponentType();
+    // force set nullable=false, and there are no side effects for 'NULL' type
+    if (type != null && type.isNullable()) {
+      type = opBinding.getTypeFactory().createTypeWithNullability(type, false);
+    }
+    requireNonNull(type, "inferred array element type");
+    return SqlTypeUtil.createArrayType(opBinding.getTypeFactory(), type, arrayType.isNullable());
+  }
+
   /** The "ARRAY_COMPACT(array)" function. */
   @LibraryOperator(libraries = {SPARK})
   public static final SqlFunction ARRAY_COMPACT =
       SqlBasicFunction.create(SqlKind.ARRAY_COMPACT,
-          ReturnTypes.ARG0_NULLABLE,
+          SqlLibraryOperators::arrayCompactReturnType,
           OperandTypes.ARRAY);
 
   /** The "ARRAY_CONCAT(array [, array]*)" function. */

diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
@@ -5970,11 +5970,15 @@ private static void checkIf(SqlOperatorFixture f) {
 
     final SqlOperatorFixture f = f0.withLibrary(SqlLibrary.SPARK);
     f.checkScalar("array_compact(array[null, 1, null, 2])", "[1, 2]",
-        "INTEGER ARRAY NOT NULL");
+        "INTEGER NOT NULL ARRAY NOT NULL");
+    f.checkScalar("array_compact(array[1, 2])", "[1, 2]",
+        "INTEGER NOT NULL ARRAY NOT NULL");
+    f.checkScalar("array_compact(array[null, 'hello', null, 'world'])", "[hello, world]",
+        "CHAR(5) NOT NULL ARRAY NOT NULL");
+    f.checkScalar("array_compact(array['hello', 'world'])", "[hello, world]",
+        "CHAR(5) NOT NULL ARRAY NOT NULL");
     f.checkScalar("array_compact(array[null])", "[]",
         "NULL ARRAY NOT NULL");
-    f.checkScalar("array_compact(array(null))", "[]",
-        "NULL ARRAY NOT NULL");
     f.checkScalar("array_compact(array())", "[]",
         "UNKNOWN NOT NULL ARRAY NOT NULL");
     f.checkNull("array_compact(null)");