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

Example 71 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project drill by axbaretto.

the class DrillReduceAggregatesRule method reduceAgg.

private RexNode reduceAgg(Aggregate oldAggRel, AggregateCall oldCall, List<AggregateCall> newCalls, Map<AggregateCall, RexNode> aggCallMapping, List<RexNode> inputExprs) {
    final SqlAggFunction sqlAggFunction = DrillCalciteWrapperUtility.extractSqlOperatorFromWrapper(oldCall.getAggregation());
    if (sqlAggFunction instanceof SqlSumAggFunction) {
        // case COUNT(x) when 0 then null else SUM0(x) end
        return reduceSum(oldAggRel, oldCall, newCalls, aggCallMapping);
    }
    if (sqlAggFunction instanceof SqlAvgAggFunction) {
        final SqlKind subtype = sqlAggFunction.getKind();
        switch(subtype) {
            case AVG:
                // replace original AVG(x) with SUM(x) / COUNT(x)
                return reduceAvg(oldAggRel, oldCall, newCalls, aggCallMapping);
            case STDDEV_POP:
                // / COUNT(x))
                return reduceStddev(oldAggRel, oldCall, true, true, newCalls, aggCallMapping, inputExprs);
            case STDDEV_SAMP:
                // / CASE COUNT(x) WHEN 1 THEN NULL ELSE COUNT(x) - 1 END)
                return reduceStddev(oldAggRel, oldCall, false, true, newCalls, aggCallMapping, inputExprs);
            case VAR_POP:
                // / COUNT(x)
                return reduceStddev(oldAggRel, oldCall, true, false, newCalls, aggCallMapping, inputExprs);
            case VAR_SAMP:
                // / CASE COUNT(x) WHEN 1 THEN NULL ELSE COUNT(x) - 1 END
                return reduceStddev(oldAggRel, oldCall, false, false, newCalls, aggCallMapping, inputExprs);
            default:
                throw Util.unexpected(subtype);
        }
    } else {
        // anything else:  preserve original call
        RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
        final int nGroups = oldAggRel.getGroupCount();
        List<RelDataType> oldArgTypes = new ArrayList<>();
        List<Integer> ordinals = oldCall.getArgList();
        assert ordinals.size() <= inputExprs.size();
        for (int ordinal : ordinals) {
            oldArgTypes.add(inputExprs.get(ordinal).getType());
        }
        return rexBuilder.addAggCall(oldCall, nGroups, oldAggRel.indicator, newCalls, aggCallMapping, oldArgTypes);
    }
}

Example 72 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project drill by axbaretto.

the class DrillMergeFilterRule method onMatch.

// ~ Methods ----------------------------------------------------------------
// implement RelOptRule
public void onMatch(RelOptRuleCall call) {
    Filter topFilter = call.rel(0);
    Filter bottomFilter = call.rel(1);
    // use RexPrograms to merge the two FilterRels into a single program
    // so we can convert the two FilterRel conditions to directly
    // reference the bottom FilterRel's child
    RexBuilder rexBuilder = topFilter.getCluster().getRexBuilder();
    RexProgram bottomProgram = createProgram(bottomFilter);
    RexProgram topProgram = createProgram(topFilter);
    RexProgram mergedProgram = RexProgramBuilder.mergePrograms(topProgram, bottomProgram, rexBuilder);
    RexNode newCondition = mergedProgram.expandLocalRef(mergedProgram.getCondition());
    // if(!RexUtil.isFlat(newCondition)){
    // RexCall newCall = (RexCall) newCondition;
    // newCondition = rexBuilder.makeFlatCall( newCall.getOperator(), newCall.getOperands());
    // }
    Filter newFilterRel = (Filter) filterFactory.createFilter(bottomFilter.getInput(), RexUtil.flatten(rexBuilder, newCondition));
    call.transformTo(newFilterRel);
}

Example 73 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project drill by axbaretto.

the class HashToRandomExchangePrel method constructMuxPrel.

/**
 * This method creates a new UnorderedMux and Demux exchanges if mux operators are enabled.
 * @param child input to the new Unordered[Mux/Demux]Prel or new HashToRandomExchange node.
 * @param options options manager to check if mux is enabled.
 */
@Override
public Prel constructMuxPrel(Prel child, OptionManager options) {
    boolean isMuxEnabled = options.getOption(PlannerSettings.MUX_EXCHANGE.getOptionName()).bool_val;
    Prel newPrel = child;
    final List<String> childFields = child.getRowType().getFieldNames();
    List<RexNode> removeUpdatedExpr = null;
    if (isMuxEnabled) {
        // Insert Project Operator with new column that will be a hash for HashToRandomExchange fields
        final List<DistributionField> distFields = getFields();
        final List<String> outputFieldNames = Lists.newArrayList(childFields);
        final RexBuilder rexBuilder = getCluster().getRexBuilder();
        final List<RelDataTypeField> childRowTypeFields = child.getRowType().getFieldList();
        final HashPrelUtil.HashExpressionCreatorHelper<RexNode> hashHelper = new HashPrelUtil.RexNodeBasedHashExpressionCreatorHelper(rexBuilder);
        final List<RexNode> distFieldRefs = Lists.newArrayListWithExpectedSize(distFields.size());
        for (int i = 0; i < distFields.size(); i++) {
            final int fieldId = distFields.get(i).getFieldId();
            distFieldRefs.add(rexBuilder.makeInputRef(childRowTypeFields.get(fieldId).getType(), fieldId));
        }
        final List<RexNode> updatedExpr = Lists.newArrayListWithExpectedSize(childRowTypeFields.size());
        removeUpdatedExpr = Lists.newArrayListWithExpectedSize(childRowTypeFields.size());
        for (RelDataTypeField field : childRowTypeFields) {
            RexNode rex = rexBuilder.makeInputRef(field.getType(), field.getIndex());
            updatedExpr.add(rex);
            removeUpdatedExpr.add(rex);
        }
        outputFieldNames.add(HashPrelUtil.HASH_EXPR_NAME);
        // distribution seed
        final RexNode distSeed = rexBuilder.makeBigintLiteral(BigDecimal.valueOf(HashPrelUtil.DIST_SEED));
        updatedExpr.add(HashPrelUtil.createHashBasedPartitionExpression(distFieldRefs, distSeed, hashHelper));
        RelDataType rowType = RexUtil.createStructType(getCluster().getTypeFactory(), updatedExpr, outputFieldNames);
        ProjectPrel addColumnprojectPrel = new ProjectPrel(child.getCluster(), child.getTraitSet(), child, updatedExpr, rowType);
        newPrel = new UnorderedMuxExchangePrel(addColumnprojectPrel.getCluster(), addColumnprojectPrel.getTraitSet(), addColumnprojectPrel);
    }
    newPrel = new HashToRandomExchangePrel(getCluster(), getTraitSet(), newPrel, getFields());
    if (options.getOption(PlannerSettings.DEMUX_EXCHANGE.getOptionName()).bool_val) {
        HashToRandomExchangePrel hashExchangePrel = (HashToRandomExchangePrel) newPrel;
        // Insert a DeMuxExchange to narrow down the number of receivers
        newPrel = new UnorderedDeMuxExchangePrel(getCluster(), getTraitSet(), hashExchangePrel, hashExchangePrel.getFields());
    }
    if (isMuxEnabled) {
        // remove earlier inserted Project Operator - since it creates issues down the road in HashJoin
        RelDataType removeRowType = RexUtil.createStructType(newPrel.getCluster().getTypeFactory(), removeUpdatedExpr, childFields);
        ProjectPrel removeColumnProjectPrel = new ProjectPrel(newPrel.getCluster(), newPrel.getTraitSet(), newPrel, removeUpdatedExpr, removeRowType);
        return removeColumnProjectPrel;
    }
    return newPrel;
}

Example 74 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project drill by axbaretto.

the class DefaultSqlHandler method addRenamedProject.

protected DrillRel addRenamedProject(DrillRel rel, RelDataType validatedRowType) {
    RelDataType t = rel.getRowType();
    RexBuilder b = rel.getCluster().getRexBuilder();
    List<RexNode> projections = Lists.newArrayList();
    int projectCount = t.getFieldList().size();
    for (int i = 0; i < projectCount; i++) {
        projections.add(b.makeInputRef(rel, i));
    }
    final List<String> fieldNames2 = SqlValidatorUtil.uniquify(validatedRowType.getFieldNames(), SqlValidatorUtil.EXPR_SUGGESTER, rel.getCluster().getTypeFactory().getTypeSystem().isSchemaCaseSensitive());
    RelDataType newRowType = RexUtil.createStructType(rel.getCluster().getTypeFactory(), projections, fieldNames2);
    DrillProjectRel topProj = DrillProjectRel.create(rel.getCluster(), rel.getTraitSet(), rel, projections, newRowType);
    // Add a final non-trivial Project to get the validatedRowType, if child is not project.
    if (rel instanceof Project && DrillRelOptUtil.isTrivialProject(topProj, true)) {
        return rel;
    } else {
        return topProj;
    }
}

Example 75 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project drill by axbaretto.

the class RewriteProjectToFlatten method visitProject.

@Override
public Prel visitProject(ProjectPrel node, Object unused) throws RelConversionException {
    ProjectPrel project = node;
    List<RexNode> exprList = new ArrayList<>();
    boolean rewrite = false;
    List<RelDataTypeField> relDataTypes = new ArrayList<>();
    int i = 0;
    RexNode flatttenExpr = null;
    for (RexNode rex : project.getChildExps()) {
        RexNode newExpr = rex;
        if (rex instanceof RexCall) {
            RexCall function = (RexCall) rex;
            String functionName = function.getOperator().getName();
            if (functionName.equalsIgnoreCase("flatten")) {
                rewrite = true;
                if (function.getOperands().size() != 1) {
                    throw new RelConversionException("Flatten expression expects a single input.");
                }
                newExpr = function.getOperands().get(0);
                RexBuilder builder = new RexBuilder(factory);
                flatttenExpr = builder.makeInputRef(new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory), i);
            }
        }
        relDataTypes.add(project.getRowType().getFieldList().get(i));
        i++;
        exprList.add(newExpr);
    }
    if (rewrite == true) {
        // TODO - figure out what is the right setting for the traits
        Prel newChild = ((Prel) project.getInput(0)).accept(this, null);
        ProjectPrel newProject = new ProjectPrel(node.getCluster(), project.getTraitSet(), newChild, exprList, new RelRecordType(relDataTypes));
        FlattenPrel flatten = new FlattenPrel(project.getCluster(), project.getTraitSet(), newProject, flatttenExpr);
        return flatten;
    }
    Prel child = ((Prel) project.getInput()).accept(this, null);
    return (Prel) project.copy(project.getTraitSet(), child, exprList, new RelRecordType(relDataTypes));
}