Examples with RexBuilder org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder used on opensource projects

Example 16 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project calcite by apache.

the class RelMdExpressionLineage method getExpressionLineage.

/**
 * Expression lineage from {@link Join}.
 *
 * <p>We only extract the lineage for INNER joins.
 */
public Set<RexNode> getExpressionLineage(Join rel, RelMetadataQuery mq, RexNode outputExpression) {
    if (rel.getJoinType() != JoinRelType.INNER) {
        // We cannot map origin of this expression.
        return null;
    }
    final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
    final RelNode leftInput = rel.getLeft();
    final RelNode rightInput = rel.getRight();
    final int nLeftColumns = leftInput.getRowType().getFieldList().size();
    // Infer column origin expressions for given references
    final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
    final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
    final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
    for (int idx = 0; idx < rel.getRowType().getFieldList().size(); idx++) {
        if (idx < nLeftColumns) {
            final RexInputRef inputRef = RexInputRef.of(idx, leftInput.getRowType().getFieldList());
            final Set<RexNode> originalExprs = mq.getExpressionLineage(leftInput, inputRef);
            if (originalExprs == null) {
                // Bail out
                return null;
            }
            // Gather table references, left input references remain unchanged
            final Set<RelTableRef> tableRefs = RexUtil.gatherTableReferences(Lists.newArrayList(originalExprs));
            for (RelTableRef leftRef : tableRefs) {
                qualifiedNamesToRefs.put(leftRef.getQualifiedName(), leftRef);
            }
            mapping.put(RexInputRef.of(idx, rel.getRowType().getFieldList()), originalExprs);
        } else {
            // Right input.
            final RexInputRef inputRef = RexInputRef.of(idx - nLeftColumns, rightInput.getRowType().getFieldList());
            final Set<RexNode> originalExprs = mq.getExpressionLineage(rightInput, inputRef);
            if (originalExprs == null) {
                // Bail out
                return null;
            }
            // Gather table references, right input references might need to be
            // updated if there are table names clashes with left input
            final Set<RelTableRef> tableRefs = RexUtil.gatherTableReferences(Lists.newArrayList(originalExprs));
            for (RelTableRef rightRef : tableRefs) {
                int shift = 0;
                Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(rightRef.getQualifiedName());
                if (lRefs != null) {
                    shift = lRefs.size();
                }
                currentTablesMapping.put(rightRef, RelTableRef.of(rightRef.getTable(), shift + rightRef.getEntityNumber()));
            }
            final Set<RexNode> updatedExprs = Sets.newHashSet(Iterables.transform(originalExprs, new Function<RexNode, RexNode>() {

                @Override
                public RexNode apply(RexNode e) {
                    return RexUtil.swapTableReferences(rexBuilder, e, currentTablesMapping);
                }
            }));
            mapping.put(RexInputRef.of(idx, rel.getRowType().getFieldList()), updatedExprs);
        }
    }
    // Return result
    return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}

Example 17 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder 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.
 */
public Set<RexNode> getExpressionLineage(Union rel, RelMetadataQuery mq, RexNode outputExpression) {
    final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
    // 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()) {
        final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
        for (int idx = 0; idx < input.getRowType().getFieldList().size(); idx++) {
            final RexInputRef inputRef = RexInputRef.of(idx, input.getRowType().getFieldList());
            final Set<RexNode> originalExprs = mq.getExpressionLineage(input, inputRef);
            if (originalExprs == null) {
                // Bail out
                return null;
            }
            final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
            // Gather table references, references might need to be
            // updated
            final Set<RelTableRef> tableRefs = RexUtil.gatherTableReferences(Lists.newArrayList(originalExprs));
            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()));
            }
            final Set<RexNode> updatedExprs = Sets.newHashSet(Iterables.transform(originalExprs, new Function<RexNode, RexNode>() {

                @Override
                public RexNode apply(RexNode e) {
                    return RexUtil.swapTableReferences(rexBuilder, e, currentTablesMapping);
                }
            }));
            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 18 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project calcite by apache.

the class RelMdPredicates method getPredicates.

/**
 * Infers predicates for a Union.
 */
public RelOptPredicateList getPredicates(Union union, RelMetadataQuery mq) {
    RexBuilder rexBuilder = union.getCluster().getRexBuilder();
    Map<String, RexNode> finalPreds = new HashMap<>();
    List<RexNode> finalResidualPreds = new ArrayList<>();
    for (int i = 0; i < union.getInputs().size(); i++) {
        RelNode input = union.getInputs().get(i);
        RelOptPredicateList info = mq.getPulledUpPredicates(input);
        if (info.pulledUpPredicates.isEmpty()) {
            return RelOptPredicateList.EMPTY;
        }
        Map<String, RexNode> preds = new HashMap<>();
        List<RexNode> residualPreds = new ArrayList<>();
        for (RexNode pred : info.pulledUpPredicates) {
            final String predDigest = pred.toString();
            if (i == 0) {
                preds.put(predDigest, pred);
                continue;
            }
            if (finalPreds.containsKey(predDigest)) {
                preds.put(predDigest, pred);
            } else {
                residualPreds.add(pred);
            }
        }
        // Add new residual preds
        finalResidualPreds.add(RexUtil.composeConjunction(rexBuilder, residualPreds, false));
        // Add those that are not part of the final set to residual
        for (Entry<String, RexNode> e : finalPreds.entrySet()) {
            if (!preds.containsKey(e.getKey())) {
                // This node was in previous union inputs, but it is not in this one
                for (int j = 0; j < i; j++) {
                    finalResidualPreds.set(j, RexUtil.composeConjunction(rexBuilder, Lists.newArrayList(finalResidualPreds.get(j), e.getValue()), false));
                }
            }
        }
        // Final preds
        finalPreds = preds;
    }
    List<RexNode> preds = new ArrayList<>(finalPreds.values());
    final RelOptCluster cluster = union.getCluster();
    final RexExecutor executor = Util.first(cluster.getPlanner().getExecutor(), RexUtil.EXECUTOR);
    final RelOptPredicateList predicates = RelOptPredicateList.EMPTY;
    final RexSimplify simplify = new RexSimplify(rexBuilder, predicates, true, executor);
    RexNode disjPred = simplify.simplifyOrs(finalResidualPreds);
    if (!disjPred.isAlwaysTrue()) {
        preds.add(disjPred);
    }
    return RelOptPredicateList.of(rexBuilder, preds);
}

Example 19 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project calcite by apache.

the class RelMdPredicates method getPredicates.

/**
 * Infers predicates for an Aggregate.
 *
 * <p>Pulls up predicates that only contains references to columns in the
 * GroupSet. For e.g.
 *
 * <blockquote><pre>
 * inputPullUpExprs : { a &gt; 7, b + c &lt; 10, a + e = 9}
 * groupSet         : { a, b}
 * pulledUpExprs    : { a &gt; 7}
 * </pre></blockquote>
 */
public RelOptPredicateList getPredicates(Aggregate agg, RelMetadataQuery mq) {
    final RelNode input = agg.getInput();
    final RexBuilder rexBuilder = agg.getCluster().getRexBuilder();
    final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
    final List<RexNode> aggPullUpPredicates = new ArrayList<>();
    ImmutableBitSet groupKeys = agg.getGroupSet();
    if (groupKeys.isEmpty()) {
        // no rows!) but not on the output (there is one row).
        return RelOptPredicateList.EMPTY;
    }
    Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION, input.getRowType().getFieldCount(), agg.getRowType().getFieldCount());
    int i = 0;
    for (int j : groupKeys) {
        m.set(j, i++);
    }
    for (RexNode r : inputInfo.pulledUpPredicates) {
        ImmutableBitSet rCols = RelOptUtil.InputFinder.bits(r);
        if (groupKeys.contains(rCols)) {
            r = r.accept(new RexPermuteInputsShuttle(m, input));
            aggPullUpPredicates.add(r);
        }
    }
    return RelOptPredicateList.of(rexBuilder, aggPullUpPredicates);
}

Example 20 with RexBuilder

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder in project calcite by apache.

the class RelMdPredicates method getPredicates.

/**
 * Infers predicates for a {@link org.apache.calcite.rel.core.Join} (including
 * {@link org.apache.calcite.rel.core.SemiJoin}).
 */
public RelOptPredicateList getPredicates(Join join, RelMetadataQuery mq) {
    RelOptCluster cluster = join.getCluster();
    RexBuilder rexBuilder = cluster.getRexBuilder();
    final RexExecutor executor = Util.first(cluster.getPlanner().getExecutor(), RexUtil.EXECUTOR);
    final RelNode left = join.getInput(0);
    final RelNode right = join.getInput(1);
    final RelOptPredicateList leftInfo = mq.getPulledUpPredicates(left);
    final RelOptPredicateList rightInfo = mq.getPulledUpPredicates(right);
    final RexSimplify simplifier = new RexSimplify(rexBuilder, RelOptPredicateList.EMPTY, true, executor);
    JoinConditionBasedPredicateInference joinInference = new JoinConditionBasedPredicateInference(join, RexUtil.composeConjunction(rexBuilder, leftInfo.pulledUpPredicates, false), RexUtil.composeConjunction(rexBuilder, rightInfo.pulledUpPredicates, false), simplifier);
    return joinInference.inferPredicates(false);
}