use of org.apache.druid.sql.calcite.planner.PlannerContext in project druid by druid-io.
the class TimeFloorOperatorConversion method toTimestampFloorOrCeilArgs.
/**
* Function that converts SQL TIME_FLOOR or TIME_CEIL args to Druid expression "timestamp_floor" or "timestamp_ceil"
* args. The main reason this function is necessary is because the handling of origin and timezone must take into
* account the SQL context timezone. It also helps with handling SQL FLOOR and CEIL, by offering handling of
* TimeUnitRange args.
*/
@Nullable
public static List<DruidExpression> toTimestampFloorOrCeilArgs(final PlannerContext plannerContext, final RowSignature rowSignature, final List<RexNode> operands) {
final List<DruidExpression> functionArgs = new ArrayList<>();
// Timestamp
functionArgs.add(Expressions.toDruidExpression(plannerContext, rowSignature, operands.get(0)));
// Period
final RexNode periodOperand = operands.get(1);
if (periodOperand.isA(SqlKind.LITERAL) && RexLiteral.value(periodOperand) instanceof TimeUnitRange) {
// TimeUnitRange literals are used by FLOOR(t TO unit) and CEIL(t TO unit)
final Period period = TimeUnits.toPeriod((TimeUnitRange) RexLiteral.value(periodOperand));
if (period == null) {
// Unrecognized time unit, bail out.
return null;
}
functionArgs.add(DruidExpression.ofStringLiteral(period.toString()));
} else {
// Other literal types are used by TIME_FLOOR and TIME_CEIL
functionArgs.add(Expressions.toDruidExpression(plannerContext, rowSignature, periodOperand));
}
// Origin
functionArgs.add(OperatorConversions.getOperandWithDefault(operands, 2, operand -> {
if (operand.isA(SqlKind.LITERAL)) {
return DruidExpression.ofLiteral(Calcites.getColumnTypeForRelDataType(operand.getType()), DruidExpression.numberLiteral(Calcites.calciteDateTimeLiteralToJoda(operand, plannerContext.getTimeZone()).getMillis()));
} else {
return Expressions.toDruidExpression(plannerContext, rowSignature, operand);
}
}, DruidExpression.ofLiteral(null, DruidExpression.nullLiteral())));
// Time zone
functionArgs.add(OperatorConversions.getOperandWithDefault(operands, 3, operand -> Expressions.toDruidExpression(plannerContext, rowSignature, operand), DruidExpression.ofStringLiteral(plannerContext.getTimeZone().getID())));
return functionArgs.stream().noneMatch(Objects::isNull) ? functionArgs : null;
}
use of org.apache.druid.sql.calcite.planner.PlannerContext 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));
}
}
use of org.apache.druid.sql.calcite.planner.PlannerContext 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))));
}
use of org.apache.druid.sql.calcite.planner.PlannerContext 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);
}
use of org.apache.druid.sql.calcite.planner.PlannerContext 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));
}
}
Aggregations