Search in sources :

Example 1 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression in project druid by druid-io.

the class DruidQuery method computeDimensions.

/**
 * Returns dimensions corresponding to {@code aggregate.getGroupSet()}, in the same order.
 *
 * @param partialQuery          partial query
 * @param plannerContext        planner context
 * @param rowSignature          source row signature
 * @param virtualColumnRegistry re-usable virtual column references
 *
 * @return dimensions
 *
 * @throws CannotBuildQueryException if dimensions cannot be computed
 */
private static List<DimensionExpression> computeDimensions(final PartialDruidQuery partialQuery, final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry) {
    final Aggregate aggregate = Preconditions.checkNotNull(partialQuery.getAggregate());
    final List<DimensionExpression> dimensions = new ArrayList<>();
    final String outputNamePrefix = Calcites.findUnusedPrefixForDigits("d", rowSignature.getColumnNames());
    int outputNameCounter = 0;
    for (int i : aggregate.getGroupSet()) {
        // Dimension might need to create virtual columns. Avoid giving it a name that would lead to colliding columns.
        final RexNode rexNode = Expressions.fromFieldAccess(rowSignature, partialQuery.getSelectProject(), i);
        final DruidExpression druidExpression = Expressions.toDruidExpression(plannerContext, rowSignature, rexNode);
        if (druidExpression == null) {
            throw new CannotBuildQueryException(aggregate, rexNode);
        }
        final RelDataType dataType = rexNode.getType();
        final ColumnType outputType = Calcites.getColumnTypeForRelDataType(dataType);
        if (Types.isNullOr(outputType, ValueType.COMPLEX)) {
            // Can't group on unknown or COMPLEX types.
            plannerContext.setPlanningError("SQL requires a group-by on a column of type %s that is unsupported.", outputType);
            throw new CannotBuildQueryException(aggregate, rexNode);
        }
        final String dimOutputName = outputNamePrefix + outputNameCounter++;
        if (!druidExpression.isSimpleExtraction()) {
            final String virtualColumn = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(druidExpression, dataType);
            dimensions.add(DimensionExpression.ofVirtualColumn(virtualColumn, dimOutputName, druidExpression, outputType));
        } else {
            dimensions.add(DimensionExpression.ofSimpleColumn(dimOutputName, druidExpression, outputType));
        }
    }
    return dimensions;
}
Also used : ColumnType(org.apache.druid.segment.column.ColumnType) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) IntArrayList(it.unimi.dsi.fastutil.ints.IntArrayList) ArrayList(java.util.ArrayList) DimensionExpression(org.apache.druid.sql.calcite.aggregation.DimensionExpression) RelDataType(org.apache.calcite.rel.type.RelDataType) Aggregate(org.apache.calcite.rel.core.Aggregate) RexNode(org.apache.calcite.rex.RexNode)

Example 2 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression in project druid by druid-io.

the class Projection method postAggregationHandleInputRefOrLiteral.

private static void postAggregationHandleInputRefOrLiteral(final Project project, final PlannerContext plannerContext, final RowSignature inputRowSignature, final RexNode postAggregatorRexNode, final List<String> rowOrder, final PostAggregatorVisitor postAggregatorVisitor) {
    // Attempt to convert to PostAggregator.
    final DruidExpression postAggregatorExpression = Expressions.toDruidExpression(plannerContext, inputRowSignature, postAggregatorRexNode);
    if (postAggregatorExpression == null) {
        throw new CannotBuildQueryException(project, postAggregatorRexNode);
    }
    handlePostAggregatorExpression(plannerContext, inputRowSignature, postAggregatorRexNode, rowOrder, postAggregatorVisitor, postAggregatorExpression);
}
Also used : DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression)

Example 3 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression in project druid by druid-io.

the class BuiltinApproxCountDistinctSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final String name, final AggregateCall aggregateCall, final Project project, final List<Aggregation> existingAggregations, final boolean finalizeAggregations) {
    // Don't use Aggregations.getArgumentsForSimpleAggregator, since it won't let us use direct column access
    // for string columns.
    final RexNode rexNode = Expressions.fromFieldAccess(rowSignature, project, Iterables.getOnlyElement(aggregateCall.getArgList()));
    final DruidExpression arg = Expressions.toDruidExpression(plannerContext, rowSignature, rexNode);
    if (arg == null) {
        return null;
    }
    final AggregatorFactory aggregatorFactory;
    final String aggregatorName = finalizeAggregations ? Calcites.makePrefixedName(name, "a") : name;
    if (arg.isDirectColumnAccess() && rowSignature.getColumnType(arg.getDirectColumn()).map(type -> type.is(ValueType.COMPLEX)).orElse(false)) {
        aggregatorFactory = new HyperUniquesAggregatorFactory(aggregatorName, arg.getDirectColumn(), false, true);
    } else {
        final RelDataType dataType = rexNode.getType();
        final ColumnType inputType = Calcites.getColumnTypeForRelDataType(dataType);
        if (inputType == null) {
            throw new ISE("Cannot translate sqlTypeName[%s] to Druid type for field[%s]", dataType.getSqlTypeName(), aggregatorName);
        }
        final DimensionSpec dimensionSpec;
        if (arg.isSimpleExtraction()) {
            dimensionSpec = arg.getSimpleExtraction().toDimensionSpec(null, inputType);
        } else {
            String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(arg, dataType);
            dimensionSpec = new DefaultDimensionSpec(virtualColumnName, null, inputType);
        }
        aggregatorFactory = new CardinalityAggregatorFactory(aggregatorName, null, ImmutableList.of(dimensionSpec), false, true);
    }
    return Aggregation.create(Collections.singletonList(aggregatorFactory), finalizeAggregations ? new HyperUniqueFinalizingPostAggregator(name, aggregatorFactory.getName()) : null);
}
Also used : DefaultDimensionSpec(org.apache.druid.query.dimension.DefaultDimensionSpec) DimensionSpec(org.apache.druid.query.dimension.DimensionSpec) ColumnType(org.apache.druid.segment.column.ColumnType) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HyperUniquesAggregatorFactory(org.apache.druid.query.aggregation.hyperloglog.HyperUniquesAggregatorFactory) RelDataType(org.apache.calcite.rel.type.RelDataType) ISE(org.apache.druid.java.util.common.ISE) HyperUniqueFinalizingPostAggregator(org.apache.druid.query.aggregation.hyperloglog.HyperUniqueFinalizingPostAggregator) CardinalityAggregatorFactory(org.apache.druid.query.aggregation.cardinality.CardinalityAggregatorFactory) HyperUniquesAggregatorFactory(org.apache.druid.query.aggregation.hyperloglog.HyperUniquesAggregatorFactory) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) CardinalityAggregatorFactory(org.apache.druid.query.aggregation.cardinality.CardinalityAggregatorFactory) DefaultDimensionSpec(org.apache.druid.query.dimension.DefaultDimensionSpec) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 4 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression in project druid by druid-io.

the class GroupingSqlAggregator method getColumnName.

@Nullable
private String getColumnName(PlannerContext plannerContext, RowSignature rowSignature, Project project, VirtualColumnRegistry virtualColumnRegistry, int fieldNumber) {
    RexNode node = Expressions.fromFieldAccess(rowSignature, project, fieldNumber);
    if (null == node) {
        return null;
    }
    DruidExpression expression = Expressions.toDruidExpression(plannerContext, rowSignature, node);
    if (null == expression) {
        return null;
    }
    if (expression.isDirectColumnAccess()) {
        return expression.getDirectColumn();
    }
    String virtualColumn = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(expression, node.getType());
    return virtualColumn;
}
Also used : DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 5 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression in project druid by druid-io.

the class SimpleSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final String name, final AggregateCall aggregateCall, final Project project, final List<Aggregation> existingAggregations, final boolean finalizeAggregations) {
    if (aggregateCall.isDistinct()) {
        return null;
    }
    final List<DruidExpression> arguments = Aggregations.getArgumentsForSimpleAggregator(plannerContext, rowSignature, aggregateCall, project);
    if (arguments == null) {
        return null;
    }
    final DruidExpression arg = Iterables.getOnlyElement(arguments);
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    final String expression;
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
        expression = null;
    } else {
        fieldName = null;
        expression = arg.getExpression();
    }
    return getAggregation(name, aggregateCall, macroTable, fieldName, expression);
}
Also used : DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) Nullable(javax.annotation.Nullable)

Aggregations

DruidExpression (org.apache.druid.sql.calcite.expression.DruidExpression)53 RexNode (org.apache.calcite.rex.RexNode)34 Nullable (javax.annotation.Nullable)30 RexCall (org.apache.calcite.rex.RexCall)20 ColumnType (org.apache.druid.segment.column.ColumnType)15 AggregatorFactory (org.apache.druid.query.aggregation.AggregatorFactory)14 RexLiteral (org.apache.calcite.rex.RexLiteral)12 Aggregation (org.apache.druid.sql.calcite.aggregation.Aggregation)12 List (java.util.List)11 RowSignature (org.apache.druid.segment.column.RowSignature)11 RelDataType (org.apache.calcite.rel.type.RelDataType)10 Expressions (org.apache.druid.sql.calcite.expression.Expressions)10 SqlKind (org.apache.calcite.sql.SqlKind)9 PlannerContext (org.apache.druid.sql.calcite.planner.PlannerContext)9 RexBuilder (org.apache.calcite.rex.RexBuilder)8 SqlAggFunction (org.apache.calcite.sql.SqlAggFunction)8 SqlFunctionCategory (org.apache.calcite.sql.SqlFunctionCategory)8 ExprMacroTable (org.apache.druid.math.expr.ExprMacroTable)8 Collectors (java.util.stream.Collectors)7 AggregateCall (org.apache.calcite.rel.core.AggregateCall)7