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);
}
}
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);
}
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);
}
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));
}
}
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));
}
Aggregations