Examples with RexCall org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall used on opensource projects

Example 86 with RexCall

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall in project druid by druid-io.

the class TimeArithmeticOperatorConversion method toDruidExpression.

@Override
public DruidExpression toDruidExpression(final PlannerContext plannerContext, final RowSignature rowSignature, final RexNode rexNode) {
    final RexCall call = (RexCall) rexNode;
    final List<RexNode> operands = call.getOperands();
    if (operands.size() != 2) {
        throw new IAE("Expected 2 args, got %s", operands.size());
    }
    final RexNode leftRexNode = operands.get(0);
    final RexNode rightRexNode = operands.get(1);
    final DruidExpression leftExpr = Expressions.toDruidExpression(plannerContext, rowSignature, leftRexNode);
    final DruidExpression rightExpr = Expressions.toDruidExpression(plannerContext, rowSignature, rightRexNode);
    if (leftExpr == null || rightExpr == null) {
        return null;
    }
    final ColumnType outputType = Calcites.getColumnTypeForRelDataType(rexNode.getType());
    if (rightRexNode.getType().getFamily() == SqlTypeFamily.INTERVAL_YEAR_MONTH) {
        // Period is a value in months.
        return DruidExpression.ofExpression(outputType, DruidExpression.functionCall("timestamp_shift"), ImmutableList.of(leftExpr, rightExpr.map(simpleExtraction -> null, expression -> rightRexNode.isA(SqlKind.LITERAL) ? StringUtils.format("'P%sM'", RexLiteral.value(rightRexNode)) : StringUtils.format("concat('P', %s, 'M')", expression)), DruidExpression.ofLiteral(ColumnType.LONG, DruidExpression.numberLiteral(direction > 0 ? 1 : -1)), DruidExpression.ofStringLiteral(plannerContext.getTimeZone().getID())));
    } else if (rightRexNode.getType().getFamily() == SqlTypeFamily.INTERVAL_DAY_TIME) {
        // Period is a value in milliseconds. Ignore time zone.
        return DruidExpression.ofExpression(outputType, (args) -> StringUtils.format("(%s %s %s)", args.get(0).getExpression(), direction > 0 ? "+" : "-", args.get(1).getExpression()), ImmutableList.of(leftExpr, rightExpr));
    } else if ((leftRexNode.getType().getFamily() == SqlTypeFamily.TIMESTAMP || leftRexNode.getType().getFamily() == SqlTypeFamily.DATE) && (rightRexNode.getType().getFamily() == SqlTypeFamily.TIMESTAMP || rightRexNode.getType().getFamily() == SqlTypeFamily.DATE)) {
        // Calcite represents both TIMESTAMP - INTERVAL and TIMESTAMPDIFF (TIMESTAMP - TIMESTAMP)
        // with a MINUS_DATE operator, so we must tell which case we're in by checking the type of
        // the second argument.
        Preconditions.checkState(direction < 0, "Time arithmetic require direction < 0");
        if (call.getType().getFamily() == SqlTypeFamily.INTERVAL_YEAR_MONTH) {
            return DruidExpression.ofExpression(outputType, DruidExpression.functionCall("subtract_months"), ImmutableList.of(leftExpr, rightExpr, DruidExpression.ofStringLiteral(plannerContext.getTimeZone().getID())));
        } else {
            return DruidExpression.ofExpression(outputType, (args) -> StringUtils.format("(%s %s %s)", args.get(0).getExpression(), "-", args.get(1).getExpression()), ImmutableList.of(leftExpr, rightExpr));
        }
    } else {
        // Shouldn't happen if subclasses are behaving.
        throw new ISE("Got unexpected type period type family[%s]", rightRexNode.getType().getFamily());
    }
}

Example 87 with RexCall

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall in project druid by druid-io.

the class TimeFormatOperatorConversion method toDruidExpression.

@Override
public DruidExpression toDruidExpression(final PlannerContext plannerContext, final RowSignature rowSignature, final RexNode rexNode) {
    final RexCall call = (RexCall) rexNode;
    final RexNode timeArg = call.getOperands().get(0);
    final DruidExpression timeExpression = Expressions.toDruidExpression(plannerContext, rowSignature, timeArg);
    if (timeExpression == null) {
        return null;
    }
    final String pattern = OperatorConversions.getOperandWithDefault(call.getOperands(), 1, RexLiteral::stringValue, DEFAULT_PATTERN);
    final DateTimeZone timeZone = OperatorConversions.getOperandWithDefault(call.getOperands(), 2, operand -> {
        try {
            return DateTimes.inferTzFromString(RexLiteral.stringValue(operand), false);
        } catch (IllegalArgumentException e) {
            throw new IAE(e.getMessage());
        }
    }, plannerContext.getTimeZone());
    return DruidExpression.ofFunctionCall(Calcites.getColumnTypeForRelDataType(rexNode.getType()), "timestamp_format", ImmutableList.of(timeExpression, DruidExpression.ofStringLiteral(pattern), DruidExpression.ofStringLiteral(timeZone.getID())));
}

Example 88 with RexCall

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall in project druid by druid-io.

the class TimeShiftOperatorConversion method toDruidExpression.

@Override
public DruidExpression toDruidExpression(final PlannerContext plannerContext, final RowSignature rowSignature, final RexNode rexNode) {
    final RexCall call = (RexCall) rexNode;
    final DruidExpression timeExpression = Expressions.toDruidExpression(plannerContext, rowSignature, call.getOperands().get(0));
    final DruidExpression periodExpression = Expressions.toDruidExpression(plannerContext, rowSignature, call.getOperands().get(1));
    final DruidExpression stepExpression = Expressions.toDruidExpression(plannerContext, rowSignature, call.getOperands().get(2));
    if (timeExpression == null || periodExpression == null || stepExpression == null) {
        return null;
    }
    final DateTimeZone timeZone = OperatorConversions.getOperandWithDefault(call.getOperands(), 3, operand -> DateTimes.inferTzFromString(RexLiteral.stringValue(operand)), plannerContext.getTimeZone());
    return DruidExpression.ofFunctionCall(Calcites.getColumnTypeForRelDataType(rexNode.getType()), "timestamp_shift", ImmutableList.of(timeExpression, periodExpression, stepExpression, DruidExpression.ofStringLiteral(timeZone.getID())));
}

Example 89 with RexCall

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall in project druid by druid-io.

the class ArrayContainsOperatorConversion method toDruidFilter.

@Nullable
@Override
public DimFilter toDruidFilter(final PlannerContext plannerContext, RowSignature rowSignature, @Nullable VirtualColumnRegistry virtualColumnRegistry, final RexNode rexNode) {
    final List<RexNode> operands = ((RexCall) rexNode).getOperands();
    final List<DruidExpression> druidExpressions = Expressions.toDruidExpressions(plannerContext, rowSignature, operands);
    if (druidExpressions == null) {
        return null;
    }
    // Converts array_contains() function into an AND of Selector filters if possible.
    final DruidExpression leftExpr = druidExpressions.get(0);
    final DruidExpression rightExpr = druidExpressions.get(1);
    if (leftExpr.isSimpleExtraction()) {
        Expr expr = Parser.parse(rightExpr.getExpression(), plannerContext.getExprMacroTable());
        // different package.
        if (expr.isLiteral()) {
            // Evaluate the expression to get out the array elements.
            // We can safely pass a noop ObjectBinding if the expression is literal.
            ExprEval<?> exprEval = expr.eval(InputBindings.nilBindings());
            String[] arrayElements = exprEval.asStringArray();
            if (arrayElements == null || arrayElements.length == 0) {
                // to create an empty array with no argument, we just return null.
                return null;
            } else if (arrayElements.length == 1) {
                return newSelectorDimFilter(leftExpr.getSimpleExtraction(), arrayElements[0]);
            } else {
                final List<DimFilter> selectFilters = Arrays.stream(arrayElements).map(val -> newSelectorDimFilter(leftExpr.getSimpleExtraction(), val)).collect(Collectors.toList());
                return new AndDimFilter(selectFilters);
            }
        }
    }
    return toExpressionFilter(plannerContext, getDruidFunctionName(), druidExpressions);
}

Example 90 with RexCall

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall in project druid by druid-io.

the class ArrayOverlapOperatorConversion method toDruidFilter.

@Nullable
@Override
public DimFilter toDruidFilter(final PlannerContext plannerContext, RowSignature rowSignature, @Nullable VirtualColumnRegistry virtualColumnRegistry, final RexNode rexNode) {
    final List<RexNode> operands = ((RexCall) rexNode).getOperands();
    final List<DruidExpression> druidExpressions = Expressions.toDruidExpressions(plannerContext, rowSignature, operands);
    if (druidExpressions == null) {
        return null;
    }
    // Converts array_overlaps() function into an OR of Selector filters if possible.
    final boolean leftSimpleExtractionExpr = druidExpressions.get(0).isSimpleExtraction();
    final boolean rightSimpleExtractionExpr = druidExpressions.get(1).isSimpleExtraction();
    final DruidExpression simpleExtractionExpr;
    final DruidExpression complexExpr;
    if (leftSimpleExtractionExpr ^ rightSimpleExtractionExpr) {
        if (leftSimpleExtractionExpr) {
            simpleExtractionExpr = druidExpressions.get(0);
            complexExpr = druidExpressions.get(1);
        } else {
            simpleExtractionExpr = druidExpressions.get(1);
            complexExpr = druidExpressions.get(0);
        }
    } else {
        return toExpressionFilter(plannerContext, getDruidFunctionName(), druidExpressions);
    }
    Expr expr = Parser.parse(complexExpr.getExpression(), plannerContext.getExprMacroTable());
    if (expr.isLiteral()) {
        // Evaluate the expression to take out the array elements.
        // We can safely pass null if the expression is literal.
        ExprEval<?> exprEval = expr.eval(InputBindings.nilBindings());
        String[] arrayElements = exprEval.asStringArray();
        if (arrayElements == null || arrayElements.length == 0) {
            // to create an empty array with no argument, we just return null.
            return null;
        } else if (arrayElements.length == 1) {
            return newSelectorDimFilter(simpleExtractionExpr.getSimpleExtraction(), arrayElements[0]);
        } else {
            return new InDimFilter(simpleExtractionExpr.getSimpleExtraction().getColumn(), Sets.newHashSet(arrayElements), simpleExtractionExpr.getSimpleExtraction().getExtractionFn(), null);
        }
    }
    return toExpressionFilter(plannerContext, getDruidFunctionName(), druidExpressions);
}