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

Example 66 with RelFieldCollation

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

the class RelFieldTrimmer method trimFields.

/**
 * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
 * {@link org.apache.calcite.rel.core.Sort}.
 */
public TrimResult trimFields(Sort sort, ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
    final RelDataType rowType = sort.getRowType();
    final int fieldCount = rowType.getFieldCount();
    final RelCollation collation = sort.getCollation();
    final RelNode input = sort.getInput();
    // We use the fields used by the consumer, plus any fields used as sort
    // keys.
    final ImmutableBitSet.Builder inputFieldsUsed = fieldsUsed.rebuild();
    for (RelFieldCollation field : collation.getFieldCollations()) {
        inputFieldsUsed.set(field.getFieldIndex());
    }
    // Create input with trimmed columns.
    final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
    TrimResult trimResult = trimChild(sort, input, inputFieldsUsed.build(), inputExtraFields);
    RelNode newInput = trimResult.left;
    final Mapping inputMapping = trimResult.right;
    // there's nothing we can do.
    if (newInput == input && inputMapping.isIdentity() && fieldsUsed.cardinality() == fieldCount) {
        return result(sort, Mappings.createIdentity(fieldCount));
    }
    // leave the Sort unchanged in case we have dynamic limits
    if (sort.offset instanceof RexDynamicParam || sort.fetch instanceof RexDynamicParam) {
        return result(sort, inputMapping);
    }
    final RelBuilder relBuilder = REL_BUILDER.get();
    relBuilder.push(newInput);
    final int offset = sort.offset == null ? 0 : RexLiteral.intValue(sort.offset);
    final int fetch = sort.fetch == null ? -1 : RexLiteral.intValue(sort.fetch);
    final ImmutableList<RexNode> fields = relBuilder.fields(RexUtil.apply(inputMapping, collation));
    relBuilder.sortLimit(offset, fetch, fields);
    // needs them for its condition.
    return result(relBuilder.build(), inputMapping);
}

Example 67 with RelFieldCollation

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

the class HiveInsertExchange4JoinRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    JoinPredicateInfo joinPredInfo;
    if (call.rel(0) instanceof HiveMultiJoin) {
        HiveMultiJoin multiJoin = call.rel(0);
        joinPredInfo = multiJoin.getJoinPredicateInfo();
    } else if (call.rel(0) instanceof HiveJoin) {
        HiveJoin hiveJoin = call.rel(0);
        joinPredInfo = hiveJoin.getJoinPredicateInfo();
    } else if (call.rel(0) instanceof Join) {
        Join join = call.rel(0);
        try {
            joinPredInfo = HiveCalciteUtil.JoinPredicateInfo.constructJoinPredicateInfo(join);
        } catch (CalciteSemanticException e) {
            throw new RuntimeException(e);
        }
    } else {
        return;
    }
    for (RelNode child : call.rel(0).getInputs()) {
        if (((HepRelVertex) child).getCurrentRel() instanceof Exchange) {
            return;
        }
    }
    // Get key columns from inputs. Those are the columns on which we will distribute on.
    // It is also the columns we will sort on.
    List<RelNode> newInputs = new ArrayList<RelNode>();
    for (int i = 0; i < call.rel(0).getInputs().size(); i++) {
        List<Integer> joinKeyPositions = new ArrayList<Integer>();
        ImmutableList.Builder<RexNode> joinExprsBuilder = new ImmutableList.Builder<RexNode>();
        Set<String> keySet = Sets.newHashSet();
        ImmutableList.Builder<RelFieldCollation> collationListBuilder = new ImmutableList.Builder<RelFieldCollation>();
        for (int j = 0; j < joinPredInfo.getEquiJoinPredicateElements().size(); j++) {
            JoinLeafPredicateInfo joinLeafPredInfo = joinPredInfo.getEquiJoinPredicateElements().get(j);
            for (RexNode joinExprNode : joinLeafPredInfo.getJoinExprs(i)) {
                if (keySet.add(joinExprNode.toString())) {
                    joinExprsBuilder.add(joinExprNode);
                }
            }
            for (int pos : joinLeafPredInfo.getProjsJoinKeysInChildSchema(i)) {
                if (!joinKeyPositions.contains(pos)) {
                    joinKeyPositions.add(pos);
                    collationListBuilder.add(new RelFieldCollation(pos));
                }
            }
        }
        HiveSortExchange exchange = HiveSortExchange.create(call.rel(0).getInput(i), new HiveRelDistribution(RelDistribution.Type.HASH_DISTRIBUTED, joinKeyPositions), new HiveRelCollation(collationListBuilder.build()), joinExprsBuilder.build());
        newInputs.add(exchange);
    }
    RelNode newOp;
    if (call.rel(0) instanceof HiveMultiJoin) {
        HiveMultiJoin multiJoin = call.rel(0);
        newOp = multiJoin.copy(multiJoin.getTraitSet(), newInputs);
    } else if (call.rel(0) instanceof Join) {
        Join join = call.rel(0);
        newOp = join.copy(join.getTraitSet(), join.getCondition(), newInputs.get(0), newInputs.get(1), join.getJoinType(), join.isSemiJoinDone());
    } else {
        return;
    }
    call.getPlanner().onCopy(call.rel(0), newOp);
    call.transformTo(newOp);
}

Example 68 with RelFieldCollation

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

the class HiveRelFieldTrimmer method trimFields.

public TrimResult trimFields(HiveSortExchange exchange, ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
    final RelDataType rowType = exchange.getRowType();
    final int fieldCount = rowType.getFieldCount();
    final RelCollation collation = exchange.getCollation();
    final RelDistribution distribution = exchange.getDistribution();
    final RelNode input = exchange.getInput();
    // We use the fields used by the consumer, plus any fields used as exchange
    // keys.
    final ImmutableBitSet.Builder inputFieldsUsed = fieldsUsed.rebuild();
    for (RelFieldCollation field : collation.getFieldCollations()) {
        inputFieldsUsed.set(field.getFieldIndex());
    }
    for (int keyIndex : distribution.getKeys()) {
        inputFieldsUsed.set(keyIndex);
    }
    // Create input with trimmed columns.
    final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
    TrimResult trimResult = trimChild(exchange, input, inputFieldsUsed.build(), inputExtraFields);
    RelNode newInput = trimResult.left;
    final Mapping inputMapping = trimResult.right;
    // there's nothing we can do.
    if (newInput == input && inputMapping.isIdentity() && fieldsUsed.cardinality() == fieldCount) {
        return result(exchange, Mappings.createIdentity(fieldCount));
    }
    final RelBuilder relBuilder = REL_BUILDER.get();
    relBuilder.push(newInput);
    RelCollation newCollation = RexUtil.apply(inputMapping, collation);
    RelDistribution newDistribution = distribution.apply(inputMapping);
    relBuilder.sortExchange(newDistribution, newCollation);
    return result(relBuilder.build(), inputMapping);
}

Example 69 with RelFieldCollation

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

the class HiveRelColumnsAlignment method align.

public RelNode align(Join rel, List<RelFieldCollation> collations) {
    ImmutableList.Builder<RelFieldCollation> propagateCollationsLeft = ImmutableList.builder();
    ImmutableList.Builder<RelFieldCollation> propagateCollationsRight = ImmutableList.builder();
    final int nLeftColumns = rel.getLeft().getRowType().getFieldList().size();
    Map<Integer, RexNode> idxToConjuncts = new HashMap<>();
    Map<Integer, Integer> refToRef = new HashMap<>();
    // 1) We extract the conditions that can be useful
    List<RexNode> conjuncts = new ArrayList<>();
    List<RexNode> otherConjuncts = new ArrayList<>();
    for (RexNode conj : RelOptUtil.conjunctions(rel.getCondition())) {
        if (conj.getKind() != SqlKind.EQUALS) {
            otherConjuncts.add(conj);
            continue;
        }
        // TODO: Currently we only support EQUAL operator on two references.
        // We might extend the logic to support other (order-preserving)
        // UDFs here.
        RexCall equals = (RexCall) conj;
        if (!(equals.getOperands().get(0) instanceof RexInputRef) || !(equals.getOperands().get(1) instanceof RexInputRef)) {
            otherConjuncts.add(conj);
            continue;
        }
        RexInputRef ref0 = (RexInputRef) equals.getOperands().get(0);
        RexInputRef ref1 = (RexInputRef) equals.getOperands().get(1);
        if ((ref0.getIndex() < nLeftColumns && ref1.getIndex() >= nLeftColumns) || (ref1.getIndex() < nLeftColumns && ref0.getIndex() >= nLeftColumns)) {
            // We made sure the references are for different join inputs
            idxToConjuncts.put(ref0.getIndex(), equals);
            idxToConjuncts.put(ref1.getIndex(), equals);
            refToRef.put(ref0.getIndex(), ref1.getIndex());
            refToRef.put(ref1.getIndex(), ref0.getIndex());
        } else {
            otherConjuncts.add(conj);
        }
    }
    // that we will propagate to the inputs of the join
    for (RelFieldCollation c : collations) {
        RexNode equals = idxToConjuncts.get(c.getFieldIndex());
        if (equals != null) {
            conjuncts.add(equals);
            idxToConjuncts.remove(c.getFieldIndex());
            idxToConjuncts.remove(refToRef.get(c.getFieldIndex()));
            if (c.getFieldIndex() < nLeftColumns) {
                propagateCollationsLeft.add(c.copy(c.getFieldIndex()));
                propagateCollationsRight.add(c.copy(refToRef.get(c.getFieldIndex()) - nLeftColumns));
            } else {
                propagateCollationsLeft.add(c.copy(refToRef.get(c.getFieldIndex())));
                propagateCollationsRight.add(c.copy(c.getFieldIndex() - nLeftColumns));
            }
        }
    }
    final Set<RexNode> visited = new HashSet<>();
    for (Entry<Integer, RexNode> e : idxToConjuncts.entrySet()) {
        if (visited.add(e.getValue())) {
            // Not included in the input collations, but can be propagated as this Join
            // might enforce it
            conjuncts.add(e.getValue());
            if (e.getKey() < nLeftColumns) {
                propagateCollationsLeft.add(new RelFieldCollation(e.getKey()));
                propagateCollationsRight.add(new RelFieldCollation(refToRef.get(e.getKey()) - nLeftColumns));
            } else {
                propagateCollationsLeft.add(new RelFieldCollation(refToRef.get(e.getKey())));
                propagateCollationsRight.add(new RelFieldCollation(e.getKey() - nLeftColumns));
            }
        }
    }
    conjuncts.addAll(otherConjuncts);
    // 3) We propagate
    final RelNode newLeftInput = dispatchAlign(rel.getLeft(), propagateCollationsLeft.build());
    final RelNode newRightInput = dispatchAlign(rel.getRight(), propagateCollationsRight.build());
    // 4) We change the Join operator to reflect this info
    final RelNode newJoin = rel.copy(rel.getTraitSet(), RexUtil.composeConjunction(relBuilder.getRexBuilder(), conjuncts, false), newLeftInput, newRightInput, rel.getJoinType(), rel.isSemiJoinDone());
    return newJoin;
}

Example 70 with RelFieldCollation

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

the class HiveRelColumnsAlignment method align.

public RelNode align(Aggregate rel, List<RelFieldCollation> collations) {
    // 1) We extract the group by positions that are part of the collations and
    // sort them so they respect it
    LinkedHashSet<Integer> aggregateColumnsOrder = new LinkedHashSet<>();
    ImmutableList.Builder<RelFieldCollation> propagateCollations = ImmutableList.builder();
    if (rel.getGroupType() == Group.SIMPLE && !collations.isEmpty()) {
        for (RelFieldCollation c : collations) {
            if (c.getFieldIndex() < rel.getGroupCount()) {
                // Group column found
                if (aggregateColumnsOrder.add(c.getFieldIndex())) {
                    propagateCollations.add(c.copy(rel.getGroupSet().nth(c.getFieldIndex())));
                }
            }
        }
    }
    for (int i = 0; i < rel.getGroupCount(); i++) {
        if (!aggregateColumnsOrder.contains(i)) {
            // Not included in the input collations, but can be propagated as this Aggregate
            // will enforce it
            propagateCollations.add(new RelFieldCollation(rel.getGroupSet().nth(i)));
        }
    }
    // 2) We propagate
    final RelNode child = dispatchAlign(rel.getInput(), propagateCollations.build());
    // 3) We annotate the Aggregate operator with this info
    final HiveAggregate newAggregate = (HiveAggregate) rel.copy(rel.getTraitSet(), ImmutableList.of(child));
    newAggregate.setAggregateColumnsOrder(aggregateColumnsOrder);
    return newAggregate;
}