fix complex types returning UNKNOWN as their SQL type inference (#16216)

* * fix

* * fix

* * address review comments

extensions-contrib/tdigestsketch/src/main/java/org/apache/druid/query/aggregation/tdigestsketch/sql/TDigestGenerateSketchSqlAggregator.java

@@ -25,9 +25,7 @@ import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.SqlKind;
-import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeFamily;
-import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.Optionality;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.query.aggregation.AggregatorFactory;
@@ -39,6 +37,7 @@ import org.apache.druid.sql.calcite.aggregation.Aggregations;
 import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
 import org.apache.druid.sql.calcite.expression.DefaultOperandTypeChecker;
 import org.apache.druid.sql.calcite.expression.DruidExpression;
+import org.apache.druid.sql.calcite.planner.Calcites;
 import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.rel.InputAccessor;
 import org.apache.druid.sql.calcite.rel.VirtualColumnRegistry;
@@ -140,7 +139,7 @@ public class TDigestGenerateSketchSqlAggregator implements SqlAggregator
           NAME,
           null,
           SqlKind.OTHER_FUNCTION,
-          ReturnTypes.explicit(SqlTypeName.OTHER),
+          Calcites.complexReturnTypeWithNullability(TDigestSketchAggregatorFactory.TYPE, false),
           null,
           // Validation for signatures like 'TDIGEST_GENERATE_SKETCH(column)' and
           // 'TDIGEST_GENERATE_SKETCH(column, compression)'

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchModule.java

@@ -35,6 +35,7 @@ import org.apache.druid.query.aggregation.datasketches.hll.sql.HllSketchEstimate
 import org.apache.druid.query.aggregation.datasketches.hll.sql.HllSketchObjectSqlAggregator;
 import org.apache.druid.query.aggregation.datasketches.hll.sql.HllSketchSetUnionOperatorConversion;
 import org.apache.druid.query.aggregation.datasketches.hll.sql.HllSketchToStringOperatorConversion;
+import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.serde.ComplexMetrics;
 import org.apache.druid.sql.guice.SqlBindings;
 
@@ -48,6 +49,7 @@ import java.util.List;
 public class HllSketchModule implements DruidModule
 {
   public static final String TYPE_NAME = "HLLSketch"; // common type name to be associated with segment data
+  public static final ColumnType TYPE = ColumnType.ofComplex(TYPE_NAME);
   public static final String BUILD_TYPE_NAME = "HLLSketchBuild";
   public static final String MERGE_TYPE_NAME = "HLLSketchMerge";
   public static final String TO_STRING_TYPE_NAME = "HLLSketchToString";

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchEstimateWithErrorBoundsOperatorConversion.java

@@ -47,7 +47,7 @@ public class HllSketchEstimateWithErrorBoundsOperatorConversion implements SqlOp
       .operatorBuilder(StringUtils.toUpperCase(FUNCTION_NAME))
       .operandTypes(SqlTypeFamily.ANY, SqlTypeFamily.INTEGER)
       .requiredOperandCount(1)
-      .returnTypeNonNull(SqlTypeName.OTHER)
+      .returnTypeNullableArrayWithNullableElements(SqlTypeName.DOUBLE)
       .build();
 
   @Override

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchObjectSqlAggregator.java

@@ -23,12 +23,13 @@ import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.type.InferTypes;
 import org.apache.calcite.sql.type.SqlTypeFamily;
-import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.datasketches.hll.HllSketchAggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.hll.HllSketchModule;
 import org.apache.druid.sql.calcite.aggregation.Aggregation;
 import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
 import org.apache.druid.sql.calcite.expression.OperatorConversions;
+import org.apache.druid.sql.calcite.planner.Calcites;
 
 import java.util.Collections;
 
@@ -42,7 +43,9 @@ public class HllSketchObjectSqlAggregator extends HllSketchBaseSqlAggregator imp
                          .operandTypeInference(InferTypes.VARCHAR_1024)
                          .requiredOperandCount(1)
                          .literalOperands(1, 2)
-                         .returnTypeNonNull(SqlTypeName.OTHER)
+                         .returnTypeInference(
+                             Calcites.complexReturnTypeWithNullability(HllSketchModule.TYPE, false)
+                         )
                          .functionCategory(SqlFunctionCategory.USER_DEFINED_FUNCTION)
                          .build();
 

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/sql/DoublesSketchObjectSqlAggregator.java

@@ -27,17 +27,18 @@ import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.type.SqlTypeFamily;
-import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.datasketches.SketchQueryContext;
 import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchAggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchModule;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.sql.calcite.aggregation.Aggregation;
 import org.apache.druid.sql.calcite.aggregation.Aggregations;
 import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
 import org.apache.druid.sql.calcite.expression.DruidExpression;
 import org.apache.druid.sql.calcite.expression.OperatorConversions;
+import org.apache.druid.sql.calcite.planner.Calcites;
 import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.rel.InputAccessor;
 import org.apache.druid.sql.calcite.rel.VirtualColumnRegistry;
@@ -52,7 +53,9 @@ public class DoublesSketchObjectSqlAggregator implements SqlAggregator
       OperatorConversions.aggregatorBuilder(NAME)
                          .operandNames("column", "k")
                          .operandTypes(SqlTypeFamily.ANY, SqlTypeFamily.EXACT_NUMERIC)
-                         .returnTypeNonNull(SqlTypeName.OTHER)
+                         .returnTypeInference(
+                             Calcites.complexReturnTypeWithNullability(DoublesSketchModule.TYPE, false)
+                         )
                          .requiredOperandCount(1)
                          .literalOperands(1)
                          .functionCategory(SqlFunctionCategory.NUMERIC)

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchModule.java

@@ -47,6 +47,7 @@ public class SketchModule implements DruidModule
 
   public static final String THETA_SKETCH_MERGE_AGG = "thetaSketchMerge";
   public static final String THETA_SKETCH_BUILD_AGG = "thetaSketchBuild";
+  public static final ColumnType THETA_SKETCH_TYPE = ColumnType.ofComplex(THETA_SKETCH);
   public static final ColumnType BUILD_TYPE = ColumnType.ofComplex(THETA_SKETCH_BUILD_AGG);
   public static final ColumnType MERGE_TYPE = ColumnType.ofComplex(THETA_SKETCH_MERGE_AGG);
 

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/sql/ThetaSketchEstimateWithErrorBoundsOperatorConversion.java

@@ -26,15 +26,16 @@ import org.apache.calcite.sql.SqlFunction;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.type.SqlTypeFamily;
-import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.query.aggregation.PostAggregator;
 import org.apache.druid.query.aggregation.datasketches.theta.SketchEstimatePostAggregator;
+import org.apache.druid.query.aggregation.datasketches.theta.SketchModule;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.sql.calcite.expression.DruidExpression;
 import org.apache.druid.sql.calcite.expression.OperatorConversions;
 import org.apache.druid.sql.calcite.expression.PostAggregatorVisitor;
 import org.apache.druid.sql.calcite.expression.SqlOperatorConversion;
+import org.apache.druid.sql.calcite.planner.Calcites;
 import org.apache.druid.sql.calcite.planner.PlannerContext;
 
 import javax.annotation.Nullable;
@@ -46,7 +47,9 @@ public class ThetaSketchEstimateWithErrorBoundsOperatorConversion implements Sql
   private static final SqlFunction SQL_FUNCTION = OperatorConversions
       .operatorBuilder(StringUtils.toUpperCase(FUNCTION_NAME))
       .operandTypes(SqlTypeFamily.ANY, SqlTypeFamily.INTEGER)
-      .returnTypeNonNull(SqlTypeName.OTHER)
+      .returnTypeInference(
+          Calcites.complexReturnTypeWithNullability(SketchModule.THETA_SKETCH_TYPE, false)
+      )
       .build();
 
   @Override

extensions-core/druid-bloom-filter/src/main/java/org/apache/druid/query/aggregation/bloom/sql/BloomFilterSqlAggregator.java

@@ -25,9 +25,7 @@ import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.SqlKind;
-import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeFamily;
-import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.Optionality;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.query.aggregation.AggregatorFactory;
@@ -175,7 +173,7 @@ public class BloomFilterSqlAggregator implements SqlAggregator
           NAME,
           null,
           SqlKind.OTHER_FUNCTION,
-          ReturnTypes.explicit(SqlTypeName.OTHER),
+          Calcites.complexReturnTypeWithNullability(BloomFilterAggregatorFactory.TYPE, false),
           null,
           // Allow signatures like 'BLOOM_FILTER(column, maxNumEntries)'
           DefaultOperandTypeChecker

sql/src/main/java/org/apache/druid/sql/calcite/planner/Calcites.java

@@ -100,6 +100,15 @@ public class Calcites
   public static final SqlReturnTypeInference
       ARG1_NULLABLE_ARRAY_RETURN_TYPE_INFERENCE = new Arg1NullableArrayTypeInference();
 
+  public static SqlReturnTypeInference complexReturnTypeWithNullability(ColumnType columnType, boolean nullable)
+  {
+    return opBinding -> RowSignatures.makeComplexType(
+        opBinding.getTypeFactory(),
+        columnType,
+        nullable
+    );
+  }
+
   private Calcites()
   {
     // No instantiation.

sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java

@@ -41,6 +41,8 @@ import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.column.ValueType;
+import org.apache.druid.segment.serde.ComplexMetricSerde;
+import org.apache.druid.segment.serde.ComplexMetrics;
 import org.apache.druid.sql.calcite.expression.SimpleExtraction;
 import org.apache.druid.sql.calcite.planner.Calcites;
 
@@ -222,7 +224,17 @@ public class RowSignatures
     )
     {
       super(typeName, isNullable, null);
+      // homogenize complex type names to common name
+      final ComplexMetricSerde serde = columnType.getComplexTypeName() != null
+          ?
+          ComplexMetrics.getSerdeForType(columnType.getComplexTypeName())
+          : null;
+
+      if (serde != null) {
+        this.columnType = ColumnType.ofComplex(serde.getTypeName());
+      } else {
         this.columnType = columnType;
+      }
       this.computeDigest();
     }