Search in sources :

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());
}
Also used : RelNode(org.apache.calcite.rel.RelNode) RexInputRef(org.apache.calcite.rex.RexInputRef) RelDataType(org.apache.calcite.rel.type.RelDataType)

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);
}
Also used : HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) TreeMap(java.util.TreeMap) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RelNode(org.apache.calcite.rel.RelNode) RexInputRef(org.apache.calcite.rex.RexInputRef) Map(java.util.Map) ImmutableSortedMap(com.google.common.collect.ImmutableSortedMap) TreeMap(java.util.TreeMap) ImmutableMap(com.google.common.collect.ImmutableMap) NavigableMap(java.util.NavigableMap) SortedMap(java.util.SortedMap) HashMap(java.util.HashMap) RexNode(org.apache.calcite.rex.RexNode)

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);
}
Also used : ArrayList(java.util.ArrayList) RelMdUtil(org.apache.calcite.rel.metadata.RelMdUtil) RexUtil(org.apache.calcite.rex.RexUtil) RelOptUtil(org.apache.calcite.plan.RelOptUtil) ReflectUtil(org.apache.calcite.util.ReflectUtil) Util(org.apache.calcite.util.Util) TreeMap(java.util.TreeMap) RelNode(org.apache.calcite.rel.RelNode) Filter(org.apache.calcite.rel.core.Filter) LogicalFilter(org.apache.calcite.rel.logical.LogicalFilter) RexInputRef(org.apache.calcite.rex.RexInputRef) RexNode(org.apache.calcite.rex.RexNode)

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);
}
Also used : Arrays(java.util.Arrays) FlinkMatchers(org.apache.flink.core.testutils.FlinkMatchers) IntType(org.apache.flink.table.types.logical.IntType) Random(java.util.Random) FlinkTypeFactory(org.apache.flink.table.planner.calcite.FlinkTypeFactory) RowType(org.apache.flink.table.types.logical.RowType) ArrayList(java.util.ArrayList) HashFunction(org.apache.flink.table.runtime.generated.HashFunction) TableConfigOptions(org.apache.flink.table.api.config.TableConfigOptions) BinaryRowWriter(org.apache.flink.table.data.writer.BinaryRowWriter) GenericRowData(org.apache.flink.table.data.GenericRowData) RexNode(org.apache.calcite.rex.RexNode) OutputStream(java.io.OutputStream) PrintStream(java.io.PrintStream) RelDataType(org.apache.calcite.rel.type.RelDataType) TableConfig(org.apache.flink.table.api.TableConfig) RecordComparator(org.apache.flink.table.runtime.generated.RecordComparator) RexBuilder(org.apache.calcite.rex.RexBuilder) Test(org.junit.Test) IOException(java.io.IOException) BinaryRowData(org.apache.flink.table.data.binary.BinaryRowData) ComparatorCodeGenerator(org.apache.flink.table.planner.codegen.sort.ComparatorCodeGenerator) RexInputRef(org.apache.calcite.rex.RexInputRef) Consumer(java.util.function.Consumer) JoinUtil(org.apache.flink.table.planner.plan.utils.JoinUtil) JoinCondition(org.apache.flink.table.runtime.generated.JoinCondition) List(java.util.List) MatcherAssert(org.hamcrest.MatcherAssert) LogicalType(org.apache.flink.table.types.logical.LogicalType) SqlStdOperatorTable(org.apache.calcite.sql.fun.SqlStdOperatorTable) Assert(org.junit.Assert) Collections(java.util.Collections) SortSpec(org.apache.flink.table.planner.plan.nodes.exec.spec.SortSpec) Projection(org.apache.flink.table.runtime.generated.Projection) RowType(org.apache.flink.table.types.logical.RowType) RelDataType(org.apache.calcite.rel.type.RelDataType) IntType(org.apache.flink.table.types.logical.IntType) JoinCondition(org.apache.flink.table.runtime.generated.JoinCondition) FlinkTypeFactory(org.apache.flink.table.planner.calcite.FlinkTypeFactory) Random(java.util.Random) RexBuilder(org.apache.calcite.rex.RexBuilder) RexInputRef(org.apache.calcite.rex.RexInputRef) GenericRowData(org.apache.flink.table.data.GenericRowData) TableConfig(org.apache.flink.table.api.TableConfig) RexNode(org.apache.calcite.rex.RexNode) Test(org.junit.Test)

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());
}
Also used : ArrayList(java.util.ArrayList) ResolvedColumnRef(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedColumnRef) JoinRelType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.JoinRelType) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) ZetaSqlUnnest(org.apache.beam.sdk.extensions.sql.zetasql.unnest.ZetaSqlUnnest) RexInputRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef) CorrelationId(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.CorrelationId) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Aggregations

RexInputRef (org.apache.calcite.rex.RexInputRef)241 RexNode (org.apache.calcite.rex.RexNode)200 ArrayList (java.util.ArrayList)105 RelNode (org.apache.calcite.rel.RelNode)85 RelDataTypeField (org.apache.calcite.rel.type.RelDataTypeField)80 RexCall (org.apache.calcite.rex.RexCall)67 RelDataType (org.apache.calcite.rel.type.RelDataType)63 RexBuilder (org.apache.calcite.rex.RexBuilder)54 ImmutableBitSet (org.apache.calcite.util.ImmutableBitSet)52 HashMap (java.util.HashMap)47 AggregateCall (org.apache.calcite.rel.core.AggregateCall)36 List (java.util.List)35 HashSet (java.util.HashSet)32 Pair (org.apache.calcite.util.Pair)32 RexLiteral (org.apache.calcite.rex.RexLiteral)29 Map (java.util.Map)24 RelOptUtil (org.apache.calcite.plan.RelOptUtil)24 Set (java.util.Set)20 ImmutableList (com.google.common.collect.ImmutableList)19 LinkedHashMap (java.util.LinkedHashMap)19