[CALCITE-6208] Update JSON_VALUE return type inference to make explic…

…it array return types be nullable with nullable elements
QAQ-Jun · Jan 30, 2024 · de39888 · de39888
1 parent 75511b8
commit de39888
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/core/src/main/java/org/apache/calcite/sql/fun/SqlJsonValueFunction.java b/core/src/main/java/org/apache/calcite/sql/fun/SqlJsonValueFunction.java
@@ -30,11 +30,10 @@
 import org.apache.calcite.sql.SqlOperatorBinding;
 import org.apache.calcite.sql.SqlWriter;
 import org.apache.calcite.sql.type.OperandTypes;
-import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlOperandCountRanges;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
-import org.apache.calcite.sql.type.SqlTypeTransforms;
+import org.apache.calcite.sql.type.SqlTypeUtil;
 
 import com.google.common.collect.ImmutableList;
 
@@ -54,20 +53,33 @@ public class SqlJsonValueFunction extends SqlFunction {
 
   public SqlJsonValueFunction(String name) {
     super(name, SqlKind.OTHER_FUNCTION,
-        ReturnTypes.cascade(
-            opBinding -> explicitTypeSpec(opBinding).orElse(getDefaultType(opBinding)),
-            SqlTypeTransforms.FORCE_NULLABLE),
+        opBinding -> explicitTypeSpec(opBinding)
+            .map(relDataType -> deriveExplicitType(opBinding, relDataType))
+            .orElseGet(() -> getDefaultType(opBinding)),
         null,
         OperandTypes.family(
             ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER),
             ordinal -> ordinal > 1),
         SqlFunctionCategory.SYSTEM);
   }
 
+  private static RelDataType deriveExplicitType(SqlOperatorBinding opBinding, RelDataType type) {
+    if (SqlTypeName.ARRAY == type.getSqlTypeName()) {
+      RelDataType elementType = Objects.requireNonNull(type.getComponentType());
+      RelDataType nullableElementType = deriveExplicitType(opBinding, elementType);
+      return SqlTypeUtil.createArrayType(
+          opBinding.getTypeFactory(),
+          nullableElementType,
+          true);
+    }
+    return opBinding.getTypeFactory().createTypeWithNullability(type, true);
+  }
+
   /** Returns VARCHAR(2000) as default. */
   private static RelDataType getDefaultType(SqlOperatorBinding opBinding) {
     final RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
-    return typeFactory.createSqlType(SqlTypeName.VARCHAR, 2000);
+    final RelDataType baseType = typeFactory.createSqlType(SqlTypeName.VARCHAR, 2000);
+    return typeFactory.createTypeWithNullability(baseType, true);
   }
 
   /**

diff --git a/core/src/test/java/org/apache/calcite/test/SqlValidatorTest.java b/core/src/test/java/org/apache/calcite/test/SqlValidatorTest.java
@@ -11417,6 +11417,14 @@ private void checkCustomColumnResolving(String table) {
     expr("json_value('{\"foo\":100}', 'lax $.foo' returning boolean"
         + " default 100 on empty)")
         .columnType("BOOLEAN");
+
+    expr("json_value('{\"foo\":[100, null, 200]}', 'lax $.foo'"
+        + "returning integer array)")
+        .columnType("INTEGER ARRAY");
+
+    expr("json_value('{\"foo\":[[100, null, 200]]}', 'lax $.foo'"
+        + "returning integer array array)")
+        .columnType("INTEGER ARRAY ARRAY");
   }
 
   @Test void testJsonQuery() {