Search in sources :

Example 1 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project calcite by apache.

the class AbstractMaterializedViewRule method replaceWithOriginalReferences.

/**
 * Given the input expression, it will replace (sub)expressions when possible
 * using the content of the mapping. In particular, the mapping contains the
 * digest of the expression and the index that the replacement input ref should
 * point to.
 */
private static RexNode replaceWithOriginalReferences(final RexBuilder rexBuilder, final RelNode node, final NodeLineage nodeLineage, final RexNode exprToRewrite) {
    // Currently we allow the following:
    // 1) compensation pred can be directly map to expression
    // 2) all references in compensation pred can be map to expressions
    // We support bypassing lossless casts.
    RexShuttle visitor = new RexShuttle() {

        @Override
        public RexNode visitCall(RexCall call) {
            RexNode rw = replace(call);
            return rw != null ? rw : super.visitCall(call);
        }

        @Override
        public RexNode visitTableInputRef(RexTableInputRef inputRef) {
            RexNode rw = replace(inputRef);
            return rw != null ? rw : super.visitTableInputRef(inputRef);
        }

        private RexNode replace(RexNode e) {
            Integer pos = nodeLineage.exprsLineage.get(e.toString());
            if (pos != null) {
                // Found it
                return rexBuilder.makeInputRef(node, pos);
            }
            pos = nodeLineage.exprsLineageLosslessCasts.get(e.toString());
            if (pos != null) {
                // Found it
                return rexBuilder.makeCast(e.getType(), rexBuilder.makeInputRef(node, pos));
            }
            return null;
        }
    };
    return visitor.apply(exprToRewrite);
}
Also used : RexCall(org.apache.calcite.rex.RexCall) RexShuttle(org.apache.calcite.rex.RexShuttle) RexTableInputRef(org.apache.calcite.rex.RexTableInputRef) RexNode(org.apache.calcite.rex.RexNode)

Example 2 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project calcite by apache.

the class JoinToCorrelateRule method onMatch.

public void onMatch(RelOptRuleCall call) {
    assert matches(call);
    final LogicalJoin join = call.rel(0);
    RelNode right = join.getRight();
    final RelNode left = join.getLeft();
    final int leftFieldCount = left.getRowType().getFieldCount();
    final RelOptCluster cluster = join.getCluster();
    final RexBuilder rexBuilder = cluster.getRexBuilder();
    final RelBuilder relBuilder = call.builder();
    final CorrelationId correlationId = cluster.createCorrel();
    final RexNode corrVar = rexBuilder.makeCorrel(left.getRowType(), correlationId);
    final ImmutableBitSet.Builder requiredColumns = ImmutableBitSet.builder();
    // Replace all references of left input with FieldAccess(corrVar, field)
    final RexNode joinCondition = join.getCondition().accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef input) {
            int field = input.getIndex();
            if (field >= leftFieldCount) {
                return rexBuilder.makeInputRef(input.getType(), input.getIndex() - leftFieldCount);
            }
            requiredColumns.set(field);
            return rexBuilder.makeFieldAccess(corrVar, field);
        }
    });
    relBuilder.push(right).filter(joinCondition);
    RelNode newRel = LogicalCorrelate.create(left, relBuilder.build(), correlationId, requiredColumns.build(), SemiJoinType.of(join.getJoinType()));
    call.transformTo(newRel);
}
Also used : RelOptCluster(org.apache.calcite.plan.RelOptCluster) RelBuilder(org.apache.calcite.tools.RelBuilder) RexShuttle(org.apache.calcite.rex.RexShuttle) ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) RelNode(org.apache.calcite.rel.RelNode) LogicalJoin(org.apache.calcite.rel.logical.LogicalJoin) RexBuilder(org.apache.calcite.rex.RexBuilder) RexInputRef(org.apache.calcite.rex.RexInputRef) CorrelationId(org.apache.calcite.rel.core.CorrelationId) RexNode(org.apache.calcite.rex.RexNode)

Example 3 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project calcite by apache.

the class ProjectWindowTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final LogicalProject project = call.rel(0);
    final LogicalWindow window = call.rel(1);
    final RelOptCluster cluster = window.getCluster();
    final List<RelDataTypeField> rowTypeWindowInput = window.getInput().getRowType().getFieldList();
    final int windowInputColumn = rowTypeWindowInput.size();
    // Record the window input columns which are actually referred
    // either in the LogicalProject above LogicalWindow or LogicalWindow itself
    // (Note that the constants used in LogicalWindow are not considered here)
    final ImmutableBitSet beReferred = findReference(project, window);
    // it is impossible to trim anyone of them out
    if (beReferred.cardinality() == windowInputColumn) {
        return;
    }
    // Put a DrillProjectRel below LogicalWindow
    final List<RexNode> exps = new ArrayList<>();
    final RelDataTypeFactory.Builder builder = cluster.getTypeFactory().builder();
    // Keep only the fields which are referred
    for (int index : BitSets.toIter(beReferred)) {
        final RelDataTypeField relDataTypeField = rowTypeWindowInput.get(index);
        exps.add(new RexInputRef(index, relDataTypeField.getType()));
        builder.add(relDataTypeField);
    }
    final LogicalProject projectBelowWindow = new LogicalProject(cluster, window.getTraitSet(), window.getInput(), exps, builder.build());
    // Create a new LogicalWindow with necessary inputs only
    final List<Window.Group> groups = new ArrayList<>();
    // As the un-referred columns are trimmed by the LogicalProject,
    // the indices specified in LogicalWindow would need to be adjusted
    final RexShuttle indexAdjustment = new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            final int newIndex = getAdjustedIndex(inputRef.getIndex(), beReferred, windowInputColumn);
            return new RexInputRef(newIndex, inputRef.getType());
        }

        @Override
        public RexNode visitCall(final RexCall call) {
            if (call instanceof Window.RexWinAggCall) {
                boolean[] update = { false };
                final List<RexNode> clonedOperands = visitList(call.operands, update);
                if (update[0]) {
                    return new Window.RexWinAggCall((SqlAggFunction) call.getOperator(), call.getType(), clonedOperands, ((Window.RexWinAggCall) call).ordinal, ((Window.RexWinAggCall) call).distinct);
                } else {
                    return call;
                }
            } else {
                return super.visitCall(call);
            }
        }
    };
    int aggCallIndex = windowInputColumn;
    final RelDataTypeFactory.Builder outputBuilder = cluster.getTypeFactory().builder();
    outputBuilder.addAll(projectBelowWindow.getRowType().getFieldList());
    for (Window.Group group : window.groups) {
        final ImmutableBitSet.Builder keys = ImmutableBitSet.builder();
        final List<RelFieldCollation> orderKeys = new ArrayList<>();
        final List<Window.RexWinAggCall> aggCalls = new ArrayList<>();
        // Adjust keys
        for (int index : group.keys) {
            keys.set(getAdjustedIndex(index, beReferred, windowInputColumn));
        }
        // Adjust orderKeys
        for (RelFieldCollation relFieldCollation : group.orderKeys.getFieldCollations()) {
            final int index = relFieldCollation.getFieldIndex();
            orderKeys.add(relFieldCollation.copy(getAdjustedIndex(index, beReferred, windowInputColumn)));
        }
        // Adjust Window Functions
        for (Window.RexWinAggCall rexWinAggCall : group.aggCalls) {
            aggCalls.add((Window.RexWinAggCall) rexWinAggCall.accept(indexAdjustment));
            final RelDataTypeField relDataTypeField = window.getRowType().getFieldList().get(aggCallIndex);
            outputBuilder.add(relDataTypeField);
            ++aggCallIndex;
        }
        groups.add(new Window.Group(keys.build(), group.isRows, group.lowerBound, group.upperBound, RelCollations.of(orderKeys), aggCalls));
    }
    final LogicalWindow newLogicalWindow = LogicalWindow.create(window.getTraitSet(), projectBelowWindow, window.constants, outputBuilder.build(), groups);
    // Modify the top LogicalProject
    final List<RexNode> topProjExps = new ArrayList<>();
    for (RexNode rexNode : project.getChildExps()) {
        topProjExps.add(rexNode.accept(indexAdjustment));
    }
    final LogicalProject newTopProj = project.copy(newLogicalWindow.getTraitSet(), newLogicalWindow, topProjExps, project.getRowType());
    if (ProjectRemoveRule.isTrivial(newTopProj)) {
        call.transformTo(newLogicalWindow);
    } else {
        call.transformTo(newTopProj);
    }
}
Also used : RelOptCluster(org.apache.calcite.plan.RelOptCluster) ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) ArrayList(java.util.ArrayList) RexCall(org.apache.calcite.rex.RexCall) LogicalWindow(org.apache.calcite.rel.logical.LogicalWindow) RelDataTypeFactory(org.apache.calcite.rel.type.RelDataTypeFactory) Window(org.apache.calcite.rel.core.Window) LogicalWindow(org.apache.calcite.rel.logical.LogicalWindow) RexShuttle(org.apache.calcite.rex.RexShuttle) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RelFieldCollation(org.apache.calcite.rel.RelFieldCollation) RexInputRef(org.apache.calcite.rex.RexInputRef) LogicalProject(org.apache.calcite.rel.logical.LogicalProject) RexNode(org.apache.calcite.rex.RexNode)

Example 4 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project calcite by apache.

the class ReduceDecimalsRule method onMatch.

// implement RelOptRule
public void onMatch(RelOptRuleCall call) {
    LogicalCalc calc = call.rel(0);
    // Expand decimals in every expression in this program. If no
    // expression changes, don't apply the rule.
    final RexProgram program = calc.getProgram();
    if (!RexUtil.requiresDecimalExpansion(program, true)) {
        return;
    }
    final RexBuilder rexBuilder = calc.getCluster().getRexBuilder();
    final RexShuttle shuttle = new DecimalShuttle(rexBuilder);
    RexProgramBuilder programBuilder = RexProgramBuilder.create(rexBuilder, calc.getInput().getRowType(), program.getExprList(), program.getProjectList(), program.getCondition(), program.getOutputRowType(), shuttle, true);
    final RexProgram newProgram = programBuilder.getProgram();
    LogicalCalc newCalc = LogicalCalc.create(calc.getInput(), newProgram);
    call.transformTo(newCalc);
}
Also used : RexShuttle(org.apache.calcite.rex.RexShuttle) RexProgram(org.apache.calcite.rex.RexProgram) RexBuilder(org.apache.calcite.rex.RexBuilder) LogicalCalc(org.apache.calcite.rel.logical.LogicalCalc) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder)

Example 5 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project flink by apache.

the class PushWatermarkIntoTableSourceScanAcrossCalcRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    FlinkLogicalWatermarkAssigner watermarkAssigner = call.rel(0);
    FlinkLogicalCalc calc = call.rel(1);
    RexProgram originProgram = calc.getProgram();
    List<RexNode> projectList = originProgram.getProjectList().stream().map(originProgram::expandLocalRef).collect(Collectors.toList());
    // get watermark expression
    RexNode rowTimeColumn = projectList.get(watermarkAssigner.rowtimeFieldIndex());
    RexNode newWatermarkExpr = watermarkAssigner.watermarkExpr().accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            return projectList.get(inputRef.getIndex());
        }
    });
    // push watermark assigner into the scan
    FlinkLogicalTableSourceScan newScan = getNewScan(watermarkAssigner, newWatermarkExpr, call.rel(2), ShortcutUtils.unwrapContext(calc).getTableConfig(), // useWatermarkAssignerRowType
    false);
    FlinkTypeFactory typeFactory = ShortcutUtils.unwrapTypeFactory(calc);
    RexBuilder builder = call.builder().getRexBuilder();
    // cast timestamp/timestamp_ltz type to rowtime type.
    RexNode newRowTimeColumn = builder.makeReinterpretCast(typeFactory.createRowtimeIndicatorType(rowTimeColumn.getType().isNullable(), rowTimeColumn.getType().getSqlTypeName() == SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE), rowTimeColumn, null);
    // build new calc program
    RexProgramBuilder programBuilder = new RexProgramBuilder(newScan.getRowType(), builder);
    List<String> outputFieldNames = originProgram.getOutputRowType().getFieldNames();
    for (int i = 0; i < projectList.size(); i++) {
        if (i == watermarkAssigner.rowtimeFieldIndex()) {
            // replace the origin computed column to keep type consistent
            programBuilder.addProject(newRowTimeColumn, outputFieldNames.get(i));
        } else {
            programBuilder.addProject(projectList.get(i), outputFieldNames.get(i));
        }
    }
    if (originProgram.getCondition() != null) {
        programBuilder.addCondition(originProgram.expandLocalRef(originProgram.getCondition()));
    }
    FlinkLogicalCalc newCalc = FlinkLogicalCalc.create(newScan, programBuilder.getProgram());
    call.transformTo(newCalc);
}
Also used : RexShuttle(org.apache.calcite.rex.RexShuttle) FlinkLogicalCalc(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc) RexProgram(org.apache.calcite.rex.RexProgram) FlinkLogicalTableSourceScan(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalTableSourceScan) FlinkLogicalWatermarkAssigner(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalWatermarkAssigner) FlinkTypeFactory(org.apache.flink.table.planner.calcite.FlinkTypeFactory) RexInputRef(org.apache.calcite.rex.RexInputRef) RexBuilder(org.apache.calcite.rex.RexBuilder) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) RexNode(org.apache.calcite.rex.RexNode)

Aggregations

RexShuttle (org.apache.calcite.rex.RexShuttle)20 RexNode (org.apache.calcite.rex.RexNode)16 RexInputRef (org.apache.calcite.rex.RexInputRef)13 ArrayList (java.util.ArrayList)9 RelNode (org.apache.calcite.rel.RelNode)8 RexCall (org.apache.calcite.rex.RexCall)7 RexBuilder (org.apache.calcite.rex.RexBuilder)6 RelOptUtil (org.apache.calcite.plan.RelOptUtil)5 List (java.util.List)4 Map (java.util.Map)4 RexLiteral (org.apache.calcite.rex.RexLiteral)4 RexLocalRef (org.apache.calcite.rex.RexLocalRef)4 RexProgram (org.apache.calcite.rex.RexProgram)4 RexProgramBuilder (org.apache.calcite.rex.RexProgramBuilder)4 Collections (java.util.Collections)3 LinkedList (java.util.LinkedList)3 RelDataTypeField (org.apache.calcite.rel.type.RelDataTypeField)3 RelBuilder (org.apache.calcite.tools.RelBuilder)3 HashMap (java.util.HashMap)2 Function (java.util.function.Function)2