Examples with RelCollation org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation used on opensource projects

Example 61 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project calcite by apache.

the class RelMetadataTest method checkCollation.

private void checkCollation(RelOptCluster cluster, RelOptTable empTable, RelOptTable deptTable) {
    final RexBuilder rexBuilder = cluster.getRexBuilder();
    final LogicalTableScan empScan = LogicalTableScan.create(cluster, empTable);
    List<RelCollation> collations = RelMdCollation.table(empScan.getTable());
    assertThat(collations.size(), equalTo(0));
    // ORDER BY field#0 ASC, field#1 ASC
    final RelCollation collation = RelCollations.of(new RelFieldCollation(0), new RelFieldCollation(1));
    collations = RelMdCollation.sort(collation);
    assertThat(collations.size(), equalTo(1));
    assertThat(collations.get(0).getFieldCollations().size(), equalTo(2));
    final Sort empSort = LogicalSort.create(empScan, collation, null, null);
    final List<RexNode> projects = ImmutableList.of(rexBuilder.makeInputRef(empSort, 1), rexBuilder.makeLiteral("foo"), rexBuilder.makeInputRef(empSort, 0), rexBuilder.makeCall(SqlStdOperatorTable.MINUS, rexBuilder.makeInputRef(empSort, 0), rexBuilder.makeInputRef(empSort, 3)));
    final RelMetadataQuery mq = RelMetadataQuery.instance();
    collations = RelMdCollation.project(mq, empSort, projects);
    assertThat(collations.size(), equalTo(1));
    assertThat(collations.get(0).getFieldCollations().size(), equalTo(2));
    assertThat(collations.get(0).getFieldCollations().get(0).getFieldIndex(), equalTo(2));
    assertThat(collations.get(0).getFieldCollations().get(1).getFieldIndex(), equalTo(0));
    final LogicalProject project = LogicalProject.create(empSort, projects, ImmutableList.of("a", "b", "c", "d"));
    final LogicalTableScan deptScan = LogicalTableScan.create(cluster, deptTable);
    final RelCollation deptCollation = RelCollations.of(new RelFieldCollation(0), new RelFieldCollation(1));
    final Sort deptSort = LogicalSort.create(deptScan, deptCollation, null, null);
    final ImmutableIntList leftKeys = ImmutableIntList.of(2);
    final ImmutableIntList rightKeys = ImmutableIntList.of(0);
    final EnumerableMergeJoin join;
    try {
        join = EnumerableMergeJoin.create(project, deptSort, rexBuilder.makeLiteral(true), leftKeys, rightKeys, JoinRelType.INNER);
    } catch (InvalidRelException e) {
        throw new RuntimeException(e);
    }
    collations = RelMdCollation.mergeJoin(mq, project, deptSort, leftKeys, rightKeys);
    assertThat(collations, equalTo(join.getTraitSet().getTraits(RelCollationTraitDef.INSTANCE)));
    // Values (empty)
    collations = RelMdCollation.values(mq, empTable.getRowType(), ImmutableList.<ImmutableList<RexLiteral>>of());
    assertThat(collations.toString(), equalTo("[[0, 1, 2, 3, 4, 5, 6, 7, 8], " + "[1, 2, 3, 4, 5, 6, 7, 8], " + "[2, 3, 4, 5, 6, 7, 8], " + "[3, 4, 5, 6, 7, 8], " + "[4, 5, 6, 7, 8], " + "[5, 6, 7, 8], " + "[6, 7, 8], " + "[7, 8], " + "[8]]"));
    final LogicalValues emptyValues = LogicalValues.createEmpty(cluster, empTable.getRowType());
    assertThat(mq.collations(emptyValues), equalTo(collations));
    // Values (non-empty)
    final RelDataType rowType = cluster.getTypeFactory().builder().add("a", SqlTypeName.INTEGER).add("b", SqlTypeName.INTEGER).add("c", SqlTypeName.INTEGER).add("d", SqlTypeName.INTEGER).build();
    final ImmutableList.Builder<ImmutableList<RexLiteral>> tuples = ImmutableList.builder();
    // sort keys are [a], [a, b], [a, b, c], [a, b, c, d], [a, c], [b], [b, a],
    // [b, d]
    // algorithm deduces [a, b, c, d], [b, d] which is a useful sub-set
    addRow(tuples, rexBuilder, 1, 1, 1, 1);
    addRow(tuples, rexBuilder, 1, 2, 0, 3);
    addRow(tuples, rexBuilder, 2, 3, 2, 2);
    addRow(tuples, rexBuilder, 3, 3, 1, 4);
    collations = RelMdCollation.values(mq, rowType, tuples.build());
    assertThat(collations.toString(), equalTo("[[0, 1, 2, 3], [1, 3]]"));
    final LogicalValues values = LogicalValues.create(cluster, rowType, tuples.build());
    assertThat(mq.collations(values), equalTo(collations));
}

Example 62 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project calcite by apache.

the class SortJoinTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Sort sort = call.rel(0);
    final Join join = call.rel(1);
    // We create a new sort operator on the corresponding input
    final RelNode newLeftInput;
    final RelNode newRightInput;
    final RelMetadataQuery mq = call.getMetadataQuery();
    if (join.getJoinType() == JoinRelType.LEFT) {
        // we bail out
        if (RelMdUtil.checkInputForCollationAndLimit(mq, join.getLeft(), sort.getCollation(), sort.offset, sort.fetch)) {
            return;
        }
        newLeftInput = sort.copy(sort.getTraitSet(), join.getLeft(), sort.getCollation(), sort.offset, sort.fetch);
        newRightInput = join.getRight();
    } else {
        final RelCollation rightCollation = RelCollationTraitDef.INSTANCE.canonize(RelCollations.shift(sort.getCollation(), -join.getLeft().getRowType().getFieldCount()));
        // we bail out
        if (RelMdUtil.checkInputForCollationAndLimit(mq, join.getRight(), rightCollation, sort.offset, sort.fetch)) {
            return;
        }
        newLeftInput = join.getLeft();
        newRightInput = sort.copy(sort.getTraitSet().replace(rightCollation), join.getRight(), rightCollation, sort.offset, sort.fetch);
    }
    // We copy the join and the top sort operator
    final RelNode joinCopy = join.copy(join.getTraitSet(), join.getCondition(), newLeftInput, newRightInput, join.getJoinType(), join.isSemiJoinDone());
    final RelNode sortCopy = sort.copy(sort.getTraitSet(), joinCopy, sort.getCollation(), sort.offset, sort.fetch);
    call.transformTo(sortCopy);
}

Example 63 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project calcite by apache.

the class RexUtil method apply.

/**
 * Applies a mapping to a collation list.
 *
 * @param mapping       Mapping
 * @param collationList Collation list
 * @return collation list with mapping applied to each field
 */
public static List<RelCollation> apply(Mappings.TargetMapping mapping, List<RelCollation> collationList) {
    final List<RelCollation> newCollationList = new ArrayList<>();
    for (RelCollation collation : collationList) {
        final List<RelFieldCollation> newFieldCollationList = new ArrayList<>();
        for (RelFieldCollation fieldCollation : collation.getFieldCollations()) {
            final RelFieldCollation newFieldCollation = apply(mapping, fieldCollation);
            if (newFieldCollation == null) {
                // if it's empty).
                break;
            }
            newFieldCollationList.add(newFieldCollation);
        }
        // and duplicate collations. Ignore these.
        if (!newFieldCollationList.isEmpty()) {
            final RelCollation newCollation = RelCollations.of(newFieldCollationList);
            if (!newCollationList.contains(newCollation)) {
                newCollationList.add(newCollation);
            }
        }
    }
    return newCollationList;
}

Example 64 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project flink by apache.

the class FlinkRelMdCollation method mergeJoin.

/**
 * Helper method to determine a {@link Join}'s collation assuming that it uses a merge-join
 * algorithm.
 *
 * <p>If the inputs are sorted on other keys <em>in addition to</em> the join key, the result
 * preserves those collations too.
 */
public static List<RelCollation> mergeJoin(RelMetadataQuery mq, RelNode left, RelNode right, ImmutableIntList leftKeys, ImmutableIntList rightKeys) {
    final com.google.common.collect.ImmutableList.Builder<RelCollation> builder = com.google.common.collect.ImmutableList.builder();
    final com.google.common.collect.ImmutableList<RelCollation> leftCollations = mq.collations(left);
    assert RelCollations.contains(leftCollations, leftKeys) : "cannot merge join: left input is not sorted on left keys";
    builder.addAll(leftCollations);
    final com.google.common.collect.ImmutableList<RelCollation> rightCollations = mq.collations(right);
    assert RelCollations.contains(rightCollations, rightKeys) : "cannot merge join: right input is not sorted on right keys";
    final int leftFieldCount = left.getRowType().getFieldCount();
    for (RelCollation collation : rightCollations) {
        builder.add(RelCollations.shift(collation, leftFieldCount));
    }
    return builder.build();
}

Example 65 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project flink by apache.

the class RelDecorrelator method decorrelateRel.

public Frame decorrelateRel(Sort rel) {
    // Sort itself should not reference corVars.
    assert !cm.mapRefRelToCorRef.containsKey(rel);
    // Sort only references field positions in collations field.
    // The collations field in the newRel now need to refer to the
    // new output positions in its input.
    // Its output does not change the input ordering, so there's no
    // need to call propagateExpr.
    final RelNode oldInput = rel.getInput();
    final Frame frame = getInvoke(oldInput, rel);
    if (frame == null) {
        // If input has not been rewritten, do not rewrite this rel.
        return null;
    }
    // BEGIN FLINK MODIFICATION
    // Reason: to de-correlate sort rel when its parent is not a correlate
    // Should be removed after CALCITE-4333 is fixed
    final RelNode newInput = frame.r;
    Mappings.TargetMapping mapping = Mappings.target(frame.oldToNewOutputs, oldInput.getRowType().getFieldCount(), newInput.getRowType().getFieldCount());
    RelCollation oldCollation = rel.getCollation();
    RelCollation newCollation = RexUtil.apply(mapping, oldCollation);
    final int offset = rel.offset == null ? -1 : RexLiteral.intValue(rel.offset);
    final int fetch = rel.fetch == null ? -1 : RexLiteral.intValue(rel.fetch);
    // END FLINK MODIFICATION
    final RelNode newSort = relBuilder.push(newInput).sortLimit(offset, fetch, relBuilder.fields(newCollation)).build();
    // Sort does not change input ordering
    return register(rel, newSort, frame.oldToNewOutputs, frame.corDefOutputs);
}