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

Example 6 with RexShuttle

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

the class WatermarkAssignerChangelogNormalizeTransposeRule method pushDownTransformedWatermark.

private RelNode pushDownTransformedWatermark(StreamPhysicalWatermarkAssigner watermark, StreamPhysicalCalc calc, StreamPhysicalChangelogNormalize changelogNormalize, StreamPhysicalExchange exchange, Mappings.TargetMapping calcMapping, RexBuilder rexBuilder) {
    Mappings.TargetMapping inversedMapping = calcMapping.inverse();
    final int newRowTimeFieldIndex = inversedMapping.getTargetOpt(watermark.rowtimeFieldIndex());
    // Updates watermark properties after push down before Calc
    // 1. rewrites watermark expression
    // 2. clears distribution
    // 3. updates row time field index
    RexNode newWatermarkExpr = watermark.watermarkExpr();
    if (watermark.watermarkExpr() != null) {
        newWatermarkExpr = RexUtil.apply(inversedMapping, watermark.watermarkExpr());
    }
    final RelNode newWatermark = watermark.copy(watermark.getTraitSet().plus(FlinkRelDistribution.DEFAULT()), exchange.getInput(), newRowTimeFieldIndex, newWatermarkExpr);
    final RelNode newChangelogNormalize = buildTreeInOrder(newWatermark, Tuple2.of(exchange, exchange.getTraitSet()), Tuple2.of(changelogNormalize, changelogNormalize.getTraitSet()));
    // Rewrites Calc program because the field type of row time
    // field is changed after watermark pushed down
    final RexProgram oldProgram = calc.getProgram();
    final RexProgramBuilder programBuilder = new RexProgramBuilder(newChangelogNormalize.getRowType(), rexBuilder);
    final Function<RexNode, RexNode> rexShuttle = e -> e.accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            if (inputRef.getIndex() == newRowTimeFieldIndex) {
                return RexInputRef.of(newRowTimeFieldIndex, newChangelogNormalize.getRowType());
            } else {
                return inputRef;
            }
        }
    });
    oldProgram.getNamedProjects().forEach(pair -> programBuilder.addProject(rexShuttle.apply(oldProgram.expandLocalRef(pair.left)), pair.right));
    if (oldProgram.getCondition() != null) {
        programBuilder.addCondition(rexShuttle.apply(oldProgram.expandLocalRef(oldProgram.getCondition())));
    }
    final RexProgram newProgram = programBuilder.getProgram();
    return calc.copy(calc.getTraitSet(), newChangelogNormalize, newProgram);
}

Example 7 with RexShuttle

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

the class RelTimeIndicatorConverter method visitJoin.

private RelNode visitJoin(FlinkLogicalJoin join) {
    RelNode newLeft = join.getLeft().accept(this);
    RelNode newRight = join.getRight().accept(this);
    int leftFieldCount = newLeft.getRowType().getFieldCount();
    // temporal table join
    if (TemporalJoinUtil.satisfyTemporalJoin(join, newLeft, newRight)) {
        RelNode rewrittenTemporalJoin = join.copy(join.getTraitSet(), join.getCondition(), newLeft, newRight, join.getJoinType(), join.isSemiJoinDone());
        // Materialize all of the time attributes from the right side of temporal join
        Set<Integer> rightIndices = IntStream.range(0, newRight.getRowType().getFieldCount()).mapToObj(startIdx -> leftFieldCount + startIdx).collect(Collectors.toSet());
        return createCalcToMaterializeTimeIndicators(rewrittenTemporalJoin, rightIndices);
    } else {
        if (JoinUtil.satisfyRegularJoin(join, newLeft, newRight)) {
            // materialize time attribute fields of regular join's inputs
            newLeft = materializeTimeIndicators(newLeft);
            newRight = materializeTimeIndicators(newRight);
        }
        List<RelDataTypeField> leftRightFields = new ArrayList<>();
        leftRightFields.addAll(newLeft.getRowType().getFieldList());
        leftRightFields.addAll(newRight.getRowType().getFieldList());
        RexNode newCondition = join.getCondition().accept(new RexShuttle() {

            @Override
            public RexNode visitInputRef(RexInputRef inputRef) {
                if (isTimeIndicatorType(inputRef.getType())) {
                    return RexInputRef.of(inputRef.getIndex(), leftRightFields);
                } else {
                    return super.visitInputRef(inputRef);
                }
            }
        });
        return FlinkLogicalJoin.create(newLeft, newRight, newCondition, join.getJoinType());
    }
}

Example 8 with RexShuttle

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

the class RelParameterizerShuttle method bindRel.

private RelNode bindRel(RelNode node, RexBuilder builder, RelDataTypeFactory typeFactory) {
    final RexShuttle binder = new RexShuttle() {

        @Override
        public RexNode visitDynamicParam(RexDynamicParam dynamicParam) {
            return bind(dynamicParam, builder, typeFactory);
        }
    };
    node = node.accept(binder);
    node.childrenAccept(new RelVisitor() {

        @Override
        public void visit(RelNode node, int ordinal, RelNode parent) {
            super.visit(node, ordinal, parent);
            RelNode transformed = node.accept(binder);
            if (!node.equals(transformed)) {
                parent.replaceInput(ordinal, transformed);
            }
        }
    });
    return node;
}

Example 9 with RexShuttle

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

the class JoinNestedLoopPhysicalRel method joinInfo.

public JetJoinInfo joinInfo(QueryParameterMetadata parameterMetadata) {
    List<Integer> leftKeys = analyzeCondition().leftKeys.toIntegerList();
    List<Integer> rightKeys = analyzeCondition().rightKeys.toIntegerList();
    HazelcastTable table = OptUtils.extractHazelcastTable(getRight());
    RexBuilder rexBuilder = getCluster().getRexBuilder();
    List<RexNode> additionalNonEquiConditions = new ArrayList<>();
    for (int i = 0; i < rightKeys.size(); i++) {
        Integer rightKeyIndex = rightKeys.get(i);
        RexNode rightExpr = table.getProjects().get(rightKeyIndex);
        if (rightExpr instanceof RexInputRef) {
            rightKeys.set(i, ((RexInputRef) rightExpr).getIndex());
        } else {
            // offset the indices in rightExp by the width of left row
            rightExpr = rightExpr.accept(new RexShuttle() {

                @Override
                public RexNode visitInputRef(RexInputRef inputRef) {
                    return rexBuilder.makeInputRef(inputRef.getType(), inputRef.getIndex() + getLeft().getRowType().getFieldCount());
                }
            });
            additionalNonEquiConditions.add(rexBuilder.makeCall(HazelcastSqlOperatorTable.EQUALS, rexBuilder.makeInputRef(getLeft(), leftKeys.get(i)), rightExpr));
            leftKeys.remove(i);
            rightKeys.remove(i);
            i--;
        }
    }
    Expression<Boolean> nonEquiCondition = filter(schema(parameterMetadata), RexUtil.composeConjunction(rexBuilder, asList(analyzeCondition().getRemaining(rexBuilder), RexUtil.composeConjunction(rexBuilder, additionalNonEquiConditions))), parameterMetadata);
    Expression<Boolean> condition = filter(schema(parameterMetadata), getCondition(), parameterMetadata);
    return new JetJoinInfo(getJoinType(), toIntArray(leftKeys), toIntArray(rightKeys), nonEquiCondition, condition);
}

Example 10 with RexShuttle

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

the class HiveSubQueryRemoveRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final RelNode relNode = call.rel(0);
    final RelBuilder builder = call.builder();
    // if subquery is in FILTER
    if (relNode instanceof HiveFilter) {
        final HiveFilter filter = call.rel(0);
        final RexSubQuery e = RexUtil.SubQueryFinder.find(filter.getCondition());
        assert e != null;
        final RelOptUtil.Logic logic = LogicVisitor.find(RelOptUtil.Logic.TRUE, ImmutableList.of(filter.getCondition()), e);
        builder.push(filter.getInput());
        final int fieldCount = builder.peek().getRowType().getFieldCount();
        SubqueryConf subqueryConfig = filter.getCluster().getPlanner().getContext().unwrap(SubqueryConf.class);
        boolean isCorrScalarQuery = subqueryConfig.getCorrScalarRexSQWithAgg().contains(e.rel);
        final RexNode target = apply(call.getMetadataQuery(), e, HiveFilter.getVariablesSet(e), logic, builder, 1, fieldCount, isCorrScalarQuery);
        final RexShuttle shuttle = new ReplaceSubQueryShuttle(e, target);
        builder.filter(shuttle.apply(filter.getCondition()));
        builder.project(fields(builder, filter.getRowType().getFieldCount()));
        RelNode newRel = builder.build();
        call.transformTo(newRel);
    } else if (relNode instanceof HiveProject) {
        // if subquery is in PROJECT
        final HiveProject project = call.rel(0);
        final RexSubQuery e = RexUtil.SubQueryFinder.find(project.getProjects());
        assert e != null;
        final RelOptUtil.Logic logic = LogicVisitor.find(RelOptUtil.Logic.TRUE_FALSE_UNKNOWN, project.getProjects(), e);
        builder.push(project.getInput());
        final int fieldCount = builder.peek().getRowType().getFieldCount();
        SubqueryConf subqueryConfig = project.getCluster().getPlanner().getContext().unwrap(SubqueryConf.class);
        boolean isCorrScalarQuery = subqueryConfig.getCorrScalarRexSQWithAgg().contains(e.rel);
        final RexNode target = apply(call.getMetadataQuery(), e, HiveFilter.getVariablesSet(e), logic, builder, 1, fieldCount, isCorrScalarQuery);
        final RexShuttle shuttle = new ReplaceSubQueryShuttle(e, target);
        builder.project(shuttle.apply(project.getProjects()), project.getRowType().getFieldNames());
        call.transformTo(builder.build());
    }
}