Examples with Aggregation org.apache.druid.sql.calcite.aggregation.Aggregation used on opensource projects

Example 11 with Aggregation

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

the class DruidQuery method computeAggregations.

/**
 * Returns aggregations corresponding to {@code aggregate.getAggCallList()}, in the same order.
 *
 * @param partialQuery          partial query
 * @param plannerContext        planner context
 * @param rowSignature          source row signature
 * @param virtualColumnRegistry re-usable virtual column references
 * @param rexBuilder            calcite RexBuilder
 * @param finalizeAggregations  true if this query should include explicit finalization for all of its
 *                              aggregators, where required. Useful for subqueries where Druid's native query layer
 *                              does not do this automatically.
 *
 * @return aggregations
 *
 * @throws CannotBuildQueryException if dimensions cannot be computed
 */
private static List<Aggregation> computeAggregations(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());
    final List<Aggregation> aggregations = new ArrayList<>();
    final String outputNamePrefix = Calcites.findUnusedPrefixForDigits("a", rowSignature.getColumnNames());
    for (int i = 0; i < aggregate.getAggCallList().size(); i++) {
        final String aggName = outputNamePrefix + i;
        final AggregateCall aggCall = aggregate.getAggCallList().get(i);
        final Aggregation aggregation = GroupByRules.translateAggregateCall(plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, partialQuery.getSelectProject(), aggregations, aggName, aggCall, finalizeAggregations);
        if (aggregation == null) {
            if (null == plannerContext.getPlanningError()) {
                plannerContext.setPlanningError("Aggregation [%s] is not supported", aggCall);
            }
            throw new CannotBuildQueryException(aggregate, aggCall);
        }
        aggregations.add(aggregation);
    }
    return aggregations;
}

Example 12 with Aggregation

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

the class DruidQuery method getVirtualColumns.

private VirtualColumns getVirtualColumns(final boolean includeDimensions) {
    // 'sourceRowSignature' could provide a list of all defined virtual columns while constructing a query, but we
    // still want to collect the set of VirtualColumns this way to ensure we only add what is still being used after
    // the various transforms and optimizations
    Set<VirtualColumn> virtualColumns = new HashSet<>();
    // rewrite any "specialized" virtual column expressions as top level virtual columns so that their native
    // implementation can be used instead of being composed as part of some expression tree in an expresson virtual
    // column
    Set<String> specialized = new HashSet<>();
    virtualColumnRegistry.visitAllSubExpressions((expression) -> {
        switch(expression.getType()) {
            case SPECIALIZED:
                // add the expression to the top level of the registry as a standalone virtual column
                final String name = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(expression, expression.getDruidType());
                specialized.add(name);
                // replace with an identifier expression of the new virtual column name
                return DruidExpression.ofColumn(expression.getDruidType(), name);
            default:
                // do nothing
                return expression;
        }
    });
    // we always want to add any virtual columns used by the query level DimFilter
    if (filter != null) {
        for (String columnName : filter.getRequiredColumns()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
            }
        }
    }
    if (selectProjection != null) {
        for (String columnName : selectProjection.getVirtualColumns()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
            }
        }
    }
    if (grouping != null) {
        if (includeDimensions) {
            for (DimensionExpression expression : grouping.getDimensions()) {
                if (virtualColumnRegistry.isVirtualColumnDefined(expression.getVirtualColumn())) {
                    virtualColumns.add(virtualColumnRegistry.getVirtualColumn(expression.getVirtualColumn()));
                }
            }
        }
        for (Aggregation aggregation : grouping.getAggregations()) {
            virtualColumns.addAll(virtualColumnRegistry.getAllVirtualColumns(aggregation.getRequiredColumns()));
        }
    }
    if (sorting != null && sorting.getProjection() != null && grouping == null) {
        for (String columnName : sorting.getProjection().getVirtualColumns()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
            }
        }
    }
    if (dataSource instanceof JoinDataSource) {
        for (String expression : ((JoinDataSource) dataSource).getVirtualColumnCandidates()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(expression)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(expression));
            }
        }
    }
    for (String columnName : specialized) {
        if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
            virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
        }
    }
    // sort for predictable output
    List<VirtualColumn> columns = new ArrayList<>(virtualColumns);
    columns.sort(Comparator.comparing(VirtualColumn::getOutputName));
    return VirtualColumns.create(columns);
}

Example 13 with Aggregation

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

the class StringSqlAggregator 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<DruidExpression> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).map(rexNode -> Expressions.toDruidExpression(plannerContext, rowSignature, rexNode)).collect(Collectors.toList());
    if (arguments.stream().anyMatch(Objects::isNull)) {
        return null;
    }
    RexNode separatorNode = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(1));
    if (!separatorNode.isA(SqlKind.LITERAL)) {
        // separator must be a literal
        return null;
    }
    String separator = RexLiteral.stringValue(separatorNode);
    if (separator == null) {
        // separator must not be null
        return null;
    }
    Integer maxSizeBytes = null;
    if (arguments.size() > 2) {
        RexNode maxBytes = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(2));
        if (!maxBytes.isA(SqlKind.LITERAL)) {
            // maxBytes must be a literal
            return null;
        }
        maxSizeBytes = ((Number) RexLiteral.value(maxBytes)).intValue();
    }
    final DruidExpression arg = arguments.get(0);
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String initialvalue = "[]";
    final ColumnType elementType = ColumnType.STRING;
    final String fieldName;
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        fieldName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(arg, elementType);
    }
    final String finalizer = StringUtils.format("if(array_length(o) == 0, null, array_to_string(o, '%s'))", separator);
    final NotDimFilter dimFilter = new NotDimFilter(new SelectorDimFilter(fieldName, null, null));
    if (aggregateCall.isDistinct()) {
        return Aggregation.create(// string_agg ignores nulls
        new FilteredAggregatorFactory(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_set_add(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name), null, finalizer, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable), dimFilter));
    } else {
        return Aggregation.create(// string_agg ignores nulls
        new FilteredAggregatorFactory(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_append(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_concat(\"__acc\", \"%s\")", name), null, finalizer, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable), dimFilter));
    }
}

Example 14 with Aggregation

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

the class BitwiseSqlAggregator 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<DruidExpression> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).map(rexNode -> Expressions.toDruidExpression(plannerContext, rowSignature, rexNode)).collect(Collectors.toList());
    if (arguments.stream().anyMatch(Objects::isNull)) {
        return null;
    }
    final DruidExpression arg = arguments.get(0);
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        fieldName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(arg, ColumnType.LONG);
    }
    return Aggregation.create(new FilteredAggregatorFactory(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, "0", null, null, false, false, StringUtils.format("%s(\"__acc\", \"%s\")", op.getDruidFunction(), fieldName), null, null, null, null, macroTable), new NotDimFilter(new SelectorDimFilter(fieldName, null, null))));
}

Example 15 with Aggregation

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

the class EarliestLatestBySqlAggregator 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 = EarliestLatestAnySqlAggregator.getColumnName(plannerContext, virtualColumnRegistry, args.get(0), rexNodes.get(0));
    final AggregatorFactory theAggFactory;
    switch(args.size()) {
        case 2:
            theAggFactory = aggregatorType.createAggregatorFactory(aggregatorName, fieldName, EarliestLatestAnySqlAggregator.getColumnName(plannerContext, virtualColumnRegistry, args.get(1), rexNodes.get(1)), outputType, -1);
            break;
        case 3:
            theAggFactory = aggregatorType.createAggregatorFactory(aggregatorName, fieldName, EarliestLatestAnySqlAggregator.getColumnName(plannerContext, virtualColumnRegistry, args.get(1), rexNodes.get(1)), outputType, RexLiteral.intValue(rexNodes.get(2)));
            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);
}