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

Example 11 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 ToLogicalConverter method visit.

@Override
public RelNode visit(RelNode relNode) {
    if (relNode instanceof Aggregate) {
        final Aggregate agg = (Aggregate) relNode;
        return relBuilder.push(visit(agg.getInput())).aggregate(relBuilder.groupKey(agg.getGroupSet(), agg.groupSets), agg.getAggCallList()).build();
    }
    if (relNode instanceof TableScan) {
        return visit((TableScan) relNode);
    }
    if (relNode instanceof Filter) {
        final Filter filter = (Filter) relNode;
        return relBuilder.push(visit(filter.getInput())).filter(filter.getCondition()).build();
    }
    if (relNode instanceof Project) {
        final Project project = (Project) relNode;
        return relBuilder.push(visit(project.getInput())).project(project.getProjects(), project.getRowType().getFieldNames()).build();
    }
    if (relNode instanceof Union) {
        final Union union = (Union) relNode;
        for (RelNode rel : union.getInputs()) {
            relBuilder.push(visit(rel));
        }
        return relBuilder.union(union.all, union.getInputs().size()).build();
    }
    if (relNode instanceof Intersect) {
        final Intersect intersect = (Intersect) relNode;
        for (RelNode rel : intersect.getInputs()) {
            relBuilder.push(visit(rel));
        }
        return relBuilder.intersect(intersect.all, intersect.getInputs().size()).build();
    }
    if (relNode instanceof Minus) {
        final Minus minus = (Minus) relNode;
        for (RelNode rel : minus.getInputs()) {
            relBuilder.push(visit(rel));
        }
        return relBuilder.minus(minus.all, minus.getInputs().size()).build();
    }
    if (relNode instanceof Join) {
        final Join join = (Join) relNode;
        return relBuilder.push(visit(join.getLeft())).push(visit(join.getRight())).join(join.getJoinType(), join.getCondition()).build();
    }
    if (relNode instanceof Correlate) {
        final Correlate corr = (Correlate) relNode;
        return relBuilder.push(visit(corr.getLeft())).push(visit(corr.getRight())).join(corr.getJoinType(), relBuilder.literal(true), corr.getVariablesSet()).build();
    }
    if (relNode instanceof Values) {
        final Values values = (Values) relNode;
        return relBuilder.values(values.tuples, values.getRowType()).build();
    }
    if (relNode instanceof Sort) {
        final Sort sort = (Sort) relNode;
        return LogicalSort.create(visit(sort.getInput()), sort.getCollation(), sort.offset, sort.fetch);
    }
    if (relNode instanceof Window) {
        final Window window = (Window) relNode;
        final RelNode input = visit(window.getInput());
        return LogicalWindow.create(input.getTraitSet(), input, window.constants, window.getRowType(), window.groups);
    }
    if (relNode instanceof Calc) {
        final Calc calc = (Calc) relNode;
        return LogicalCalc.create(visit(calc.getInput()), calc.getProgram());
    }
    if (relNode instanceof TableModify) {
        final TableModify tableModify = (TableModify) relNode;
        final RelNode input = visit(tableModify.getInput());
        return LogicalTableModify.create(tableModify.getTable(), tableModify.getCatalogReader(), input, tableModify.getOperation(), tableModify.getUpdateColumnList(), tableModify.getSourceExpressionList(), tableModify.isFlattened());
    }
    if (relNode instanceof EnumerableInterpreter || relNode instanceof JdbcToEnumerableConverter) {
        return visit(((SingleRel) relNode).getInput());
    }
    if (relNode instanceof EnumerableLimit) {
        final EnumerableLimit limit = (EnumerableLimit) relNode;
        RelNode logicalInput = visit(limit.getInput());
        RelCollation collation = RelCollations.of();
        if (logicalInput instanceof Sort) {
            collation = ((Sort) logicalInput).collation;
            logicalInput = ((Sort) logicalInput).getInput();
        }
        return LogicalSort.create(logicalInput, collation, limit.offset, limit.fetch);
    }
    if (relNode instanceof Uncollect) {
        final Uncollect uncollect = (Uncollect) relNode;
        final RelNode input = visit(uncollect.getInput());
        return Uncollect.create(input.getTraitSet(), input, uncollect.withOrdinality, Collections.emptyList());
    }
    throw new AssertionError("Need to implement logical converter for " + relNode.getClass().getName());
}

Example 12 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 RelMdExpressionLineage method getExpressionLineage.

/**
 * Expression lineage from {@link Union}.
 *
 * <p>For Union operator, we might be able to extract multiple origins for the
 * references in the given expression.
 */
@Nullable
public Set<RexNode> getExpressionLineage(Union rel, RelMetadataQuery mq, RexNode outputExpression) {
    final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
    // Extract input fields referenced by expression
    final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression);
    // Infer column origin expressions for given references
    final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
    final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
    for (RelNode input : rel.getInputs()) {
        // Gather table references
        final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
        final Set<RelTableRef> tableRefs = mq.getTableReferences(input);
        if (tableRefs == null) {
            // Bail out
            return null;
        }
        for (RelTableRef tableRef : tableRefs) {
            int shift = 0;
            Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(tableRef.getQualifiedName());
            if (lRefs != null) {
                shift = lRefs.size();
            }
            currentTablesMapping.put(tableRef, RelTableRef.of(tableRef.getTable(), shift + tableRef.getEntityNumber()));
        }
        // Map references
        for (int idx : inputFieldsUsed) {
            final RexInputRef inputRef = RexInputRef.of(idx, input.getRowType().getFieldList());
            final Set<RexNode> originalExprs = mq.getExpressionLineage(input, inputRef);
            if (originalExprs == null) {
                // Bail out
                return null;
            }
            // References might need to be updated
            final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
            final Set<RexNode> updatedExprs = originalExprs.stream().map(e -> RexUtil.swapTableReferences(rexBuilder, e, currentTablesMapping)).collect(Collectors.toSet());
            final Set<RexNode> set = mapping.get(ref);
            if (set != null) {
                set.addAll(updatedExprs);
            } else {
                mapping.put(ref, updatedExprs);
            }
        }
        // Add to existing qualified names
        for (RelTableRef newRef : currentTablesMapping.values()) {
            qualifiedNamesToRefs.put(newRef.getQualifiedName(), newRef);
        }
    }
    // Return result
    return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}

Example 13 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 JoinUnionTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Join join = call.rel(0);
    final Union unionRel;
    RelNode otherInput;
    boolean unionOnLeft;
    if (call.rel(1) instanceof Union) {
        unionRel = call.rel(1);
        otherInput = call.rel(2);
        unionOnLeft = true;
    } else {
        otherInput = call.rel(1);
        unionRel = call.rel(2);
        unionOnLeft = false;
    }
    if (!unionRel.all) {
        return;
    }
    if (!join.getVariablesSet().isEmpty()) {
        return;
    }
    // in one or both branches of the union)
    if (unionOnLeft) {
        if (join.getJoinType().generatesNullsOnLeft()) {
            return;
        }
    } else {
        if (join.getJoinType().generatesNullsOnRight() || !join.getJoinType().projectsRight()) {
            return;
        }
    }
    List<RelNode> newUnionInputs = new ArrayList<>();
    for (RelNode input : unionRel.getInputs()) {
        RelNode joinLeft;
        RelNode joinRight;
        if (unionOnLeft) {
            joinLeft = input;
            joinRight = otherInput;
        } else {
            joinLeft = otherInput;
            joinRight = input;
        }
        newUnionInputs.add(join.copy(join.getTraitSet(), join.getCondition(), joinLeft, joinRight, join.getJoinType(), join.isSemiJoinDone()));
    }
    final SetOp newUnionRel = unionRel.copy(unionRel.getTraitSet(), newUnionInputs, true);
    call.transformTo(newUnionRel);
}

Example 14 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 AggregateUnionTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    Aggregate aggRel = call.rel(0);
    Union union = call.rel(1);
    if (!union.all) {
        // which yields 25 (incorrect).
        return;
    }
    int groupCount = aggRel.getGroupSet().cardinality();
    List<AggregateCall> transformedAggCalls = transformAggCalls(aggRel.copy(aggRel.getTraitSet(), aggRel.getInput(), aggRel.getGroupSet(), null, aggRel.getAggCallList()), groupCount, aggRel.getAggCallList());
    if (transformedAggCalls == null) {
        // which we can't handle
        return;
    }
    boolean hasUniqueKeyInAllInputs = true;
    final RelMetadataQuery mq = call.getMetadataQuery();
    for (RelNode input : union.getInputs()) {
        boolean alreadyUnique = RelMdUtil.areColumnsDefinitelyUnique(mq, input, aggRel.getGroupSet());
        if (!alreadyUnique) {
            hasUniqueKeyInAllInputs = false;
            break;
        }
    }
    if (hasUniqueKeyInAllInputs) {
        // planners would succumb)
        return;
    }
    // create corresponding aggregates on top of each union child
    final RelBuilder relBuilder = call.builder();
    for (RelNode input : union.getInputs()) {
        relBuilder.push(input);
        relBuilder.aggregate(relBuilder.groupKey(aggRel.getGroupSet()), aggRel.getAggCallList());
    }
    // create a new union whose children are the aggregates created above
    relBuilder.union(true, union.getInputs().size());
    // Create the top aggregate. We must adjust group key indexes of the
    // original aggregate. E.g., if the original tree was:
    // 
    // Aggregate[groupSet=$1, ...]
    // Union[...]
    // 
    // Then the new tree should be:
    // Aggregate[groupSet=$0, ...]
    // Union[...]
    // Aggregate[groupSet=$1, ...]
    ImmutableBitSet groupSet = aggRel.getGroupSet();
    Mapping topGroupMapping = Mappings.create(MappingType.INVERSE_SURJECTION, union.getRowType().getFieldCount(), aggRel.getGroupCount());
    for (int i = 0; i < groupSet.cardinality(); i++) {
        topGroupMapping.set(groupSet.nth(i), i);
    }
    ImmutableBitSet topGroupSet = Mappings.apply(topGroupMapping, groupSet);
    ImmutableList<ImmutableBitSet> topGroupSets = Mappings.apply2(topGroupMapping, aggRel.getGroupSets());
    relBuilder.aggregate(relBuilder.groupKey(topGroupSet, topGroupSets), transformedAggCalls);
    call.transformTo(relBuilder.build());
}

Example 15 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 UnionEliminatorRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    Union union = call.rel(0);
    call.transformTo(union.getInputs().get(0));
}