Handle array operandTypeChecker bug

chucheng92 · Dec 14, 2023 · 2e56171 · 2e56171
1 parent afd0def
commit 2e56171
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 5 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
@@ -1119,7 +1119,7 @@ private static RelDataType arrayReturnType(SqlOperatorBinding opBinding) {
   public static final SqlFunction ARRAY =
       SqlBasicFunction.create("ARRAY",
           SqlLibraryOperators::arrayReturnType,
-          OperandTypes.SAME_VARIADIC,
+          OperandTypes.ARRAY_FUNCTION,
           SqlFunctionCategory.SYSTEM);
 
   private static RelDataType mapReturnType(SqlOperatorBinding opBinding) {

diff --git a/core/src/main/java/org/apache/calcite/sql/type/OperandTypes.java b/core/src/main/java/org/apache/calcite/sql/type/OperandTypes.java
@@ -26,6 +26,7 @@
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.SqlOperandCountRange;
 import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlOperatorBinding;
 import org.apache.calcite.sql.SqlUtil;
 import org.apache.calcite.sql.validate.SqlValidatorScope;
 import org.apache.calcite.util.ImmutableIntList;
@@ -560,6 +561,9 @@ public static SqlOperandTypeChecker variadic(
   public static final SqlSingleOperandTypeChecker MAP =
       family(SqlTypeFamily.MAP);
 
+  public static final SqlOperandTypeChecker ARRAY_FUNCTION =
+      new ArrayFunctionOperandTypeChecker();
+
   public static final SqlOperandTypeChecker ARRAY_ELEMENT =
       new ArrayElementOperandTypeChecker();
 
@@ -1225,6 +1229,62 @@ private static class MapFromEntriesOperandTypeChecker
     }
   }
 
+  /**
+   * Operand type-checking strategy for a ARRAY function, it allows empty array.
+   *
+   * <p> The reason it overrides SameOperandTypeChecker#checkOperandTypesImpl is that it needs to handle
+   * the scenario where row/struct type and NULL exist simultaneously in array.
+   * This scenario is supported by spark array, but will be rejected by the current checkOperandTypesImpl.
+   */
+  private static class ArrayFunctionOperandTypeChecker
+      extends SameOperandTypeChecker {
+
+    ArrayFunctionOperandTypeChecker() {
+      // The args of array are non-fixed, so we set to -1 here. then operandCount
+      // can dynamically set according to the number of input args.
+      // details please see SameOperandTypeChecker#getOperandList.
+      super(-1);
+    }
+
+    @Override
+    protected boolean checkOperandTypesImpl(SqlOperatorBinding operatorBinding, boolean throwOnFailure,
+                                            @Nullable SqlCallBinding callBinding) {
+      if (throwOnFailure && callBinding == null) {
+        throw new IllegalArgumentException(
+            "callBinding must be non-null in case throwOnFailure=true");
+      }
+      int nOperandsActual = nOperands;
+      if (nOperandsActual == -1) {
+        nOperandsActual = operatorBinding.getOperandCount();
+      }
+      RelDataType[] types = new RelDataType[nOperandsActual];
+      final List<Integer> operandList =
+          getOperandList(operatorBinding.getOperandCount());
+      for (int i : operandList) {
+        types[i] = operatorBinding.getOperandType(i);
+      }
+      int prev = -1;
+      for (int i : operandList) {
+        if (prev >= 0) {
+          // we replace SqlTypeUtil.isComparable with SqlTypeUtil.leastRestrictiveForComparison
+          // to handle struct type and NULL constant.
+          // details please see: https://issues.apache.org/jira/browse/CALCITE-6163
+          RelDataType type =
+              SqlTypeUtil.leastRestrictiveForComparison(operatorBinding.getTypeFactory(), types[i], types[prev]);
+          if (type == null) {
+            if (!throwOnFailure) {
+              return false;
+            }
+            throw requireNonNull(callBinding, "callBinding").newValidationError(
+                RESOURCE.needSameTypeParameter());
+          }
+        }
+        prev = i;
+      }
+      return true;
+    }
+  }
+
   /**
    * Operand type-checking strategy for a MAP function, it allows empty map.
    */

diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
@@ -10546,10 +10546,10 @@ private static void checkArrayConcatAggFuncFails(SqlOperatorFixture t) {
         "RecordType(NULL EXPR$0, INTEGER NOT NULL EXPR$1) NOT NULL ARRAY NOT NULL");
     f2.checkScalar("array(row(1, 2))", "[{1, 2}]",
         "RecordType(INTEGER NOT NULL EXPR$0, INTEGER NOT NULL EXPR$1) NOT NULL ARRAY NOT NULL");
-    f2.checkFails("^array(row(1, 2), null)^",
-        "Parameters must be of the same type", false);
-    f2.checkFails("^array(null, row(1, 2))^",
-        "Parameters must be of the same type", false);
+    f2.checkScalar("array(row(1, 2), null)",
+        "[{1, 2}, null]", "RecordType(INTEGER EXPR$0, INTEGER EXPR$1) ARRAY NOT NULL");
+    f2.checkScalar("array(null, row(1, 2))",
+        "[null, {1, 2}]", "RecordType(INTEGER EXPR$0, INTEGER EXPR$1) ARRAY NOT NULL");
     f2.checkScalar("array(row(1, null), row(2, null))", "[{1, null}, {2, null}]",
         "RecordType(INTEGER NOT NULL EXPR$0, NULL EXPR$1) NOT NULL ARRAY NOT NULL");
     f2.checkScalar("array(row(null, 1), row(null, 2))", "[{null, 1}, {null, 2}]",
@@ -10560,6 +10560,11 @@ private static void checkArrayConcatAggFuncFails(SqlOperatorFixture t) {
         "RecordType(INTEGER EXPR$0, INTEGER EXPR$1) NOT NULL ARRAY NOT NULL");
     f2.checkScalar("array(row(1, 2), row(3, 4))", "[{1, 2}, {3, 4}]",
         "RecordType(INTEGER NOT NULL EXPR$0, INTEGER NOT NULL EXPR$1) NOT NULL ARRAY NOT NULL");
+    // checkFails
+    f2.checkFails("^array(row(1), row(2, 3))^",
+        "Parameters must be of the same type", false);
+    f2.checkFails("^array(row(1), row(2, 3), null)^",
+        "Parameters must be of the same type", false);
     // calcite default cast char type will fill extra spaces
     f2.checkScalar("array(1, 2, 'Hi')",
         "[1 , 2 , Hi]", "CHAR(2) NOT NULL ARRAY NOT NULL");