Search in sources :

Example 1 with Aggregation

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

the class DruidQuery method computeGrouping.

@Nonnull
private static Grouping computeGrouping(final PartialDruidQuery partialQuery, final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final boolean finalizeAggregations) {
    final Aggregate aggregate = Preconditions.checkNotNull(partialQuery.getAggregate(), "aggregate");
    final Project aggregateProject = partialQuery.getAggregateProject();
    final List<DimensionExpression> dimensions = computeDimensions(partialQuery, plannerContext, rowSignature, virtualColumnRegistry);
    final Subtotals subtotals = computeSubtotals(partialQuery, rowSignature);
    final List<Aggregation> aggregations = computeAggregations(partialQuery, plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, finalizeAggregations);
    final RowSignature aggregateRowSignature = RowSignatures.fromRelDataType(ImmutableList.copyOf(Iterators.concat(dimensions.stream().map(DimensionExpression::getOutputName).iterator(), aggregations.stream().map(Aggregation::getOutputName).iterator())), aggregate.getRowType());
    final DimFilter havingFilter = computeHavingFilter(partialQuery, plannerContext, aggregateRowSignature);
    final Grouping grouping = Grouping.create(dimensions, subtotals, aggregations, havingFilter, aggregateRowSignature);
    if (aggregateProject == null) {
        return grouping;
    } else {
        return grouping.applyProject(plannerContext, aggregateProject);
    }
}
Also used : Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Project(org.apache.calcite.rel.core.Project) DimensionExpression(org.apache.druid.sql.calcite.aggregation.DimensionExpression) Aggregate(org.apache.calcite.rel.core.Aggregate) RowSignature(org.apache.druid.segment.column.RowSignature) DimFilter(org.apache.druid.query.filter.DimFilter) Nonnull(javax.annotation.Nonnull)

Example 2 with Aggregation

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

the class Grouping method applyProject.

/**
 * Applies a post-grouping projection.
 *
 * @see DruidQuery#computeGrouping which uses this
 */
public Grouping applyProject(final PlannerContext plannerContext, final Project project) {
    final List<DimensionExpression> newDimensions = new ArrayList<>();
    final List<Aggregation> newAggregations = new ArrayList<>(aggregations);
    final Subtotals newSubtotals;
    final Projection postAggregationProjection = Projection.postAggregation(project, plannerContext, outputRowSignature, "p");
    postAggregationProjection.getPostAggregators().forEach(postAggregator -> newAggregations.add(Aggregation.create(postAggregator)));
    // Remove literal dimensions that did not appear in the projection. This is useful for queries
    // like "SELECT COUNT(*) FROM tbl GROUP BY 'dummy'" which some tools can generate, and for which we don't
    // actually want to include a dimension 'dummy'.
    final ImmutableBitSet aggregateProjectBits = RelOptUtil.InputFinder.bits(project.getChildExps(), null);
    final int[] newDimIndexes = new int[dimensions.size()];
    boolean droppedDimensions = false;
    for (int i = 0; i < dimensions.size(); i++) {
        final DimensionExpression dimension = dimensions.get(i);
        if (Parser.parse(dimension.getDruidExpression().getExpression(), plannerContext.getExprMacroTable()).isLiteral() && !aggregateProjectBits.get(i)) {
            droppedDimensions = true;
            newDimIndexes[i] = -1;
        } else {
            newDimIndexes[i] = newDimensions.size();
            newDimensions.add(dimension);
        }
    }
    // Renumber subtotals, if needed, to account for removed dummy dimensions.
    if (newDimensions.size() != dimensions.size()) {
        final List<IntList> newSubtotalsList = new ArrayList<>();
        for (IntList subtotal : subtotals.getSubtotals()) {
            final IntList newSubtotal = new IntArrayList();
            for (int dimIndex : subtotal) {
                final int newDimIndex = newDimIndexes[dimIndex];
                if (newDimIndex >= 0) {
                    newSubtotal.add(newDimIndex);
                }
            }
            newSubtotalsList.add(newSubtotal);
        }
        newSubtotals = new Subtotals(newSubtotalsList);
    } else {
        newSubtotals = subtotals;
    }
    return Grouping.create(newDimensions, newSubtotals, newAggregations, havingFilter, postAggregationProjection.getOutputRowSignature(), droppedDimensions);
}
Also used : ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) ArrayList(java.util.ArrayList) IntArrayList(it.unimi.dsi.fastutil.ints.IntArrayList) DimensionExpression(org.apache.druid.sql.calcite.aggregation.DimensionExpression) IntList(it.unimi.dsi.fastutil.ints.IntList) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) IntArrayList(it.unimi.dsi.fastutil.ints.IntArrayList)

Example 3 with Aggregation

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

the class GroupingSqlAggregator 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) {
    List<String> arguments = aggregateCall.getArgList().stream().map(i -> getColumnName(plannerContext, rowSignature, project, virtualColumnRegistry, i)).filter(Objects::nonNull).collect(Collectors.toList());
    if (arguments.size() < aggregateCall.getArgList().size()) {
        return null;
    }
    for (Aggregation existing : existingAggregations) {
        for (AggregatorFactory factory : existing.getAggregatorFactories()) {
            if (!(factory instanceof GroupingAggregatorFactory)) {
                continue;
            }
            GroupingAggregatorFactory groupingFactory = (GroupingAggregatorFactory) factory;
            if (groupingFactory.getGroupings().equals(arguments) && groupingFactory.getName().equals(name)) {
                return Aggregation.create(groupingFactory);
            }
        }
    }
    AggregatorFactory factory = new GroupingAggregatorFactory(name, arguments);
    return Aggregation.create(factory);
}
Also used : Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) GroupingAggregatorFactory(org.apache.druid.query.aggregation.GroupingAggregatorFactory) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) GroupingAggregatorFactory(org.apache.druid.query.aggregation.GroupingAggregatorFactory) Nullable(javax.annotation.Nullable)

Example 4 with Aggregation

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

the class ArrayConcatSqlAggregator 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));
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    final ColumnType druidType = Calcites.getValueTypeForRelDataTypeFull(aggregateCall.getType());
    if (druidType == null || !druidType.isArray()) {
        // must be an array
        return null;
    }
    final String initialvalue = ExpressionType.fromColumnTypeStrict(druidType).asTypeString() + "[]";
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        VirtualColumn vc = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(plannerContext, arg, druidType);
        fieldName = vc.getOutputName();
    }
    if (aggregateCall.isDistinct()) {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_set_add_all(\"__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, false, false, StringUtils.format("array_concat(\"__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) ReturnTypes(org.apache.calcite.sql.type.ReturnTypes) 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) ImmutableSet(com.google.common.collect.ImmutableSet) SqlKind(org.apache.calcite.sql.SqlKind) 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) VirtualColumn(org.apache.druid.segment.VirtualColumn) 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) 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) VirtualColumn(org.apache.druid.segment.VirtualColumn) 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 5 with Aggregation

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

the class TDigestSketchQuantileSqlAggregator 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) {
    // This is expected to be a tdigest sketch
    final DruidExpression input = Aggregations.toDruidExpressionForNumericAggregator(plannerContext, rowSignature, Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(0)));
    if (input == null) {
        return null;
    }
    final AggregatorFactory aggregatorFactory;
    final String sketchName = StringUtils.format("%s:agg", name);
    // this is expected to be quantile fraction
    final RexNode quantileArg = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(1));
    if (!quantileArg.isA(SqlKind.LITERAL)) {
        // Quantile must be a literal in order to plan.
        return null;
    }
    final double quantile = ((Number) RexLiteral.value(quantileArg)).floatValue();
    Integer compression = TDigestSketchAggregatorFactory.DEFAULT_COMPRESSION;
    if (aggregateCall.getArgList().size() > 2) {
        final RexNode compressionArg = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(2));
        compression = ((Number) RexLiteral.value(compressionArg)).intValue();
    }
    // Look for existing matching aggregatorFactory.
    for (final Aggregation existing : existingAggregations) {
        for (AggregatorFactory factory : existing.getAggregatorFactories()) {
            if (factory instanceof TDigestSketchAggregatorFactory) {
                final boolean matches = TDigestSketchUtils.matchingAggregatorFactoryExists(virtualColumnRegistry, input, compression, (TDigestSketchAggregatorFactory) factory);
                if (matches) {
                    // Found existing one. Use this.
                    return Aggregation.create(ImmutableList.of(), new TDigestSketchToQuantilePostAggregator(name, new FieldAccessPostAggregator(factory.getName(), factory.getName()), quantile));
                }
            }
        }
    }
    // No existing match found. Create a new one.
    if (input.isDirectColumnAccess()) {
        aggregatorFactory = new TDigestSketchAggregatorFactory(sketchName, input.getDirectColumn(), compression);
    } else {
        String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(input, ColumnType.FLOAT);
        aggregatorFactory = new TDigestSketchAggregatorFactory(sketchName, virtualColumnName, compression);
    }
    return Aggregation.create(ImmutableList.of(aggregatorFactory), new TDigestSketchToQuantilePostAggregator(name, new FieldAccessPostAggregator(sketchName, sketchName), quantile));
}
Also used : Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) TDigestSketchAggregatorFactory(org.apache.druid.query.aggregation.tdigestsketch.TDigestSketchAggregatorFactory) FieldAccessPostAggregator(org.apache.druid.query.aggregation.post.FieldAccessPostAggregator) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) TDigestSketchToQuantilePostAggregator(org.apache.druid.query.aggregation.tdigestsketch.TDigestSketchToQuantilePostAggregator) TDigestSketchAggregatorFactory(org.apache.druid.query.aggregation.tdigestsketch.TDigestSketchAggregatorFactory) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

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