Examples with RelOptCluster org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptCluster used on opensource projects

Example 16 with RelOptCluster

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptCluster 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);
}

Example 17 with RelOptCluster

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptCluster in project calcite by apache.

the class AggregateStarTableRule method apply.

protected void apply(RelOptRuleCall call, Project postProject, final Aggregate aggregate, StarTable.StarTableScan scan) {
    final RelOptCluster cluster = scan.getCluster();
    final RelOptTable table = scan.getTable();
    final RelOptLattice lattice = call.getPlanner().getLattice(table);
    final List<Lattice.Measure> measures = lattice.lattice.toMeasures(aggregate.getAggCallList());
    final Pair<CalciteSchema.TableEntry, TileKey> pair = lattice.getAggregate(call.getPlanner(), aggregate.getGroupSet(), measures);
    if (pair == null) {
        return;
    }
    final RelBuilder relBuilder = call.builder();
    final CalciteSchema.TableEntry tableEntry = pair.left;
    final TileKey tileKey = pair.right;
    final RelMetadataQuery mq = call.getMetadataQuery();
    final double rowCount = aggregate.estimateRowCount(mq);
    final Table aggregateTable = tableEntry.getTable();
    final RelDataType aggregateTableRowType = aggregateTable.getRowType(cluster.getTypeFactory());
    final RelOptTable aggregateRelOptTable = RelOptTableImpl.create(table.getRelOptSchema(), aggregateTableRowType, tableEntry, rowCount);
    relBuilder.push(aggregateRelOptTable.toRel(RelOptUtil.getContext(cluster)));
    if (tileKey == null) {
        if (CalcitePrepareImpl.DEBUG) {
            System.out.println("Using materialization " + aggregateRelOptTable.getQualifiedName() + " (exact match)");
        }
    } else if (!tileKey.dimensions.equals(aggregate.getGroupSet())) {
        // Aggregate has finer granularity than we need. Roll up.
        if (CalcitePrepareImpl.DEBUG) {
            System.out.println("Using materialization " + aggregateRelOptTable.getQualifiedName() + ", rolling up " + tileKey.dimensions + " to " + aggregate.getGroupSet());
        }
        assert tileKey.dimensions.contains(aggregate.getGroupSet());
        final List<AggregateCall> aggCalls = Lists.newArrayList();
        ImmutableBitSet.Builder groupSet = ImmutableBitSet.builder();
        for (int key : aggregate.getGroupSet()) {
            groupSet.set(tileKey.dimensions.indexOf(key));
        }
        for (AggregateCall aggCall : aggregate.getAggCallList()) {
            final AggregateCall copy = rollUp(groupSet.cardinality(), relBuilder, aggCall, tileKey);
            if (copy == null) {
                return;
            }
            aggCalls.add(copy);
        }
        relBuilder.push(aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), false, groupSet.build(), null, aggCalls));
    } else if (!tileKey.measures.equals(measures)) {
        if (CalcitePrepareImpl.DEBUG) {
            System.out.println("Using materialization " + aggregateRelOptTable.getQualifiedName() + ", right granularity, but different measures " + aggregate.getAggCallList());
        }
        relBuilder.project(relBuilder.fields(new AbstractSourceMapping(tileKey.dimensions.cardinality() + tileKey.measures.size(), aggregate.getRowType().getFieldCount()) {

            public int getSourceOpt(int source) {
                assert aggregate.getIndicatorCount() == 0;
                if (source < aggregate.getGroupCount()) {
                    int in = tileKey.dimensions.nth(source);
                    return aggregate.getGroupSet().indexOf(in);
                }
                Lattice.Measure measure = measures.get(source - aggregate.getGroupCount());
                int i = tileKey.measures.indexOf(measure);
                assert i >= 0;
                return tileKey.dimensions.cardinality() + i;
            }
        }.inverse()));
    }
    if (postProject != null) {
        relBuilder.push(postProject.copy(postProject.getTraitSet(), ImmutableList.of(relBuilder.peek())));
    }
    call.transformTo(relBuilder.build());
}

Example 18 with RelOptCluster

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptCluster in project calcite by apache.

the class IntersectToDistinctRule method onMatch.

// ~ Methods ----------------------------------------------------------------
public void onMatch(RelOptRuleCall call) {
    final Intersect intersect = call.rel(0);
    if (intersect.all) {
        // nothing we can do
        return;
    }
    final RelOptCluster cluster = intersect.getCluster();
    final RexBuilder rexBuilder = cluster.getRexBuilder();
    final RelBuilder relBuilder = call.builder();
    // 1st level GB: create a GB (col0, col1, count() as c) for each branch
    for (RelNode input : intersect.getInputs()) {
        relBuilder.push(input);
        relBuilder.aggregate(relBuilder.groupKey(relBuilder.fields()), relBuilder.countStar(null));
    }
    // create a union above all the branches
    final int branchCount = intersect.getInputs().size();
    relBuilder.union(true, branchCount);
    final RelNode union = relBuilder.peek();
    // 2nd level GB: create a GB (col0, col1, count(c)) for each branch
    // the index of c is union.getRowType().getFieldList().size() - 1
    final int fieldCount = union.getRowType().getFieldCount();
    final ImmutableBitSet groupSet = ImmutableBitSet.range(fieldCount - 1);
    relBuilder.aggregate(relBuilder.groupKey(groupSet, null), relBuilder.countStar(null));
    // add a filter count(c) = #branches
    relBuilder.filter(relBuilder.equals(relBuilder.field(fieldCount - 1), rexBuilder.makeBigintLiteral(new BigDecimal(branchCount))));
    // Project all but the last field
    relBuilder.project(Util.skipLast(relBuilder.fields()));
    // the schema for intersect distinct is like this
    // R3 on all attributes + count(c) as cnt
    // finally add a project to project out the last column
    call.transformTo(relBuilder.build());
}

Example 19 with RelOptCluster

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptCluster in project calcite by apache.

the class JoinAssociateRule method onMatch.

// ~ Methods ----------------------------------------------------------------
public void onMatch(final RelOptRuleCall call) {
    final Join topJoin = call.rel(0);
    final Join bottomJoin = call.rel(1);
    final RelNode relA = bottomJoin.getLeft();
    final RelNode relB = bottomJoin.getRight();
    final RelSubset relC = call.rel(2);
    final RelOptCluster cluster = topJoin.getCluster();
    final RexBuilder rexBuilder = cluster.getRexBuilder();
    if (relC.getConvention() != relA.getConvention()) {
        // EnumerableConvention, we're only interested in enumerable subsets.
        return;
    }
    // topJoin
    // /     \
    // bottomJoin  C
    // /    \
    // A      B
    final int aCount = relA.getRowType().getFieldCount();
    final int bCount = relB.getRowType().getFieldCount();
    final int cCount = relC.getRowType().getFieldCount();
    final ImmutableBitSet aBitSet = ImmutableBitSet.range(0, aCount);
    final ImmutableBitSet bBitSet = ImmutableBitSet.range(aCount, aCount + bCount);
    if (!topJoin.getSystemFieldList().isEmpty()) {
        // FIXME Enable this rule for joins with system fields
        return;
    }
    // (Is this too strict?)
    if (topJoin.getJoinType() != JoinRelType.INNER || bottomJoin.getJoinType() != JoinRelType.INNER) {
        return;
    }
    // Goal is to transform to
    // 
    // newTopJoin
    // /     \
    // A   newBottomJoin
    // /    \
    // B      C
    // Split the condition of topJoin and bottomJoin into a conjunctions. A
    // condition can be pushed down if it does not use columns from A.
    final List<RexNode> top = Lists.newArrayList();
    final List<RexNode> bottom = Lists.newArrayList();
    JoinPushThroughJoinRule.split(topJoin.getCondition(), aBitSet, top, bottom);
    JoinPushThroughJoinRule.split(bottomJoin.getCondition(), aBitSet, top, bottom);
    // Mapping for moving conditions from topJoin or bottomJoin to
    // newBottomJoin.
    // target: | B | C      |
    // source: | A       | B | C      |
    final Mappings.TargetMapping bottomMapping = Mappings.createShiftMapping(aCount + bCount + cCount, 0, aCount, bCount, bCount, aCount + bCount, cCount);
    final List<RexNode> newBottomList = Lists.newArrayList();
    new RexPermuteInputsShuttle(bottomMapping, relB, relC).visitList(bottom, newBottomList);
    RexNode newBottomCondition = RexUtil.composeConjunction(rexBuilder, newBottomList, false);
    final Join newBottomJoin = bottomJoin.copy(bottomJoin.getTraitSet(), newBottomCondition, relB, relC, JoinRelType.INNER, false);
    // Condition for newTopJoin consists of pieces from bottomJoin and topJoin.
    // Field ordinals do not need to be changed.
    RexNode newTopCondition = RexUtil.composeConjunction(rexBuilder, top, false);
    final Join newTopJoin = topJoin.copy(topJoin.getTraitSet(), newTopCondition, relA, newBottomJoin, JoinRelType.INNER, false);
    call.transformTo(newTopJoin);
}

Example 20 with RelOptCluster

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptCluster in project calcite by apache.

the class RelDecorrelator method createValueGenerator.

/**
 * Create RelNode tree that produces a list of correlated variables.
 *
 * @param correlations         correlated variables to generate
 * @param valueGenFieldOffset  offset in the output that generated columns
 *                             will start
 * @param corDefOutputs        output positions for the correlated variables
 *                             generated
 * @return RelNode the root of the resultant RelNode tree
 */
private RelNode createValueGenerator(Iterable<CorRef> correlations, int valueGenFieldOffset, SortedMap<CorDef, Integer> corDefOutputs) {
    final Map<RelNode, List<Integer>> mapNewInputToOutputs = new HashMap<>();
    final Map<RelNode, Integer> mapNewInputToNewOffset = new HashMap<>();
    // Add to map all the referenced positions (relative to each input rel).
    for (CorRef corVar : correlations) {
        final int oldCorVarOffset = corVar.field;
        final RelNode oldInput = getCorRel(corVar);
        assert oldInput != null;
        final Frame frame = getFrame(oldInput, true);
        assert frame != null;
        final RelNode newInput = frame.r;
        final List<Integer> newLocalOutputs;
        if (!mapNewInputToOutputs.containsKey(newInput)) {
            newLocalOutputs = new ArrayList<>();
        } else {
            newLocalOutputs = mapNewInputToOutputs.get(newInput);
        }
        final int newCorVarOffset = frame.oldToNewOutputs.get(oldCorVarOffset);
        // Add all unique positions referenced.
        if (!newLocalOutputs.contains(newCorVarOffset)) {
            newLocalOutputs.add(newCorVarOffset);
        }
        mapNewInputToOutputs.put(newInput, newLocalOutputs);
    }
    int offset = 0;
    // Project only the correlated fields out of each input
    // and join the project together.
    // To make sure the plan does not change in terms of join order,
    // join these rels based on their occurrence in corVar list which
    // is sorted.
    final Set<RelNode> joinedInputs = new HashSet<>();
    RelNode r = null;
    for (CorRef corVar : correlations) {
        final RelNode oldInput = getCorRel(corVar);
        assert oldInput != null;
        final RelNode newInput = getFrame(oldInput, true).r;
        assert newInput != null;
        if (!joinedInputs.contains(newInput)) {
            RelNode project = RelOptUtil.createProject(newInput, mapNewInputToOutputs.get(newInput));
            RelNode distinct = relBuilder.push(project).distinct().build();
            RelOptCluster cluster = distinct.getCluster();
            joinedInputs.add(newInput);
            mapNewInputToNewOffset.put(newInput, offset);
            offset += distinct.getRowType().getFieldCount();
            if (r == null) {
                r = distinct;
            } else {
                r = LogicalJoin.create(r, distinct, cluster.getRexBuilder().makeLiteral(true), ImmutableSet.<CorrelationId>of(), JoinRelType.INNER);
            }
        }
    }
    // referencing correlated variables.
    for (CorRef corRef : correlations) {
        // The first input of a Correlate is always the rel defining
        // the correlated variables.
        final RelNode oldInput = getCorRel(corRef);
        assert oldInput != null;
        final Frame frame = getFrame(oldInput, true);
        final RelNode newInput = frame.r;
        assert newInput != null;
        final List<Integer> newLocalOutputs = mapNewInputToOutputs.get(newInput);
        final int newLocalOutput = frame.oldToNewOutputs.get(corRef.field);
        // newOutput is the index of the corVar in the referenced
        // position list plus the offset of referenced position list of
        // each newInput.
        final int newOutput = newLocalOutputs.indexOf(newLocalOutput) + mapNewInputToNewOffset.get(newInput) + valueGenFieldOffset;
        corDefOutputs.put(corRef.def(), newOutput);
    }
    return r;
}