Examples with ImmutableBitSet org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet used on opensource projects

Example 51 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class RelMetadataTest method assertUniqueConsistent.

/**
 * Asserts that {@link RelMetadataQuery#getUniqueKeys(RelNode)}
 * and {@link RelMetadataQuery#areColumnsUnique(RelNode, ImmutableBitSet)}
 * return consistent results.
 */
private void assertUniqueConsistent(RelNode rel) {
    final RelMetadataQuery mq = RelMetadataQuery.instance();
    final Set<ImmutableBitSet> uniqueKeys = mq.getUniqueKeys(rel);
    final ImmutableBitSet allCols = ImmutableBitSet.range(0, rel.getRowType().getFieldCount());
    for (ImmutableBitSet key : allCols.powerSet()) {
        Boolean result2 = mq.areColumnsUnique(rel, key);
        assertTrue(result2 == null || result2 == isUnique(uniqueKeys, key));
    }
}

Example 52 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class RelMetadataTest method testUnion.

@Test
public void testUnion() {
    RelNode rel = convertSql("select deptno from emp\n" + "union\n" + "select deptno from dept");
    final RelMetadataQuery mq = RelMetadataQuery.instance();
    final Set<ImmutableBitSet> result = mq.getUniqueKeys(rel);
    assertThat(result, CoreMatchers.<Set<ImmutableBitSet>>equalTo(ImmutableSet.of(ImmutableBitSet.of(0))));
    assertUniqueConsistent(rel);
}

Example 53 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class ProjectWindowTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final LogicalProject project = call.rel(0);
    final LogicalWindow window = call.rel(1);
    final RelOptCluster cluster = window.getCluster();
    final List<RelDataTypeField> rowTypeWindowInput = window.getInput().getRowType().getFieldList();
    final int windowInputColumn = rowTypeWindowInput.size();
    // Record the window input columns which are actually referred
    // either in the LogicalProject above LogicalWindow or LogicalWindow itself
    // (Note that the constants used in LogicalWindow are not considered here)
    final ImmutableBitSet beReferred = findReference(project, window);
    // it is impossible to trim anyone of them out
    if (beReferred.cardinality() == windowInputColumn) {
        return;
    }
    // Put a DrillProjectRel below LogicalWindow
    final List<RexNode> exps = new ArrayList<>();
    final RelDataTypeFactory.Builder builder = cluster.getTypeFactory().builder();
    // Keep only the fields which are referred
    for (int index : BitSets.toIter(beReferred)) {
        final RelDataTypeField relDataTypeField = rowTypeWindowInput.get(index);
        exps.add(new RexInputRef(index, relDataTypeField.getType()));
        builder.add(relDataTypeField);
    }
    final LogicalProject projectBelowWindow = new LogicalProject(cluster, window.getTraitSet(), window.getInput(), exps, builder.build());
    // Create a new LogicalWindow with necessary inputs only
    final List<Window.Group> groups = new ArrayList<>();
    // As the un-referred columns are trimmed by the LogicalProject,
    // the indices specified in LogicalWindow would need to be adjusted
    final RexShuttle indexAdjustment = new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            final int newIndex = getAdjustedIndex(inputRef.getIndex(), beReferred, windowInputColumn);
            return new RexInputRef(newIndex, inputRef.getType());
        }

        @Override
        public RexNode visitCall(final RexCall call) {
            if (call instanceof Window.RexWinAggCall) {
                boolean[] update = { false };
                final List<RexNode> clonedOperands = visitList(call.operands, update);
                if (update[0]) {
                    return new Window.RexWinAggCall((SqlAggFunction) call.getOperator(), call.getType(), clonedOperands, ((Window.RexWinAggCall) call).ordinal, ((Window.RexWinAggCall) call).distinct);
                } else {
                    return call;
                }
            } else {
                return super.visitCall(call);
            }
        }
    };
    int aggCallIndex = windowInputColumn;
    final RelDataTypeFactory.Builder outputBuilder = cluster.getTypeFactory().builder();
    outputBuilder.addAll(projectBelowWindow.getRowType().getFieldList());
    for (Window.Group group : window.groups) {
        final ImmutableBitSet.Builder keys = ImmutableBitSet.builder();
        final List<RelFieldCollation> orderKeys = new ArrayList<>();
        final List<Window.RexWinAggCall> aggCalls = new ArrayList<>();
        // Adjust keys
        for (int index : group.keys) {
            keys.set(getAdjustedIndex(index, beReferred, windowInputColumn));
        }
        // Adjust orderKeys
        for (RelFieldCollation relFieldCollation : group.orderKeys.getFieldCollations()) {
            final int index = relFieldCollation.getFieldIndex();
            orderKeys.add(relFieldCollation.copy(getAdjustedIndex(index, beReferred, windowInputColumn)));
        }
        // Adjust Window Functions
        for (Window.RexWinAggCall rexWinAggCall : group.aggCalls) {
            aggCalls.add((Window.RexWinAggCall) rexWinAggCall.accept(indexAdjustment));
            final RelDataTypeField relDataTypeField = window.getRowType().getFieldList().get(aggCallIndex);
            outputBuilder.add(relDataTypeField);
            ++aggCallIndex;
        }
        groups.add(new Window.Group(keys.build(), group.isRows, group.lowerBound, group.upperBound, RelCollations.of(orderKeys), aggCalls));
    }
    final LogicalWindow newLogicalWindow = LogicalWindow.create(window.getTraitSet(), projectBelowWindow, window.constants, outputBuilder.build(), groups);
    // Modify the top LogicalProject
    final List<RexNode> topProjExps = new ArrayList<>();
    for (RexNode rexNode : project.getChildExps()) {
        topProjExps.add(rexNode.accept(indexAdjustment));
    }
    final LogicalProject newTopProj = project.copy(newLogicalWindow.getTraitSet(), newLogicalWindow, topProjExps, project.getRowType());
    if (ProjectRemoveRule.isTrivial(newTopProj)) {
        call.transformTo(newLogicalWindow);
    } else {
        call.transformTo(newTopProj);
    }
}

Example 54 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class ReduceExpressionsRule method pushPredicateIntoCase.

/**
 * Pushes predicates into a CASE.
 *
 * <p>We have a loose definition of 'predicate': any boolean expression will
 * do, except CASE. For example '(CASE ...) = 5' or '(CASE ...) IS NULL'.
 */
public static RexCall pushPredicateIntoCase(RexCall call) {
    if (call.getType().getSqlTypeName() != SqlTypeName.BOOLEAN) {
        return call;
    }
    switch(call.getKind()) {
        case CASE:
        case AND:
        case OR:
            // don't push CASE into CASE!
            return call;
        case EQUALS:
            {
                // checks that the EQUALS operands may be splitted and
                // doesn't push EQUALS into CASE
                List<RexNode> equalsOperands = call.getOperands();
                ImmutableBitSet left = RelOptUtil.InputFinder.bits(equalsOperands.get(0));
                ImmutableBitSet right = RelOptUtil.InputFinder.bits(equalsOperands.get(1));
                if (!left.isEmpty() && !right.isEmpty() && left.intersect(right).isEmpty()) {
                    return call;
                }
            }
    }
    int caseOrdinal = -1;
    final List<RexNode> operands = call.getOperands();
    for (int i = 0; i < operands.size(); i++) {
        RexNode operand = operands.get(i);
        switch(operand.getKind()) {
            case CASE:
                caseOrdinal = i;
        }
    }
    if (caseOrdinal < 0) {
        return call;
    }
    // Convert
    // f(CASE WHEN p1 THEN v1 ... END, arg)
    // to
    // CASE WHEN p1 THEN f(v1, arg) ... END
    final RexCall case_ = (RexCall) operands.get(caseOrdinal);
    final List<RexNode> nodes = new ArrayList<>();
    for (int i = 0; i < case_.getOperands().size(); i++) {
        RexNode node = case_.getOperands().get(i);
        if (!RexUtil.isCasePredicate(case_, i)) {
            node = substitute(call, caseOrdinal, node);
        }
        nodes.add(node);
    }
    return case_.clone(call.getType(), nodes);
}

Example 55 with ImmutableBitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class SemiJoinRule method perform.

protected void perform(RelOptRuleCall call, Project project, Join join, RelNode left, Aggregate aggregate) {
    final RelOptCluster cluster = join.getCluster();
    final RexBuilder rexBuilder = cluster.getRexBuilder();
    if (project != null) {
        final ImmutableBitSet bits = RelOptUtil.InputFinder.bits(project.getProjects(), null);
        final ImmutableBitSet rightBits = ImmutableBitSet.range(left.getRowType().getFieldCount(), join.getRowType().getFieldCount());
        if (bits.intersects(rightBits)) {
            return;
        }
    }
    final JoinInfo joinInfo = join.analyzeCondition();
    if (!joinInfo.rightSet().equals(ImmutableBitSet.range(aggregate.getGroupCount()))) {
        // By the way, neither a super-set nor a sub-set would work.
        return;
    }
    if (!joinInfo.isEqui()) {
        return;
    }
    final RelBuilder relBuilder = call.builder();
    relBuilder.push(left);
    switch(join.getJoinType()) {
        case INNER:
            final List<Integer> newRightKeyBuilder = Lists.newArrayList();
            final List<Integer> aggregateKeys = aggregate.getGroupSet().asList();
            for (int key : joinInfo.rightKeys) {
                newRightKeyBuilder.add(aggregateKeys.get(key));
            }
            final ImmutableIntList newRightKeys = ImmutableIntList.copyOf(newRightKeyBuilder);
            relBuilder.push(aggregate.getInput());
            final RexNode newCondition = RelOptUtil.createEquiJoinCondition(relBuilder.peek(2, 0), joinInfo.leftKeys, relBuilder.peek(2, 1), newRightKeys, rexBuilder);
            relBuilder.semiJoin(newCondition);
            break;
        case LEFT:
            // we can eliminate the semi-join.
            break;
        default:
            throw new AssertionError(join.getJoinType());
    }
    if (project != null) {
        relBuilder.project(project.getProjects(), project.getRowType().getFieldNames());
    }
    call.transformTo(relBuilder.build());
}