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

Example 6 with DruidExpression

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

Example 7 with DruidExpression

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

the class AvgSqlAggregator 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<DruidExpression> arguments = Aggregations.getArgumentsForSimpleAggregator(plannerContext, rowSignature, aggregateCall, project);
    if (arguments == null) {
        return null;
    }
    final String countName = Calcites.makePrefixedName(name, "count");
    final AggregatorFactory count = CountSqlAggregator.createCountAggregatorFactory(countName, plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, aggregateCall, project);
    final String fieldName;
    final String expression;
    final DruidExpression arg = Iterables.getOnlyElement(arguments);
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final ValueType sumType;
    // Use 64-bit sum regardless of the type of the AVG aggregator.
    if (SqlTypeName.INT_TYPES.contains(aggregateCall.getType().getSqlTypeName())) {
        sumType = ValueType.LONG;
    } else {
        sumType = ValueType.DOUBLE;
    }
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
        expression = null;
    } else {
        // if the filter or anywhere else defined a virtual column for us, re-use it
        final RexNode resolutionArg = Expressions.fromFieldAccess(rowSignature, project, Iterables.getOnlyElement(aggregateCall.getArgList()));
        String vc = virtualColumnRegistry.getVirtualColumnByExpression(arg, resolutionArg.getType());
        fieldName = vc != null ? vc : null;
        expression = vc != null ? null : arg.getExpression();
    }
    final String sumName = Calcites.makePrefixedName(name, "sum");
    final AggregatorFactory sum = SumSqlAggregator.createSumAggregatorFactory(sumType, sumName, fieldName, expression, macroTable);
    return Aggregation.create(ImmutableList.of(sum, count), new ArithmeticPostAggregator(name, "quotient", ImmutableList.of(new FieldAccessPostAggregator(null, sumName), new FieldAccessPostAggregator(null, countName))));
}

Example 8 with DruidExpression

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

the class MultiColumnSqlAggregator 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) {
    if (aggregateCall.isDistinct()) {
        return null;
    }
    final List<DruidExpression> arguments = Aggregations.getArgumentsForSimpleAggregator(plannerContext, rowSignature, aggregateCall, project);
    if (arguments == null) {
        return null;
    }
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final List<FieldInfo> fieldInfoList = new ArrayList<>();
    // Convert arguments to concise field information
    for (DruidExpression argument : arguments) {
        if (argument.isDirectColumnAccess()) {
            fieldInfoList.add(FieldInfo.fromFieldName(argument.getDirectColumn()));
        } else {
            fieldInfoList.add(FieldInfo.fromExpression(argument.getExpression()));
        }
    }
    Preconditions.checkArgument(!fieldInfoList.isEmpty(), "FieldInfoList should not be empty");
    return getAggregation(name, aggregateCall, macroTable, fieldInfoList);
}

Example 9 with DruidExpression

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

the class DruidRexExecutor method reduce.

@Override
public void reduce(final RexBuilder rexBuilder, final List<RexNode> constExps, final List<RexNode> reducedValues) {
    for (RexNode constExp : constExps) {
        final DruidExpression druidExpression = Expressions.toDruidExpression(plannerContext, EMPTY_ROW_SIGNATURE, constExp);
        if (druidExpression == null) {
            reducedValues.add(constExp);
        } else {
            final SqlTypeName sqlTypeName = constExp.getType().getSqlTypeName();
            final Expr expr = Parser.parse(druidExpression.getExpression(), plannerContext.getExprMacroTable());
            final ExprEval exprResult = expr.eval(InputBindings.forFunction(name -> {
                // Sanity check. Bindings should not be used for a constant expression.
                throw new UnsupportedOperationException();
            }));
            final RexNode literal;
            if (sqlTypeName == SqlTypeName.BOOLEAN) {
                literal = rexBuilder.makeLiteral(exprResult.asBoolean(), constExp.getType(), true);
            } else if (sqlTypeName == SqlTypeName.DATE) {
                // ExprEval.isNumericNull checks whether the parsed primitive value is null or not.
                if (!constExp.getType().isNullable() && exprResult.isNumericNull()) {
                    throw new UnsupportedSQLQueryException("Illegal DATE constant: %s", constExp);
                }
                literal = rexBuilder.makeDateLiteral(Calcites.jodaToCalciteDateString(DateTimes.utc(exprResult.asLong()), plannerContext.getTimeZone()));
            } else if (sqlTypeName == SqlTypeName.TIMESTAMP) {
                // ExprEval.isNumericNull checks whether the parsed primitive value is null or not.
                if (!constExp.getType().isNullable() && exprResult.isNumericNull()) {
                    throw new UnsupportedSQLQueryException("Illegal TIMESTAMP constant: %s", constExp);
                }
                literal = rexBuilder.makeTimestampLiteral(Calcites.jodaToCalciteTimestampString(DateTimes.utc(exprResult.asLong()), plannerContext.getTimeZone()), RelDataType.PRECISION_NOT_SPECIFIED);
            } else if (SqlTypeName.NUMERIC_TYPES.contains(sqlTypeName)) {
                final BigDecimal bigDecimal;
                if (exprResult.isNumericNull()) {
                    literal = rexBuilder.makeNullLiteral(constExp.getType());
                } else {
                    if (exprResult.type().is(ExprType.LONG)) {
                        bigDecimal = BigDecimal.valueOf(exprResult.asLong());
                    } else {
                        // if exprResult evaluates to Nan or infinity, this will throw a NumberFormatException.
                        // If you find yourself in such a position, consider casting the literal to a BIGINT so that
                        // the query can execute.
                        double exprResultDouble = exprResult.asDouble();
                        if (Double.isNaN(exprResultDouble) || Double.isInfinite(exprResultDouble)) {
                            String expression = druidExpression.getExpression();
                            throw new UnsupportedSQLQueryException("'%s' evaluates to '%s' that is not supported in SQL. You can either cast the expression as bigint ('cast(%s as bigint)') or char ('cast(%s as char)') or change the expression itself", expression, Double.toString(exprResultDouble), expression, expression);
                        }
                        bigDecimal = BigDecimal.valueOf(exprResult.asDouble());
                    }
                    literal = rexBuilder.makeLiteral(bigDecimal, constExp.getType(), true);
                }
            } else if (sqlTypeName == SqlTypeName.ARRAY) {
                assert exprResult.isArray();
                if (SqlTypeName.NUMERIC_TYPES.contains(constExp.getType().getComponentType().getSqlTypeName())) {
                    if (exprResult.type().getElementType().is(ExprType.LONG)) {
                        List<BigDecimal> resultAsBigDecimalList = Arrays.stream(exprResult.asLongArray()).map(BigDecimal::valueOf).collect(Collectors.toList());
                        literal = rexBuilder.makeLiteral(resultAsBigDecimalList, constExp.getType(), true);
                    } else {
                        List<BigDecimal> resultAsBigDecimalList = Arrays.stream(exprResult.asDoubleArray()).map(doubleVal -> {
                            if (Double.isNaN(doubleVal) || Double.isInfinite(doubleVal)) {
                                String expression = druidExpression.getExpression();
                                throw new UnsupportedSQLQueryException("'%s' contains an element that evaluates to '%s' which is not supported in SQL. You can either cast the element in the array to bigint or char or change the expression itself", expression, Double.toString(doubleVal));
                            }
                            return BigDecimal.valueOf(doubleVal);
                        }).collect(Collectors.toList());
                        literal = rexBuilder.makeLiteral(resultAsBigDecimalList, constExp.getType(), true);
                    }
                } else {
                    literal = rexBuilder.makeLiteral(Arrays.asList(exprResult.asArray()), constExp.getType(), true);
                }
            } else if (sqlTypeName == SqlTypeName.OTHER && constExp.getType() instanceof RowSignatures.ComplexSqlType) {
                // complex constant is not reducible, so just leave it as an expression
                literal = constExp;
            } else {
                if (exprResult.isArray()) {
                    // just leave array expressions on multi-value strings alone, we're going to push them down into a virtual
                    // column selector anyway
                    literal = constExp;
                } else {
                    literal = rexBuilder.makeLiteral(exprResult.value(), constExp.getType(), true);
                }
            }
            reducedValues.add(literal);
        }
    }
}

Example 10 with DruidExpression

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

the class TimeCeilOperatorConversion method toDruidExpression.

@Override
@Nullable
public DruidExpression toDruidExpression(final PlannerContext plannerContext, final RowSignature rowSignature, final RexNode rexNode) {
    final RexCall call = (RexCall) rexNode;
    final List<DruidExpression> functionArgs = TimeFloorOperatorConversion.toTimestampFloorOrCeilArgs(plannerContext, rowSignature, call.getOperands());
    if (functionArgs == null) {
        return null;
    }
    return DruidExpression.ofFunctionCall(Calcites.getColumnTypeForRelDataType(rexNode.getType()), "timestamp_ceil", functionArgs);
}