Examples with DruidExpression org.apache.druid.sql.calcite.expression.DruidExpression used on opensource projects

Example 36 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression 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;
}

Example 37 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression in project druid by druid-io.

the class Projection method postAggregationHandleOtherKinds.

private static void postAggregationHandleOtherKinds(final Project project, final PlannerContext plannerContext, final RowSignature inputRowSignature, final RexNode postAggregatorRexNode, final List<String> rowOrder, final PostAggregatorVisitor postAggregatorVisitor) {
    PostAggregator pagg = OperatorConversions.toPostAggregator(plannerContext, inputRowSignature, postAggregatorRexNode, postAggregatorVisitor);
    if (pagg != null) {
        postAggregatorVisitor.addPostAgg(pagg);
        rowOrder.add(pagg.getName());
    } else {
        final DruidExpression postAggregatorExpression = Expressions.toDruidExpressionWithPostAggOperands(plannerContext, inputRowSignature, postAggregatorRexNode, postAggregatorVisitor);
        if (postAggregatorExpression == null) {
            throw new CannotBuildQueryException(project, postAggregatorRexNode);
        }
        handlePostAggregatorExpression(plannerContext, inputRowSignature, postAggregatorRexNode, rowOrder, postAggregatorVisitor, postAggregatorExpression);
    }
}

Example 38 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression in project druid by druid-io.

the class Projection method preAggregation.

public static Projection preAggregation(final Project project, final PlannerContext plannerContext, final RowSignature inputRowSignature, final VirtualColumnRegistry virtualColumnRegistry) {
    final List<DruidExpression> expressions = new ArrayList<>();
    for (final RexNode rexNode : project.getChildExps()) {
        final DruidExpression expression = Expressions.toDruidExpression(plannerContext, inputRowSignature, rexNode);
        if (expression == null) {
            throw new CannotBuildQueryException(project, rexNode);
        } else {
            expressions.add(expression);
        }
    }
    final Set<String> virtualColumns = new HashSet<>();
    final List<String> rowOrder = new ArrayList<>();
    for (int i = 0; i < expressions.size(); i++) {
        final DruidExpression expression = expressions.get(i);
        final RelDataType dataType = project.getRowType().getFieldList().get(i).getType();
        if (expression.isDirectColumnAccess() && Objects.equals(inputRowSignature.getColumnType(expression.getDirectColumn()).orElse(null), Calcites.getColumnTypeForRelDataType(dataType))) {
            // Refer to column directly when it's a direct access with matching type.
            rowOrder.add(expression.getDirectColumn());
        } else {
            String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(expression, project.getChildExps().get(i).getType());
            virtualColumns.add(virtualColumnName);
            rowOrder.add(virtualColumnName);
        }
    }
    return new Projection(null, ImmutableList.copyOf(virtualColumns), RowSignatures.fromRelDataType(rowOrder, project.getRowType()));
}

Example 39 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression 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 40 with DruidExpression

use of org.apache.druid.sql.calcite.expression.DruidExpression 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))));
}