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

Example 91 with RexInputRef

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

the class RelDecorrelator method getNewForOldInputRef.

private static RexInputRef getNewForOldInputRef(RelNode currentRel, Map<RelNode, Frame> map, RexInputRef oldInputRef) {
    assert currentRel != null;
    int oldOrdinal = oldInputRef.getIndex();
    int newOrdinal = 0;
    // determine which input rel oldOrdinal references, and adjust
    // oldOrdinal to be relative to that input rel
    RelNode oldInput = null;
    for (RelNode oldInput0 : currentRel.getInputs()) {
        RelDataType oldInputType = oldInput0.getRowType();
        int n = oldInputType.getFieldCount();
        if (oldOrdinal < n) {
            oldInput = oldInput0;
            break;
        }
        RelNode newInput = map.get(oldInput0).r;
        newOrdinal += newInput.getRowType().getFieldCount();
        oldOrdinal -= n;
    }
    assert oldInput != null;
    final Frame frame = map.get(oldInput);
    assert frame != null;
    // now oldOrdinal is relative to oldInput
    int oldLocalOrdinal = oldOrdinal;
    // figure out the newLocalOrdinal, relative to the newInput.
    int newLocalOrdinal = oldLocalOrdinal;
    if (!frame.oldToNewOutputs.isEmpty()) {
        newLocalOrdinal = frame.oldToNewOutputs.get(oldLocalOrdinal);
    }
    newOrdinal += newLocalOrdinal;
    return new RexInputRef(newOrdinal, frame.r.getRowType().getFieldList().get(newLocalOrdinal).getType());
}

Example 92 with RexInputRef

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

the class RelDecorrelator method decorrelateRel.

public Frame decorrelateRel(Correlate rel) {
    // 
    // Rewrite logic:
    // 
    // The original left input will be joined with the new right input that
    // has generated correlated variables propagated up. For any generated
    // corVars that are not used in the join key, pass them along to be
    // joined later with the Correlates that produce them.
    // 
    // the right input to Correlate should produce correlated variables
    final RelNode oldLeft = rel.getInput(0);
    final RelNode oldRight = rel.getInput(1);
    final Frame leftFrame = getInvoke(oldLeft, rel);
    final Frame rightFrame = getInvoke(oldRight, rel);
    if (leftFrame == null || rightFrame == null) {
        // If any input has not been rewritten, do not rewrite this rel.
        return null;
    }
    if (rightFrame.corDefOutputs.isEmpty()) {
        return null;
    }
    assert rel.getRequiredColumns().cardinality() <= rightFrame.corDefOutputs.keySet().size();
    // Change correlator rel into a join.
    // Join all the correlated variables produced by this correlator rel
    // with the values generated and propagated from the right input
    final SortedMap<CorDef, Integer> corDefOutputs = new TreeMap<>(rightFrame.corDefOutputs);
    final List<RexNode> conditions = new ArrayList<>();
    final List<RelDataTypeField> newLeftOutput = leftFrame.r.getRowType().getFieldList();
    int newLeftFieldCount = newLeftOutput.size();
    final List<RelDataTypeField> newRightOutput = rightFrame.r.getRowType().getFieldList();
    for (Map.Entry<CorDef, Integer> rightOutput : new ArrayList<>(corDefOutputs.entrySet())) {
        final CorDef corDef = rightOutput.getKey();
        if (!corDef.corr.equals(rel.getCorrelationId())) {
            continue;
        }
        final int newLeftPos = leftFrame.oldToNewOutputs.get(corDef.field);
        final int newRightPos = rightOutput.getValue();
        conditions.add(relBuilder.call(SqlStdOperatorTable.EQUALS, RexInputRef.of(newLeftPos, newLeftOutput), new RexInputRef(newLeftFieldCount + newRightPos, newRightOutput.get(newRightPos).getType())));
        // remove this corVar from output position mapping
        corDefOutputs.remove(corDef);
    }
    // vars that are not used in the join key.
    for (CorDef corDef : corDefOutputs.keySet()) {
        int newPos = corDefOutputs.get(corDef) + newLeftFieldCount;
        corDefOutputs.put(corDef, newPos);
    }
    // then add any corVar from the left input. Do not need to change
    // output positions.
    corDefOutputs.putAll(leftFrame.corDefOutputs);
    // Create the mapping between the output of the old correlation rel
    // and the new join rel
    final Map<Integer, Integer> mapOldToNewOutputs = new HashMap<>();
    int oldLeftFieldCount = oldLeft.getRowType().getFieldCount();
    int oldRightFieldCount = oldRight.getRowType().getFieldCount();
    // noinspection AssertWithSideEffects
    assert rel.getRowType().getFieldCount() == oldLeftFieldCount + oldRightFieldCount;
    // Left input positions are not changed.
    mapOldToNewOutputs.putAll(leftFrame.oldToNewOutputs);
    // Right input positions are shifted by newLeftFieldCount.
    for (int i = 0; i < oldRightFieldCount; i++) {
        mapOldToNewOutputs.put(i + oldLeftFieldCount, rightFrame.oldToNewOutputs.get(i) + newLeftFieldCount);
    }
    final RexNode condition = RexUtil.composeConjunction(relBuilder.getRexBuilder(), conditions);
    RelNode newJoin = relBuilder.push(leftFrame.r).push(rightFrame.r).join(rel.getJoinType(), condition).build();
    return register(rel, newJoin, mapOldToNewOutputs, corDefOutputs);
}

Example 93 with RexInputRef

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

the class RelDecorrelator method decorrelateInputWithValueGenerator.

private Frame decorrelateInputWithValueGenerator(RelNode rel, Frame frame) {
    // currently only handles one input
    assert rel.getInputs().size() == 1;
    RelNode oldInput = frame.r;
    final SortedMap<CorDef, Integer> corDefOutputs = new TreeMap<>(frame.corDefOutputs);
    final Collection<CorRef> corVarList = cm.mapRefRelToCorRef.get(rel);
    // This means that we do not need a value generator.
    if (rel instanceof Filter) {
        SortedMap<CorDef, Integer> map = new TreeMap<>();
        List<RexNode> projects = new ArrayList<>();
        for (CorRef correlation : corVarList) {
            final CorDef def = correlation.def();
            if (corDefOutputs.containsKey(def) || map.containsKey(def)) {
                continue;
            }
            try {
                findCorrelationEquivalent(correlation, ((Filter) rel).getCondition());
            } catch (Util.FoundOne e) {
                if (e.getNode() instanceof RexInputRef) {
                    map.put(def, ((RexInputRef) e.getNode()).getIndex());
                } else {
                    map.put(def, frame.r.getRowType().getFieldCount() + projects.size());
                    projects.add((RexNode) e.getNode());
                }
            }
        }
        // generator.
        if (map.size() == corVarList.size()) {
            map.putAll(frame.corDefOutputs);
            final RelNode r;
            if (!projects.isEmpty()) {
                relBuilder.push(oldInput).project(Iterables.concat(relBuilder.fields(), projects));
                r = relBuilder.build();
            } else {
                r = oldInput;
            }
            return register(rel.getInput(0), r, frame.oldToNewOutputs, map);
        }
    }
    int leftInputOutputCount = frame.r.getRowType().getFieldCount();
    // can directly add positions into corDefOutputs since join
    // does not change the output ordering from the inputs.
    RelNode valueGen = createValueGenerator(corVarList, leftInputOutputCount, corDefOutputs);
    RelNode join = relBuilder.push(frame.r).push(valueGen).join(JoinRelType.INNER, relBuilder.literal(true), ImmutableSet.of()).build();
    // Filter) are in the output and in the same position.
    return register(rel.getInput(0), join, frame.oldToNewOutputs, corDefOutputs);
}

Example 94 with RexInputRef

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

the class CodeSplitTest method testJoinCondition.

@Test
public void testJoinCondition() {
    int numFields = 200;
    FlinkTypeFactory typeFactory = FlinkTypeFactory.INSTANCE();
    RexBuilder builder = new RexBuilder(typeFactory);
    RelDataType intType = typeFactory.createFieldTypeFromLogicalType(new IntType());
    RexNode[] conditions = new RexNode[numFields];
    for (int i = 0; i < numFields; i++) {
        conditions[i] = builder.makeCall(SqlStdOperatorTable.LESS_THAN, new RexInputRef(i, intType), new RexInputRef(numFields + i, intType));
    }
    RexNode joinCondition = builder.makeCall(SqlStdOperatorTable.AND, conditions);
    RowType rowType = getIntRowType(numFields);
    GenericRowData rowData1 = new GenericRowData(numFields);
    GenericRowData rowData2 = new GenericRowData(numFields);
    Random random = new Random();
    for (int i = 0; i < numFields; i++) {
        rowData1.setField(i, 0);
        rowData2.setField(i, 1);
    }
    boolean result = random.nextBoolean();
    if (!result) {
        rowData1.setField(random.nextInt(numFields), 1);
    }
    Consumer<TableConfig> consumer = tableConfig -> {
        JoinCondition instance = JoinUtil.generateConditionFunction(tableConfig, joinCondition, rowType, rowType).newInstance(classLoader);
        for (int i = 0; i < 100; i++) {
            Assert.assertEquals(result, instance.apply(rowData1, rowData2));
        }
    };
    runTest(consumer);
}

Example 95 with RexInputRef

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

the class ArrayScanToJoinConverter method convert.

/**
 * Returns a LogicJoin.
 */
@Override
public RelNode convert(ResolvedArrayScan zetaNode, List<RelNode> inputs) {
    List<RexNode> projects = new ArrayList<>();
    RelNode leftInput = inputs.get(0);
    ResolvedColumnRef columnRef = (ResolvedColumnRef) zetaNode.getArrayExpr();
    CorrelationId correlationId = getCluster().createCorrel();
    getCluster().getQuery().mapCorrel(correlationId.getName(), leftInput);
    String columnName = String.format("%s%s", zetaNode.getElementColumn().getTableName(), zetaNode.getElementColumn().getName());
    projects.add(getCluster().getRexBuilder().makeFieldAccess(getCluster().getRexBuilder().makeCorrel(leftInput.getRowType(), correlationId), getExpressionConverter().indexOfProjectionColumnRef(columnRef.getColumn().getId(), zetaNode.getInputScan().getColumnList())));
    RelNode projectNode = LogicalProject.create(createOneRow(getCluster()), ImmutableList.of(), projects, ImmutableList.of(columnName));
    // Create an UnCollect
    boolean ordinality = (zetaNode.getArrayOffsetColumn() != null);
    // If they aren't true we need the Project to reorder columns.
    assert zetaNode.getElementColumn().getId() == 1;
    assert !ordinality || zetaNode.getArrayOffsetColumn().getColumn().getId() == 2;
    ZetaSqlUnnest uncollectNode = ZetaSqlUnnest.create(projectNode.getTraitSet(), projectNode, ordinality);
    List<RexInputRef> rightProjects = new ArrayList<>();
    List<String> rightNames = new ArrayList<>();
    rightProjects.add(getCluster().getRexBuilder().makeInputRef(uncollectNode, 0));
    rightNames.add(columnName);
    if (ordinality) {
        rightProjects.add(getCluster().getRexBuilder().makeInputRef(uncollectNode, 1));
        rightNames.add(String.format(zetaNode.getArrayOffsetColumn().getColumn().getTableName(), zetaNode.getArrayOffsetColumn().getColumn().getName()));
    }
    RelNode rightInput = LogicalProject.create(uncollectNode, ImmutableList.of(), rightProjects, rightNames);
    // Join condition should be a RexNode converted from join_expr.
    RexNode condition = getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getJoinExpr());
    JoinRelType joinRelType = zetaNode.getIsOuter() ? JoinRelType.LEFT : JoinRelType.INNER;
    return LogicalJoin.create(leftInput, rightInput, ImmutableList.of(), condition, ImmutableSet.of(), joinRelType, false, ImmutableList.of());
}