Add all supported types to return type inference from string literal …

…for JSON extract functions in the multi-stage query engine (#14289)
apache · Oct 31, 2024 · e61d503 · e61d503
1 parent 65f658d
commit e61d503
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java
@@ -315,12 +315,26 @@ private static RelDataType inferTypeFromStringLiteral(String operandTypeStr, Rel
     switch (operandTypeStr) {
       case "INT":
         return typeFactory.createSqlType(SqlTypeName.INTEGER);
+      case "INT_ARRAY":
+        return typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.INTEGER), -1);
       case "LONG":
         return typeFactory.createSqlType(SqlTypeName.BIGINT);
+      case "LONG_ARRAY":
+        return typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.BIGINT), -1);
+      case "FLOAT":
+        return typeFactory.createSqlType(SqlTypeName.REAL);
+      case "FLOAT_ARRAY":
+        return typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.REAL), -1);
+      case "DOUBLE_ARRAY":
+        return typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.DOUBLE), -1);
       case "STRING":
         return typeFactory.createSqlType(SqlTypeName.VARCHAR);
+      case "STRING_ARRAY":
+        return typeFactory.createArrayType(typeFactory.createSqlType(SqlTypeName.VARCHAR), -1);
       case "BYTES":
         return typeFactory.createSqlType(SqlTypeName.VARBINARY);
+      case "BIG_DECIMAL":
+        return typeFactory.createSqlType(SqlTypeName.DECIMAL);
       default:
         SqlTypeName sqlTypeName = SqlTypeName.get(operandTypeStr);
         if (sqlTypeName == null) {

diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
@@ -236,7 +236,20 @@ protected Object[][] provideQueries() {
             "SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */ a.col2, a.col3 FROM a JOIN b "
                 + "ON a.col1 = b.col1  WHERE a.col3 >= 0 GROUP BY a.col2, a.col3"
         },
-        new Object[]{"SELECT ROUND(ts_timestamp, 10000) FROM a"}
+        new Object[]{"SELECT ROUND(ts_timestamp, 10000) FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'INT') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'LONG') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'FLOAT') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'DOUBLE') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'BOOLEAN') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'BIG_DECIMAL') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'TIMESTAMP') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'STRING') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'INT_ARRAY') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'LONG_ARRAY') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'FLOAT_ARRAY') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'DOUBLE_ARRAY') FROM a"},
+        new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'STRING_ARRAY') FROM a"},
     };
   }