Search in sources :

Example 16 with Aggregation

use of org.apache.druid.sql.calcite.aggregation.Aggregation in project druid by druid-io.

the class ArraySqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(PlannerContext plannerContext, RowSignature rowSignature, VirtualColumnRegistry virtualColumnRegistry, RexBuilder rexBuilder, String name, AggregateCall aggregateCall, Project project, List<Aggregation> existingAggregations, boolean finalizeAggregations) {
    final List<RexNode> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).collect(Collectors.toList());
    Integer maxSizeBytes = null;
    if (arguments.size() > 1) {
        RexNode maxBytes = arguments.get(1);
        if (!maxBytes.isA(SqlKind.LITERAL)) {
            // maxBytes must be a literal
            return null;
        }
        maxSizeBytes = ((Number) RexLiteral.value(maxBytes)).intValue();
    }
    final DruidExpression arg = Expressions.toDruidExpression(plannerContext, rowSignature, arguments.get(0));
    if (arg == null) {
        // can't translate argument
        return null;
    }
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    final String initialvalue;
    final ColumnType druidType = Calcites.getValueTypeForRelDataTypeFull(aggregateCall.getType());
    final ColumnType elementType;
    if (druidType == null || !druidType.isArray()) {
        initialvalue = "[]";
        elementType = ColumnType.STRING;
    } else {
        initialvalue = ExpressionType.fromColumnTypeStrict(druidType).asTypeString() + "[]";
        elementType = (ColumnType) druidType.getElementType();
    }
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        fieldName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(arg, elementType);
    }
    if (aggregateCall.isDistinct()) {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, true, false, StringUtils.format("array_set_add(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    } else {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, true, false, StringUtils.format("array_append(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_concat(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    }
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) RowSignatures(org.apache.druid.sql.calcite.table.RowSignatures) UnsupportedSQLQueryException(org.apache.druid.sql.calcite.planner.UnsupportedSQLQueryException) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) Optionality(org.apache.calcite.util.Optionality) RexNode(org.apache.calcite.rex.RexNode) ExpressionType(org.apache.druid.math.expr.ExpressionType) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Nullable(javax.annotation.Nullable) SqlOperatorBinding(org.apache.calcite.sql.SqlOperatorBinding) RelDataType(org.apache.calcite.rel.type.RelDataType) ImmutableSet(com.google.common.collect.ImmutableSet) SqlKind(org.apache.calcite.sql.SqlKind) SqlTypeFamily(org.apache.calcite.sql.type.SqlTypeFamily) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) InferTypes(org.apache.calcite.sql.type.InferTypes) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) StringUtils(org.apache.druid.java.util.common.StringUtils) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Collectors(java.util.stream.Collectors) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) SqlReturnTypeInference(org.apache.calcite.sql.type.SqlReturnTypeInference) List(java.util.List) RowSignature(org.apache.druid.segment.column.RowSignature) OperandTypes(org.apache.calcite.sql.type.OperandTypes) ColumnType(org.apache.druid.segment.column.ColumnType) AggregateCall(org.apache.calcite.rel.core.AggregateCall) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) Calcites(org.apache.druid.sql.calcite.planner.Calcites) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ColumnType(org.apache.druid.segment.column.ColumnType) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 17 with Aggregation

use of org.apache.druid.sql.calcite.aggregation.Aggregation in project druid by druid-io.

the class EarliestLatestAnySqlAggregator 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) {
    final List<RexNode> rexNodes = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).collect(Collectors.toList());
    final List<DruidExpression> args = Expressions.toDruidExpressions(plannerContext, rowSignature, rexNodes);
    if (args == null) {
        return null;
    }
    final String aggregatorName = finalizeAggregations ? Calcites.makePrefixedName(name, "a") : name;
    final ColumnType outputType = Calcites.getColumnTypeForRelDataType(aggregateCall.getType());
    if (outputType == null) {
        throw new ISE("Cannot translate output sqlTypeName[%s] to Druid type for aggregator[%s]", aggregateCall.getType().getSqlTypeName(), aggregateCall.getName());
    }
    final String fieldName = getColumnName(plannerContext, virtualColumnRegistry, args.get(0), rexNodes.get(0));
    final AggregatorFactory theAggFactory;
    switch(args.size()) {
        case 1:
            theAggFactory = aggregatorType.createAggregatorFactory(aggregatorName, fieldName, null, outputType, -1);
            break;
        case 2:
            theAggFactory = aggregatorType.createAggregatorFactory(aggregatorName, fieldName, null, outputType, RexLiteral.intValue(rexNodes.get(1)));
            break;
        default:
            throw new IAE("aggregation[%s], Invalid number of arguments[%,d] to [%s] operator", aggregatorName, args.size(), aggregatorType.name());
    }
    return Aggregation.create(Collections.singletonList(theAggFactory), finalizeAggregations ? new FinalizingFieldAccessPostAggregator(name, aggregatorName) : null);
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) StringLastAggregatorFactory(org.apache.druid.query.aggregation.last.StringLastAggregatorFactory) UnsupportedSQLQueryException(org.apache.druid.sql.calcite.planner.UnsupportedSQLQueryException) FinalizingFieldAccessPostAggregator(org.apache.druid.query.aggregation.post.FinalizingFieldAccessPostAggregator) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) Optionality(org.apache.calcite.util.Optionality) StringAnyAggregatorFactory(org.apache.druid.query.aggregation.any.StringAnyAggregatorFactory) RexNode(org.apache.calcite.rex.RexNode) FloatAnyAggregatorFactory(org.apache.druid.query.aggregation.any.FloatAnyAggregatorFactory) LongAnyAggregatorFactory(org.apache.druid.query.aggregation.any.LongAnyAggregatorFactory) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) LongLastAggregatorFactory(org.apache.druid.query.aggregation.last.LongLastAggregatorFactory) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) LongFirstAggregatorFactory(org.apache.druid.query.aggregation.first.LongFirstAggregatorFactory) DoubleAnyAggregatorFactory(org.apache.druid.query.aggregation.any.DoubleAnyAggregatorFactory) IAE(org.apache.druid.java.util.common.IAE) FloatFirstAggregatorFactory(org.apache.druid.query.aggregation.first.FloatFirstAggregatorFactory) Nullable(javax.annotation.Nullable) SqlOperatorBinding(org.apache.calcite.sql.SqlOperatorBinding) RelDataType(org.apache.calcite.rel.type.RelDataType) SqlKind(org.apache.calcite.sql.SqlKind) SqlTypeName(org.apache.calcite.sql.type.SqlTypeName) InferTypes(org.apache.calcite.sql.type.InferTypes) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) ISE(org.apache.druid.java.util.common.ISE) SqlTypeUtil(org.apache.calcite.sql.type.SqlTypeUtil) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) FloatLastAggregatorFactory(org.apache.druid.query.aggregation.last.FloatLastAggregatorFactory) Collectors(java.util.stream.Collectors) SqlReturnTypeInference(org.apache.calcite.sql.type.SqlReturnTypeInference) List(java.util.List) RowSignature(org.apache.druid.segment.column.RowSignature) OperandTypes(org.apache.calcite.sql.type.OperandTypes) ColumnType(org.apache.druid.segment.column.ColumnType) AggregateCall(org.apache.calcite.rel.core.AggregateCall) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) DoubleFirstAggregatorFactory(org.apache.druid.query.aggregation.first.DoubleFirstAggregatorFactory) Calcites(org.apache.druid.sql.calcite.planner.Calcites) Collections(java.util.Collections) StringFirstAggregatorFactory(org.apache.druid.query.aggregation.first.StringFirstAggregatorFactory) DoubleLastAggregatorFactory(org.apache.druid.query.aggregation.last.DoubleLastAggregatorFactory) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ColumnType(org.apache.druid.segment.column.ColumnType) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ISE(org.apache.druid.java.util.common.ISE) FinalizingFieldAccessPostAggregator(org.apache.druid.query.aggregation.post.FinalizingFieldAccessPostAggregator) StringLastAggregatorFactory(org.apache.druid.query.aggregation.last.StringLastAggregatorFactory) StringAnyAggregatorFactory(org.apache.druid.query.aggregation.any.StringAnyAggregatorFactory) FloatAnyAggregatorFactory(org.apache.druid.query.aggregation.any.FloatAnyAggregatorFactory) LongAnyAggregatorFactory(org.apache.druid.query.aggregation.any.LongAnyAggregatorFactory) LongLastAggregatorFactory(org.apache.druid.query.aggregation.last.LongLastAggregatorFactory) LongFirstAggregatorFactory(org.apache.druid.query.aggregation.first.LongFirstAggregatorFactory) DoubleAnyAggregatorFactory(org.apache.druid.query.aggregation.any.DoubleAnyAggregatorFactory) FloatFirstAggregatorFactory(org.apache.druid.query.aggregation.first.FloatFirstAggregatorFactory) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) FloatLastAggregatorFactory(org.apache.druid.query.aggregation.last.FloatLastAggregatorFactory) DoubleFirstAggregatorFactory(org.apache.druid.query.aggregation.first.DoubleFirstAggregatorFactory) StringFirstAggregatorFactory(org.apache.druid.query.aggregation.first.StringFirstAggregatorFactory) DoubleLastAggregatorFactory(org.apache.druid.query.aggregation.last.DoubleLastAggregatorFactory) IAE(org.apache.druid.java.util.common.IAE) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 18 with Aggregation

use of org.apache.druid.sql.calcite.aggregation.Aggregation in project druid by druid-io.

the class GroupByRules method translateAggregateCall.

/**
 * Translate an AggregateCall to Druid equivalents.
 *
 * @return translated aggregation, or null if translation failed.
 */
public static Aggregation translateAggregateCall(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final Project project, final List<Aggregation> existingAggregations, final String name, final AggregateCall call, final boolean finalizeAggregations) {
    final DimFilter filter;
    if (call.filterArg >= 0) {
        // AGG(xxx) FILTER(WHERE yyy)
        if (project == null) {
            // We need some kind of projection to support filtered aggregations.
            return null;
        }
        final RexNode expression = project.getChildExps().get(call.filterArg);
        final DimFilter nonOptimizedFilter = Expressions.toFilter(plannerContext, rowSignature, virtualColumnRegistry, expression);
        if (nonOptimizedFilter == null) {
            return null;
        } else {
            filter = Filtration.create(nonOptimizedFilter).optimizeFilterOnly(virtualColumnRegistry.getFullRowSignature()).getDimFilter();
        }
    } else {
        filter = null;
    }
    final SqlAggregator sqlAggregator = plannerContext.getOperatorTable().lookupAggregator(call.getAggregation());
    if (sqlAggregator == null) {
        return null;
    }
    // Compute existingAggregations for SqlAggregator impls that want it.
    final List<Aggregation> existingAggregationsWithSameFilter = new ArrayList<>();
    for (Aggregation existingAggregation : existingAggregations) {
        if (filter == null) {
            final boolean doesMatch = existingAggregation.getAggregatorFactories().stream().noneMatch(factory -> factory instanceof FilteredAggregatorFactory);
            if (doesMatch) {
                existingAggregationsWithSameFilter.add(existingAggregation);
            }
        } else {
            final boolean doesMatch = existingAggregation.getAggregatorFactories().stream().allMatch(factory -> factory instanceof FilteredAggregatorFactory && ((FilteredAggregatorFactory) factory).getFilter().equals(filter));
            if (doesMatch) {
                existingAggregationsWithSameFilter.add(Aggregation.create(existingAggregation.getAggregatorFactories().stream().map(factory -> ((FilteredAggregatorFactory) factory).getAggregator()).collect(Collectors.toList()), existingAggregation.getPostAggregator()));
            }
        }
    }
    final Aggregation retVal = sqlAggregator.toDruidAggregation(plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, name, call, project, existingAggregationsWithSameFilter, finalizeAggregations);
    if (retVal == null) {
        return null;
    } else {
        // Check if this refers to the existingAggregationsWithSameFilter. If so, no need to apply the filter.
        if (isUsingExistingAggregation(retVal, existingAggregationsWithSameFilter)) {
            return retVal;
        } else {
            return retVal.filter(rowSignature, virtualColumnRegistry, filter);
        }
    }
}
Also used : Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) FilteredAggregatorFactory(org.apache.druid.query.aggregation.FilteredAggregatorFactory) Project(org.apache.calcite.rel.core.Project) RexBuilder(org.apache.calcite.rex.RexBuilder) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) FilteredAggregatorFactory(org.apache.druid.query.aggregation.FilteredAggregatorFactory) Set(java.util.Set) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Collectors(java.util.stream.Collectors) ArrayList(java.util.ArrayList) List(java.util.List) DimFilter(org.apache.druid.query.filter.DimFilter) RexNode(org.apache.calcite.rex.RexNode) RowSignature(org.apache.druid.segment.column.RowSignature) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) AggregateCall(org.apache.calcite.rel.core.AggregateCall) Expressions(org.apache.druid.sql.calcite.expression.Expressions) Filtration(org.apache.druid.sql.calcite.filtration.Filtration) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) ArrayList(java.util.ArrayList) DimFilter(org.apache.druid.query.filter.DimFilter) RexNode(org.apache.calcite.rex.RexNode)

Aggregations

Aggregation (org.apache.druid.sql.calcite.aggregation.Aggregation)18 Nullable (javax.annotation.Nullable)13 RexNode (org.apache.calcite.rex.RexNode)12 DruidExpression (org.apache.druid.sql.calcite.expression.DruidExpression)12 AggregatorFactory (org.apache.druid.query.aggregation.AggregatorFactory)10 AggregateCall (org.apache.calcite.rel.core.AggregateCall)9 Project (org.apache.calcite.rel.core.Project)9 RowSignature (org.apache.druid.segment.column.RowSignature)9 List (java.util.List)8 RexBuilder (org.apache.calcite.rex.RexBuilder)8 ColumnType (org.apache.druid.segment.column.ColumnType)8 SqlAggregator (org.apache.druid.sql.calcite.aggregation.SqlAggregator)8 Expressions (org.apache.druid.sql.calcite.expression.Expressions)8 PlannerContext (org.apache.druid.sql.calcite.planner.PlannerContext)8 VirtualColumnRegistry (org.apache.druid.sql.calcite.rel.VirtualColumnRegistry)8 Collectors (java.util.stream.Collectors)7 SqlAggFunction (org.apache.calcite.sql.SqlAggFunction)7 SqlFunctionCategory (org.apache.calcite.sql.SqlFunctionCategory)7 SqlKind (org.apache.calcite.sql.SqlKind)7 OperandTypes (org.apache.calcite.sql.type.OperandTypes)7