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