Examples with RexLiteral org.apache.calcite.rex.RexLiteral used on opensource projects

Example 6 with RexLiteral

use of org.apache.calcite.rex.RexLiteral 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.
   */
private static Aggregation translateAggregateCall(final PlannerContext plannerContext, final RowSignature sourceRowSignature, final Project project, final AggregateCall call, final DruidOperatorTable operatorTable, final List<Aggregation> existingAggregations, final int aggNumber, final boolean approximateCountDistinct) {
    final List<DimFilter> filters = Lists.newArrayList();
    final List<String> rowOrder = sourceRowSignature.getRowOrder();
    final String name = aggOutputName(aggNumber);
    final SqlKind kind = call.getAggregation().getKind();
    final SqlTypeName outputType = call.getType().getSqlTypeName();
    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 filter = Expressions.toFilter(operatorTable, plannerContext, sourceRowSignature, expression);
        if (filter == null) {
            return null;
        }
        filters.add(filter);
    }
    if (kind == SqlKind.COUNT && call.getArgList().isEmpty()) {
        // COUNT(*)
        return Aggregation.create(new CountAggregatorFactory(name)).filter(makeFilter(filters, sourceRowSignature));
    } else if (kind == SqlKind.COUNT && call.isDistinct()) {
        // COUNT(DISTINCT x)
        return approximateCountDistinct ? APPROX_COUNT_DISTINCT.toDruidAggregation(name, sourceRowSignature, operatorTable, plannerContext, existingAggregations, project, call, makeFilter(filters, sourceRowSignature)) : null;
    } else if (kind == SqlKind.COUNT || kind == SqlKind.SUM || kind == SqlKind.SUM0 || kind == SqlKind.MIN || kind == SqlKind.MAX || kind == SqlKind.AVG) {
        // Built-in agg, not distinct, not COUNT(*)
        boolean forceCount = false;
        final FieldOrExpression input;
        final int inputField = Iterables.getOnlyElement(call.getArgList());
        final RexNode rexNode = Expressions.fromFieldAccess(sourceRowSignature, project, inputField);
        final FieldOrExpression foe = FieldOrExpression.fromRexNode(operatorTable, plannerContext, rowOrder, rexNode);
        if (foe != null) {
            input = foe;
        } else if (rexNode.getKind() == SqlKind.CASE && ((RexCall) rexNode).getOperands().size() == 3) {
            // Possibly a CASE-style filtered aggregation. Styles supported:
            // A: SUM(CASE WHEN x = 'foo' THEN cnt END) => operands (x = 'foo', cnt, null)
            // B: SUM(CASE WHEN x = 'foo' THEN 1 ELSE 0 END) => operands (x = 'foo', 1, 0)
            // C: COUNT(CASE WHEN x = 'foo' THEN 'dummy' END) => operands (x = 'foo', 'dummy', null)
            // If the null and non-null args are switched, "flip" is set, which negates the filter.
            final RexCall caseCall = (RexCall) rexNode;
            final boolean flip = RexLiteral.isNullLiteral(caseCall.getOperands().get(1)) && !RexLiteral.isNullLiteral(caseCall.getOperands().get(2));
            final RexNode arg1 = caseCall.getOperands().get(flip ? 2 : 1);
            final RexNode arg2 = caseCall.getOperands().get(flip ? 1 : 2);
            // Operand 1: Filter
            final DimFilter filter = Expressions.toFilter(operatorTable, plannerContext, sourceRowSignature, caseCall.getOperands().get(0));
            if (filter == null) {
                return null;
            } else {
                filters.add(flip ? new NotDimFilter(filter) : filter);
            }
            if (call.getAggregation().getKind() == SqlKind.COUNT && arg1 instanceof RexLiteral && !RexLiteral.isNullLiteral(arg1) && RexLiteral.isNullLiteral(arg2)) {
                // Case C
                forceCount = true;
                input = null;
            } else if (call.getAggregation().getKind() == SqlKind.SUM && arg1 instanceof RexLiteral && ((Number) RexLiteral.value(arg1)).intValue() == 1 && arg2 instanceof RexLiteral && ((Number) RexLiteral.value(arg2)).intValue() == 0) {
                // Case B
                forceCount = true;
                input = null;
            } else if (RexLiteral.isNullLiteral(arg2)) {
                // Maybe case A
                input = FieldOrExpression.fromRexNode(operatorTable, plannerContext, rowOrder, arg1);
                if (input == null) {
                    return null;
                }
            } else {
                // Can't translate CASE into a filter.
                return null;
            }
        } else {
            // Can't translate operand.
            return null;
        }
        if (!forceCount) {
            Preconditions.checkNotNull(input, "WTF?! input was null for non-COUNT aggregation");
        }
        if (forceCount || kind == SqlKind.COUNT) {
            // COUNT(x)
            return Aggregation.create(new CountAggregatorFactory(name)).filter(makeFilter(filters, sourceRowSignature));
        } else {
            // Built-in aggregator that is not COUNT.
            final Aggregation retVal;
            final String fieldName = input.getFieldName();
            final String expression = input.getExpression();
            final boolean isLong = SqlTypeName.INT_TYPES.contains(outputType) || SqlTypeName.TIMESTAMP == outputType || SqlTypeName.DATE == outputType;
            if (kind == SqlKind.SUM || kind == SqlKind.SUM0) {
                retVal = isLong ? Aggregation.create(new LongSumAggregatorFactory(name, fieldName, expression)) : Aggregation.create(new DoubleSumAggregatorFactory(name, fieldName, expression));
            } else if (kind == SqlKind.MIN) {
                retVal = isLong ? Aggregation.create(new LongMinAggregatorFactory(name, fieldName, expression)) : Aggregation.create(new DoubleMinAggregatorFactory(name, fieldName, expression));
            } else if (kind == SqlKind.MAX) {
                retVal = isLong ? Aggregation.create(new LongMaxAggregatorFactory(name, fieldName, expression)) : Aggregation.create(new DoubleMaxAggregatorFactory(name, fieldName, expression));
            } else if (kind == SqlKind.AVG) {
                final String sumName = aggInternalName(aggNumber, "sum");
                final String countName = aggInternalName(aggNumber, "count");
                final AggregatorFactory sum = isLong ? new LongSumAggregatorFactory(sumName, fieldName, expression) : new DoubleSumAggregatorFactory(sumName, fieldName, expression);
                final AggregatorFactory count = new CountAggregatorFactory(countName);
                retVal = Aggregation.create(ImmutableList.of(sum, count), new ArithmeticPostAggregator(name, "quotient", ImmutableList.<PostAggregator>of(new FieldAccessPostAggregator(null, sumName), new FieldAccessPostAggregator(null, countName))));
            } else {
                // Not reached.
                throw new ISE("WTF?! Kind[%s] got into the built-in aggregator path somehow?!", kind);
            }
            return retVal.filter(makeFilter(filters, sourceRowSignature));
        }
    } else {
        // Not a built-in aggregator, check operator table.
        final SqlAggregator sqlAggregator = operatorTable.lookupAggregator(call.getAggregation().getName());
        return sqlAggregator != null ? sqlAggregator.toDruidAggregation(name, sourceRowSignature, operatorTable, plannerContext, existingAggregations, project, call, makeFilter(filters, sourceRowSignature)) : null;
    }
}

Example 7 with RexLiteral

use of org.apache.calcite.rex.RexLiteral in project druid by druid-io.

the class FloorExtractionOperator method convert.

@Override
public RowExtraction convert(final DruidOperatorTable operatorTable, final PlannerContext plannerContext, final List<String> rowOrder, final RexNode expression) {
    final RexCall call = (RexCall) expression;
    final RexNode arg = call.getOperands().get(0);
    final RowExtraction rex = Expressions.toRowExtraction(operatorTable, plannerContext, rowOrder, arg);
    if (rex == null) {
        return null;
    } else if (call.getOperands().size() == 1) {
        // FLOOR(expr)
        return RowExtraction.of(rex.getColumn(), ExtractionFns.compose(new BucketExtractionFn(1.0, 0.0), rex.getExtractionFn()));
    } else if (call.getOperands().size() == 2) {
        // FLOOR(expr TO timeUnit)
        final RexLiteral flag = (RexLiteral) call.getOperands().get(1);
        final TimeUnitRange timeUnit = (TimeUnitRange) flag.getValue();
        return applyTimestampFloor(rex, TimeUnits.toQueryGranularity(timeUnit, plannerContext.getTimeZone()));
    } else {
        // WTF? FLOOR with 3 arguments?
        return null;
    }
}

Example 8 with RexLiteral

use of org.apache.calcite.rex.RexLiteral in project drill by apache.

the class DrillReduceAggregatesRule method reduceStddev.

private RexNode reduceStddev(Aggregate oldAggRel, AggregateCall oldCall, boolean biased, boolean sqrt, List<AggregateCall> newCalls, Map<AggregateCall, RexNode> aggCallMapping, List<RexNode> inputExprs) {
    // stddev_pop(x) ==>
    //   power(
    //     (sum(x * x) - sum(x) * sum(x) / count(x))
    //     / count(x),
    //     .5)
    //
    // stddev_samp(x) ==>
    //   power(
    //     (sum(x * x) - sum(x) * sum(x) / count(x))
    //     / nullif(count(x) - 1, 0),
    //     .5)
    final PlannerSettings plannerSettings = (PlannerSettings) oldAggRel.getCluster().getPlanner().getContext();
    final boolean isInferenceEnabled = plannerSettings.isTypeInferenceEnabled();
    final int nGroups = oldAggRel.getGroupCount();
    RelDataTypeFactory typeFactory = oldAggRel.getCluster().getTypeFactory();
    final RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
    assert oldCall.getArgList().size() == 1 : oldCall.getArgList();
    final int argOrdinal = oldCall.getArgList().get(0);
    final RelDataType argType = getFieldType(oldAggRel.getInput(), argOrdinal);
    // final RexNode argRef = inputExprs.get(argOrdinal);
    RexNode argRef = rexBuilder.makeCall(CastHighOp, inputExprs.get(argOrdinal));
    inputExprs.set(argOrdinal, argRef);
    final RexNode argSquared = rexBuilder.makeCall(SqlStdOperatorTable.MULTIPLY, argRef, argRef);
    final int argSquaredOrdinal = lookupOrAdd(inputExprs, argSquared);
    final RelDataType sumType = typeFactory.createTypeWithNullability(argType, true);
    final AggregateCall sumArgSquaredAggCall = AggregateCall.create(new SqlSumAggFunction(sumType), oldCall.isDistinct(), ImmutableIntList.of(argSquaredOrdinal), -1, sumType, null);
    final RexNode sumArgSquared = rexBuilder.addAggCall(sumArgSquaredAggCall, nGroups, oldAggRel.indicator, newCalls, aggCallMapping, ImmutableList.of(argType));
    final AggregateCall sumArgAggCall = AggregateCall.create(new SqlSumAggFunction(sumType), oldCall.isDistinct(), ImmutableIntList.of(argOrdinal), -1, sumType, null);
    final RexNode sumArg = rexBuilder.addAggCall(sumArgAggCall, nGroups, oldAggRel.indicator, newCalls, aggCallMapping, ImmutableList.of(argType));
    final RexNode sumSquaredArg = rexBuilder.makeCall(SqlStdOperatorTable.MULTIPLY, sumArg, sumArg);
    final SqlCountAggFunction countAgg = (SqlCountAggFunction) SqlStdOperatorTable.COUNT;
    final RelDataType countType = countAgg.getReturnType(typeFactory);
    final AggregateCall countArgAggCall = AggregateCall.create(countAgg, oldCall.isDistinct(), oldCall.getArgList(), -1, countType, null);
    final RexNode countArg = rexBuilder.addAggCall(countArgAggCall, nGroups, oldAggRel.indicator, newCalls, aggCallMapping, ImmutableList.of(argType));
    final RexNode avgSumSquaredArg = rexBuilder.makeCall(SqlStdOperatorTable.DIVIDE, sumSquaredArg, countArg);
    final RexNode diff = rexBuilder.makeCall(SqlStdOperatorTable.MINUS, sumArgSquared, avgSumSquaredArg);
    final RexNode denominator;
    if (biased) {
        denominator = countArg;
    } else {
        final RexLiteral one = rexBuilder.makeExactLiteral(BigDecimal.ONE);
        final RexNode nul = rexBuilder.makeNullLiteral(countArg.getType().getSqlTypeName());
        final RexNode countMinusOne = rexBuilder.makeCall(SqlStdOperatorTable.MINUS, countArg, one);
        final RexNode countEqOne = rexBuilder.makeCall(SqlStdOperatorTable.EQUALS, countArg, one);
        denominator = rexBuilder.makeCall(SqlStdOperatorTable.CASE, countEqOne, nul, countMinusOne);
    }
    final SqlOperator divide;
    if (isInferenceEnabled) {
        divide = new DrillSqlOperator("divide", 2, true, oldCall.getType(), false);
    } else {
        divide = SqlStdOperatorTable.DIVIDE;
    }
    final RexNode div = rexBuilder.makeCall(divide, diff, denominator);
    RexNode result = div;
    if (sqrt) {
        final RexNode half = rexBuilder.makeExactLiteral(new BigDecimal("0.5"));
        result = rexBuilder.makeCall(SqlStdOperatorTable.POWER, div, half);
    }
    if (isInferenceEnabled) {
        return result;
    } else {
        /*
      * Currently calcite's strategy to infer the return type of aggregate functions
      * is wrong because it uses the first known argument to determine output type. For
      * instance if we are performing stddev on an integer column then it interprets the
      * output type to be integer which is incorrect as it should be double. So based on
      * this if we add cast after rewriting the aggregate we add an additional cast which
      * would cause wrong results. So we simply add a cast to ANY.
      */
        return rexBuilder.makeCast(typeFactory.createSqlType(SqlTypeName.ANY), result);
    }
}

Example 9 with RexLiteral

use of org.apache.calcite.rex.RexLiteral in project flink by apache.

the class FlinkRelDecorrelator method projectedLiteral.

/**
	 * Returns a literal output field, or null if it is not literal.
	 */
private static RexLiteral projectedLiteral(RelNode rel, int i) {
    if (rel instanceof Project) {
        final Project project = (Project) rel;
        final RexNode node = project.getProjects().get(i);
        if (node instanceof RexLiteral) {
            return (RexLiteral) node;
        }
    }
    return null;
}

Example 10 with RexLiteral

use of org.apache.calcite.rex.RexLiteral in project hive by apache.

the class ASTConverter method convert.

private ASTNode convert() throws CalciteSemanticException {
    /*
     * 1. Walk RelNode Graph; note from, where, gBy.. nodes.
     */
    new QBVisitor().go(root);
    /*
     * 2. convert from node.
     */
    QueryBlockInfo qb = convertSource(from);
    schema = qb.schema;
    hiveAST.from = ASTBuilder.construct(HiveParser.TOK_FROM, "TOK_FROM").add(qb.ast).node();
    /*
     * 3. convert filterNode
     */
    if (where != null) {
        ASTNode cond = where.getCondition().accept(new RexVisitor(schema));
        hiveAST.where = ASTBuilder.where(cond);
    }
    /*
     * 4. GBy
     */
    if (groupBy != null) {
        ASTBuilder b;
        boolean groupingSetsExpression = false;
        if (groupBy.indicator) {
            Group aggregateType = Aggregate.Group.induce(groupBy.getGroupSet(), groupBy.getGroupSets());
            if (aggregateType == Group.ROLLUP) {
                b = ASTBuilder.construct(HiveParser.TOK_ROLLUP_GROUPBY, "TOK_ROLLUP_GROUPBY");
            } else if (aggregateType == Group.CUBE) {
                b = ASTBuilder.construct(HiveParser.TOK_CUBE_GROUPBY, "TOK_CUBE_GROUPBY");
            } else {
                b = ASTBuilder.construct(HiveParser.TOK_GROUPING_SETS, "TOK_GROUPING_SETS");
                groupingSetsExpression = true;
            }
        } else {
            b = ASTBuilder.construct(HiveParser.TOK_GROUPBY, "TOK_GROUPBY");
        }
        HiveAggregate hiveAgg = (HiveAggregate) groupBy;
        for (int pos : hiveAgg.getAggregateColumnsOrder()) {
            RexInputRef iRef = new RexInputRef(groupBy.getGroupSet().nth(pos), groupBy.getCluster().getTypeFactory().createSqlType(SqlTypeName.ANY));
            b.add(iRef.accept(new RexVisitor(schema)));
        }
        for (int pos = 0; pos < groupBy.getGroupCount(); pos++) {
            if (!hiveAgg.getAggregateColumnsOrder().contains(pos)) {
                RexInputRef iRef = new RexInputRef(groupBy.getGroupSet().nth(pos), groupBy.getCluster().getTypeFactory().createSqlType(SqlTypeName.ANY));
                b.add(iRef.accept(new RexVisitor(schema)));
            }
        }
        //Grouping sets expressions
        if (groupingSetsExpression) {
            for (ImmutableBitSet groupSet : groupBy.getGroupSets()) {
                ASTBuilder expression = ASTBuilder.construct(HiveParser.TOK_GROUPING_SETS_EXPRESSION, "TOK_GROUPING_SETS_EXPRESSION");
                for (int i : groupSet) {
                    RexInputRef iRef = new RexInputRef(i, groupBy.getCluster().getTypeFactory().createSqlType(SqlTypeName.ANY));
                    expression.add(iRef.accept(new RexVisitor(schema)));
                }
                b.add(expression);
            }
        }
        if (!groupBy.getGroupSet().isEmpty()) {
            hiveAST.groupBy = b.node();
        }
        schema = new Schema(schema, groupBy);
    }
    /*
     * 5. Having
     */
    if (having != null) {
        ASTNode cond = having.getCondition().accept(new RexVisitor(schema));
        hiveAST.having = ASTBuilder.having(cond);
    }
    /*
     * 6. Project
     */
    ASTBuilder b = ASTBuilder.construct(HiveParser.TOK_SELECT, "TOK_SELECT");
    if (select instanceof Project) {
        if (select.getChildExps().isEmpty()) {
            RexLiteral r = select.getCluster().getRexBuilder().makeExactLiteral(new BigDecimal(1));
            ASTNode selectExpr = ASTBuilder.selectExpr(ASTBuilder.literal(r), "1");
            b.add(selectExpr);
        } else {
            int i = 0;
            for (RexNode r : select.getChildExps()) {
                if (RexUtil.isNull(r) && r.getType().getSqlTypeName() != SqlTypeName.NULL) {
                    // It is NULL value with different type, we need to introduce a CAST
                    // to keep it
                    r = select.getCluster().getRexBuilder().makeAbstractCast(r.getType(), r);
                }
                ASTNode expr = r.accept(new RexVisitor(schema, r instanceof RexLiteral));
                String alias = select.getRowType().getFieldNames().get(i++);
                ASTNode selectExpr = ASTBuilder.selectExpr(expr, alias);
                b.add(selectExpr);
            }
        }
        hiveAST.select = b.node();
    } else {
        // select is UDTF
        HiveTableFunctionScan udtf = (HiveTableFunctionScan) select;
        List<ASTNode> children = new ArrayList<>();
        RexCall call = (RexCall) udtf.getCall();
        for (RexNode r : call.getOperands()) {
            if (RexUtil.isNull(r) && r.getType().getSqlTypeName() != SqlTypeName.NULL) {
                // It is NULL value with different type, we need to introduce a CAST
                // to keep it
                r = select.getCluster().getRexBuilder().makeAbstractCast(r.getType(), r);
            }
            ASTNode expr = r.accept(new RexVisitor(schema, r instanceof RexLiteral));
            children.add(expr);
        }
        ASTBuilder sel = ASTBuilder.construct(HiveParser.TOK_SELEXPR, "TOK_SELEXPR");
        ASTNode function = buildUDTFAST(call.getOperator().getName(), children);
        sel.add(function);
        for (String alias : udtf.getRowType().getFieldNames()) {
            sel.add(HiveParser.Identifier, alias);
        }
        b.add(sel);
        hiveAST.select = b.node();
    }
    /*
     * 7. Order Use in Order By from the block above. RelNode has no pointer to
     * parent hence we need to go top down; but OB at each block really belong
     * to its src/from. Hence the need to pass in sort for each block from
     * its parent.
     * 8. Limit
     */
    convertOrderLimitToASTNode((HiveSortLimit) orderLimit);
    return hiveAST.getAST();
}