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

Example 31 with Union

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

the class AggregateUnionAggregateRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Aggregate topAggRel = call.rel(0);
    final Union union = call.rel(1);
    // If distincts haven't been removed yet, defer invoking this rule
    if (!union.all) {
        return;
    }
    final RelBuilder relBuilder = call.builder();
    final Aggregate bottomAggRel;
    if (call.rel(3) instanceof Aggregate) {
        // Aggregate is the second input
        bottomAggRel = call.rel(3);
        relBuilder.push(call.rel(2)).push(getInputWithSameRowType(bottomAggRel));
    } else if (call.rel(2) instanceof Aggregate) {
        // Aggregate is the first input
        bottomAggRel = call.rel(2);
        relBuilder.push(getInputWithSameRowType(bottomAggRel)).push(call.rel(3));
    } else {
        return;
    }
    // Only pull up aggregates if they are there just to remove distincts
    if (!topAggRel.getAggCallList().isEmpty() || !bottomAggRel.getAggCallList().isEmpty()) {
        return;
    }
    relBuilder.union(true);
    relBuilder.rename(union.getRowType().getFieldNames());
    relBuilder.aggregate(relBuilder.groupKey(topAggRel.getGroupSet()), topAggRel.getAggCallList());
    call.transformTo(relBuilder.build());
}

Example 32 with Union

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

the class UnionPullUpConstantsRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Union union = call.rel(0);
    final RexBuilder rexBuilder = union.getCluster().getRexBuilder();
    final RelMetadataQuery mq = call.getMetadataQuery();
    final RelOptPredicateList predicates = mq.getPulledUpPredicates(union);
    if (RelOptPredicateList.isEmpty(predicates)) {
        return;
    }
    final Map<Integer, RexNode> constants = new HashMap<>();
    for (Map.Entry<RexNode, RexNode> e : predicates.constantMap.entrySet()) {
        if (e.getKey() instanceof RexInputRef) {
            constants.put(((RexInputRef) e.getKey()).getIndex(), e.getValue());
        }
    }
    // None of the expressions are constant. Nothing to do.
    if (constants.isEmpty()) {
        return;
    }
    // Create expressions for Project operators before and after the Union
    List<RelDataTypeField> fields = union.getRowType().getFieldList();
    List<RexNode> topChildExprs = new ArrayList<>();
    List<String> topChildExprsFields = new ArrayList<>();
    List<RexNode> refs = new ArrayList<>();
    ImmutableBitSet.Builder refsIndexBuilder = ImmutableBitSet.builder();
    for (RelDataTypeField field : fields) {
        final RexNode constant = constants.get(field.getIndex());
        if (constant != null) {
            topChildExprs.add(constant);
            topChildExprsFields.add(field.getName());
        } else {
            final RexNode expr = rexBuilder.makeInputRef(union, field.getIndex());
            topChildExprs.add(expr);
            topChildExprsFields.add(field.getName());
            refs.add(expr);
            refsIndexBuilder.set(field.getIndex());
        }
    }
    ImmutableBitSet refsIndex = refsIndexBuilder.build();
    // Update top Project positions
    final Mappings.TargetMapping mapping = RelOptUtil.permutation(refs, union.getInput(0).getRowType()).inverse();
    topChildExprs = RexUtil.apply(mapping, topChildExprs);
    // Create new Project-Union-Project sequences
    final RelBuilder relBuilder = call.builder();
    for (RelNode input : union.getInputs()) {
        List<Pair<RexNode, String>> newChildExprs = new ArrayList<>();
        for (int j : refsIndex) {
            newChildExprs.add(Pair.of(rexBuilder.makeInputRef(input, j), input.getRowType().getFieldList().get(j).getName()));
        }
        if (newChildExprs.isEmpty()) {
            // At least a single item in project is required.
            newChildExprs.add(Pair.of(topChildExprs.get(0), topChildExprsFields.get(0)));
        }
        // Add the input with project on top
        relBuilder.push(input);
        relBuilder.project(Pair.left(newChildExprs), Pair.right(newChildExprs));
    }
    relBuilder.union(union.all, union.getInputs().size());
    // Create top Project fixing nullability of fields
    relBuilder.project(topChildExprs, topChildExprsFields);
    relBuilder.convert(union.getRowType(), false);
    call.transformTo(relBuilder.build());
}

Example 33 with Union

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

the class SortUnionTransposeRule method matches.

// ~ Methods ----------------------------------------------------------------
@Override
public boolean matches(RelOptRuleCall call) {
    final Sort sort = call.rel(0);
    final Union union = call.rel(1);
    // Sort.fetch is null.
    return union.all && sort.offset == null && (config.matchNullFetch() || sort.fetch != null);
}

Example 34 with Union

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

the class SortUnionTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Sort sort = call.rel(0);
    final Union union = call.rel(1);
    List<RelNode> inputs = new ArrayList<>();
    // Thus we use 'ret' as a flag to identify if we have finished pushing the
    // sort past a union.
    boolean ret = true;
    final RelMetadataQuery mq = call.getMetadataQuery();
    for (RelNode input : union.getInputs()) {
        if (!RelMdUtil.checkInputForCollationAndLimit(mq, input, sort.getCollation(), sort.offset, sort.fetch)) {
            ret = false;
            Sort branchSort = sort.copy(sort.getTraitSet(), input, sort.getCollation(), sort.offset, sort.fetch);
            inputs.add(branchSort);
        } else {
            inputs.add(input);
        }
    }
    // there is nothing to change
    if (ret) {
        return;
    }
    // create new union and sort
    Union unionCopy = (Union) union.copy(union.getTraitSet(), inputs, union.all);
    Sort result = sort.copy(sort.getTraitSet(), unionCopy, sort.getCollation(), sort.offset, sort.fetch);
    call.transformTo(result);
}

Example 35 with Union

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

the class EnumerableUnionRule method convert.

@Override
public RelNode convert(RelNode rel) {
    final Union union = (Union) rel;
    final EnumerableConvention out = EnumerableConvention.INSTANCE;
    final RelTraitSet traitSet = rel.getCluster().traitSet().replace(out);
    final List<RelNode> newInputs = Util.transform(union.getInputs(), n -> convert(n, traitSet));
    return new EnumerableUnion(rel.getCluster(), traitSet, newInputs, union.all);
}