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

* * fix

* * fix

* * address review comments
This commit is contained in:
zachjsh 2024-04-02 14:36:01 -04:00 committed by GitHub
parent 26f9b174de
commit 9b52c909e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 44 additions and 14 deletions

View File

@ -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)'

View File

@ -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";

View File

@ -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

View File

@ -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();

View File

@ -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)

View File

@ -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);

View File

@ -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

View File

@ -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

View File

@ -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.

View File

@ -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();
}