From 9b52c909e08849acfdce8c59b54f350aad940bed Mon Sep 17 00:00:00 2001 From: zachjsh Date: Tue, 2 Apr 2024 14:36:01 -0400 Subject: [PATCH] fix complex types returning UNKNOWN as their SQL type inference (#16216) * * fix * * fix * * address review comments --- .../sql/TDigestGenerateSketchSqlAggregator.java | 5 ++--- .../datasketches/hll/HllSketchModule.java | 2 ++ ...hEstimateWithErrorBoundsOperatorConversion.java | 2 +- .../hll/sql/HllSketchObjectSqlAggregator.java | 7 +++++-- .../sql/DoublesSketchObjectSqlAggregator.java | 7 +++++-- .../datasketches/theta/SketchModule.java | 1 + ...hEstimateWithErrorBoundsOperatorConversion.java | 7 +++++-- .../bloom/sql/BloomFilterSqlAggregator.java | 4 +--- .../apache/druid/sql/calcite/planner/Calcites.java | 9 +++++++++ .../druid/sql/calcite/table/RowSignatures.java | 14 +++++++++++++- 10 files changed, 44 insertions(+), 14 deletions(-) diff --git a/extensions-contrib/tdigestsketch/src/main/java/org/apache/druid/query/aggregation/tdigestsketch/sql/TDigestGenerateSketchSqlAggregator.java b/extensions-contrib/tdigestsketch/src/main/java/org/apache/druid/query/aggregation/tdigestsketch/sql/TDigestGenerateSketchSqlAggregator.java index cb9b95423a6..1777ce5c544 100644 --- a/extensions-contrib/tdigestsketch/src/main/java/org/apache/druid/query/aggregation/tdigestsketch/sql/TDigestGenerateSketchSqlAggregator.java +++ b/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)' diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchModule.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchModule.java index ea2f11ca785..a44f0cee663 100644 --- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchModule.java +++ b/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"; diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchEstimateWithErrorBoundsOperatorConversion.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchEstimateWithErrorBoundsOperatorConversion.java index cc89d4c7a56..0c2e8e0d668 100644 --- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchEstimateWithErrorBoundsOperatorConversion.java +++ b/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 diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchObjectSqlAggregator.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchObjectSqlAggregator.java index 9d8ade636f1..0e466bcaf0a 100644 --- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchObjectSqlAggregator.java +++ b/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(); diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/sql/DoublesSketchObjectSqlAggregator.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/sql/DoublesSketchObjectSqlAggregator.java index 8331ab72064..15b15b0dc21 100644 --- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/sql/DoublesSketchObjectSqlAggregator.java +++ b/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) diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchModule.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchModule.java index 979f3f2579f..5a67aa3d2f1 100644 --- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchModule.java +++ b/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); diff --git a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/sql/ThetaSketchEstimateWithErrorBoundsOperatorConversion.java b/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/sql/ThetaSketchEstimateWithErrorBoundsOperatorConversion.java index 459a12867bc..6b07f98d668 100644 --- a/extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/sql/ThetaSketchEstimateWithErrorBoundsOperatorConversion.java +++ b/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 diff --git a/extensions-core/druid-bloom-filter/src/main/java/org/apache/druid/query/aggregation/bloom/sql/BloomFilterSqlAggregator.java b/extensions-core/druid-bloom-filter/src/main/java/org/apache/druid/query/aggregation/bloom/sql/BloomFilterSqlAggregator.java index cb73d94ef7e..5beb6c64232 100644 --- a/extensions-core/druid-bloom-filter/src/main/java/org/apache/druid/query/aggregation/bloom/sql/BloomFilterSqlAggregator.java +++ b/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 diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/Calcites.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/Calcites.java index b56a6d5bed5..64929c6445e 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/Calcites.java +++ b/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. diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java b/sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java index 0234d6b5319..e02cd645894 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java +++ b/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); - this.columnType = columnType; + // 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(); }