org.apache.druid.sql.calcite.table.RowSignatures Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.apache.druid.sql.calcite.table;
import com.google.common.base.Preconditions;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeComparability;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.sql.SqlCallBinding;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlOperandCountRange;
import org.apache.calcite.sql.SqlOperator;
import org.apache.calcite.sql.type.AbstractSqlType;
import org.apache.calcite.sql.type.SqlOperandCountRanges;
import org.apache.calcite.sql.type.SqlSingleOperandTypeChecker;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.query.ordering.StringComparator;
import org.apache.druid.query.ordering.StringComparators;
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;
import javax.annotation.Nonnull;
import java.util.List;
/**
* Utility functions for working with {@link RowSignature}.
*/
public class RowSignatures
{
private RowSignatures()
{
// No instantiation.
}
public static RowSignature fromRelDataType(final List rowOrder, final RelDataType rowType)
{
if (rowOrder.size() != rowType.getFieldCount()) {
throw new IAE("Field count %d != %d", rowOrder.size(), rowType.getFieldCount());
}
final RowSignature.Builder rowSignatureBuilder = RowSignature.builder();
for (int i = 0; i < rowOrder.size(); i++) {
final RelDataType dataType = rowType.getFieldList().get(i).getType();
final ColumnType valueType = Calcites.getColumnTypeForRelDataType(dataType);
rowSignatureBuilder.add(rowOrder.get(i), valueType);
}
return rowSignatureBuilder.build();
}
/**
* Return the "natural" {@link StringComparator} for an extraction from a row signature. This will be a
* lexicographic comparator for String types and a numeric comparator for Number types.
*/
@Nonnull
public static StringComparator getNaturalStringComparator(
final RowSignature rowSignature,
final SimpleExtraction simpleExtraction
)
{
Preconditions.checkNotNull(simpleExtraction, "simpleExtraction");
if (simpleExtraction.getExtractionFn() != null
|| rowSignature.getColumnType(simpleExtraction.getColumn())
.map(type -> type.is(ValueType.STRING))
.orElse(false)) {
return StringComparators.LEXICOGRAPHIC;
} else {
return StringComparators.NUMERIC;
}
}
/**
* Returns a Calcite {@link RelDataType} corresponding to a {@link RowSignature}. It will typecast __time column to
* TIMESTAMP irrespective of the type present in the row signature
*/
public static RelDataType toRelDataType(final RowSignature rowSignature, final RelDataTypeFactory typeFactory)
{
return toRelDataType(rowSignature, typeFactory, true);
}
/**
* Returns a Calcite {@link RelDataType} corresponding to a {@link RowSignature}. For columns that are named
* "__time", it automatically casts it to TIMESTAMP if typecastTimeColumn is set to true
*/
public static RelDataType toRelDataType(
final RowSignature rowSignature,
final RelDataTypeFactory typeFactory,
boolean typecastTimeColumn
)
{
final RelDataTypeFactory.Builder builder = typeFactory.builder();
final boolean nullNumeric = !NullHandling.replaceWithDefault();
for (final String columnName : rowSignature.getColumnNames()) {
final RelDataType type;
if (typecastTimeColumn && ColumnHolder.TIME_COLUMN_NAME.equals(columnName)) {
type = Calcites.createSqlType(typeFactory, SqlTypeName.TIMESTAMP);
} else {
final ColumnType columnType =
rowSignature.getColumnType(columnName)
.orElseThrow(() -> new ISE("Encountered null type for column[%s]", columnName));
type = columnTypeToRelDataType(typeFactory, columnType, nullNumeric);
}
builder.add(columnName, type);
}
return builder.build();
}
/**
* Returns a Calcite {@link RelDataType} corresponding to a {@link ColumnType}
*/
public static RelDataType columnTypeToRelDataType(
RelDataTypeFactory typeFactory,
ColumnType columnType,
boolean nullNumeric
)
{
final RelDataType type;
switch (columnType.getType()) {
case STRING:
// Note that there is no attempt here to handle multi-value in any special way. Maybe one day...
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.VARCHAR, true);
break;
case LONG:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.BIGINT, nullNumeric);
break;
case FLOAT:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.FLOAT, nullNumeric);
break;
case DOUBLE:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.DOUBLE, nullNumeric);
break;
case ARRAY:
final RelDataType elementType = columnTypeToRelDataType(
typeFactory,
(ColumnType) columnType.getElementType(),
nullNumeric
);
type = typeFactory.createTypeWithNullability(
typeFactory.createArrayType(elementType, -1),
true
);
break;
case COMPLEX:
type = makeComplexType(typeFactory, columnType, true);
break;
default:
throw new ISE("valueType[%s] not translatable", columnType);
} return type;
}
/**
* Creates a {@link ComplexSqlType} using the supplied {@link RelDataTypeFactory} to ensure that the
* {@link ComplexSqlType} is interned. This is important because Calcite checks that the references are equal
* instead of the objects being equivalent.
*
* This method uses {@link RelDataTypeFactory#createTypeWithNullability(RelDataType, boolean) ensures that if the
* type factory is a {@link org.apache.calcite.rel.type.RelDataTypeFactoryImpl} that the type is passed through
* {@link org.apache.calcite.rel.type.RelDataTypeFactoryImpl#canonize(RelDataType)} which interns the type.
*/
public static RelDataType makeComplexType(RelDataTypeFactory typeFactory, ColumnType columnType, boolean isNullable)
{
return typeFactory.createTypeWithNullability(
new ComplexSqlType(SqlTypeName.OTHER, columnType, isNullable),
isNullable
);
}
/**
* Calcite {@link RelDataType} for Druid complex columns, to preserve complex type information.
*
* If using with other operations of a {@link RelDataTypeFactory}, consider wrapping the creation of this type in
* {@link RelDataTypeFactory#createTypeWithNullability(RelDataType, boolean) to ensure that if the type factory is a
* {@link org.apache.calcite.rel.type.RelDataTypeFactoryImpl} that the type is passed through
* {@link org.apache.calcite.rel.type.RelDataTypeFactoryImpl#canonize(RelDataType)} which interns the type.
*
* If {@link SqlTypeName} is going to be {@link SqlTypeName#OTHER} and a {@link RelDataTypeFactory} is available,
* consider using {@link #makeComplexType(RelDataTypeFactory, ColumnType, boolean)}.
*
* This type does not work well with {@link org.apache.calcite.sql.type.ReturnTypes#explicit(RelDataType)}, which
* will create new {@link RelDataType} using {@link SqlTypeName} during return type inference, so implementors of
* {@link org.apache.druid.sql.calcite.expression.SqlOperatorConversion} should implement the
* {@link org.apache.calcite.sql.type.SqlReturnTypeInference} directly for best results.
*/
public static final class ComplexSqlType extends AbstractSqlType
{
private final ColumnType columnType;
public ComplexSqlType(
SqlTypeName typeName,
ColumnType columnType,
boolean isNullable
)
{
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();
}
@Override
public RelDataTypeComparability getComparability()
{
return RelDataTypeComparability.UNORDERED;
}
@Override
protected void generateTypeString(StringBuilder sb, boolean withDetail)
{
sb.append(columnType.asTypeString());
}
public ColumnType getColumnType()
{
return columnType;
}
public String getComplexTypeName()
{
return columnType.getComplexTypeName();
}
public String asTypeString()
{
return columnType.asTypeString();
}
}
public static ComplexSqlSingleOperandTypeChecker complexTypeChecker(ColumnType complexType)
{
return new ComplexSqlSingleOperandTypeChecker(
new ComplexSqlType(SqlTypeName.OTHER, complexType, true)
);
}
public static final class ComplexSqlSingleOperandTypeChecker implements SqlSingleOperandTypeChecker
{
private final ComplexSqlType type;
public ComplexSqlSingleOperandTypeChecker(
ComplexSqlType type
)
{
this.type = type;
}
@Override
public boolean checkSingleOperandType(
SqlCallBinding callBinding,
SqlNode operand,
int iFormalOperand,
boolean throwOnFailure
)
{
return type.equals(callBinding.getValidator().deriveType(callBinding.getScope(), operand));
}
@Override
public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFailure)
{
if (callBinding.getOperandCount() != 1) {
return false;
}
return checkSingleOperandType(callBinding, callBinding.operand(0), 0, throwOnFailure);
}
@Override
public SqlOperandCountRange getOperandCountRange()
{
return SqlOperandCountRanges.of(1);
}
@Override
public String getAllowedSignatures(SqlOperator op, String opName)
{
return StringUtils.format("'%s'(%s)", opName, type);
}
@Override
public Consistency getConsistency()
{
return Consistency.NONE;
}
@Override
public boolean isOptional(int i)
{
return false;
}
}
}