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

Example 71 with ImmutableBitSet

use of org.apache.calcite.util.ImmutableBitSet in project ignite-3 by apache.

the class HashAggregateSingleGroupExecutionTest method mapReduceMin.

@Test
public void mapReduceMin() {
    ExecutionContext<Object[]> ctx = executionContext();
    IgniteTypeFactory tf = ctx.getTypeFactory();
    RelDataType rowType = TypeUtils.createRowType(tf, String.class, int.class);
    ScanNode<Object[]> scan = new ScanNode<>(ctx, rowType, Arrays.asList(row("Igor", 200), row("Roman", 300), row("Ivan", 1400), row("Alexey", 1000)));
    AggregateCall call = AggregateCall.create(SqlStdOperatorTable.MIN, false, false, false, ImmutableIntList.of(1), -1, RelCollations.EMPTY, tf.createJavaType(int.class), null);
    List<ImmutableBitSet> grpSets = List.of(ImmutableBitSet.of());
    RelDataType mapType = IgniteMapHashAggregate.rowType(tf, true);
    HashAggregateNode<Object[]> map = new HashAggregateNode<>(ctx, mapType, MAP, grpSets, accFactory(ctx, call, MAP, rowType), rowFactory());
    map.register(scan);
    RelDataType reduceType = TypeUtils.createRowType(tf, int.class);
    HashAggregateNode<Object[]> reduce = new HashAggregateNode<>(ctx, reduceType, REDUCE, grpSets, accFactory(ctx, call, REDUCE, null), rowFactory());
    reduce.register(map);
    RootNode<Object[]> root = new RootNode<>(ctx, reduceType);
    root.register(reduce);
    assertTrue(root.hasNext());
    assertEquals(200, root.next()[0]);
    assertFalse(root.hasNext());
}

Example 72 with ImmutableBitSet

use of org.apache.calcite.util.ImmutableBitSet in project ignite-3 by apache.

the class IgniteMergeJoin method extendCollation.

/**
 * This function extends collation by appending new collation fields defined on keys.
 */
private static RelCollation extendCollation(RelCollation collation, List<Integer> keys) {
    List<RelFieldCollation> fieldsForNewCollation = new ArrayList<>(keys.size());
    fieldsForNewCollation.addAll(collation.getFieldCollations());
    ImmutableBitSet keysBitset = ImmutableBitSet.of(keys);
    ImmutableBitSet colKeysBitset = ImmutableBitSet.of(collation.getKeys());
    ImmutableBitSet exceptBitset = keysBitset.except(colKeysBitset);
    for (Integer i : exceptBitset) {
        fieldsForNewCollation.add(new RelFieldCollation(i));
    }
    return RelCollations.of(fieldsForNewCollation);
}

Example 73 with ImmutableBitSet

use of org.apache.calcite.util.ImmutableBitSet in project ignite-3 by apache.

the class IgniteMergeJoin method passThroughCollation.

/**
 * {@inheritDoc}
 */
@Override
public Pair<RelTraitSet, List<RelTraitSet>> passThroughCollation(RelTraitSet required, List<RelTraitSet> inputTraits) {
    RelCollation collation = TraitUtils.collation(required);
    RelTraitSet left = inputTraits.get(0);
    RelTraitSet right = inputTraits.get(1);
    if (joinType == FULL) {
        return defaultCollationPair(required, left, right);
    }
    int leftInputFieldCount = this.left.getRowType().getFieldCount();
    List<Integer> reqKeys = RelCollations.ordinals(collation);
    List<Integer> leftKeys = joinInfo.leftKeys.toIntegerList();
    List<Integer> rightKeys = joinInfo.rightKeys.incr(leftInputFieldCount).toIntegerList();
    ImmutableBitSet reqKeySet = ImmutableBitSet.of(reqKeys);
    ImmutableBitSet leftKeySet = ImmutableBitSet.of(joinInfo.leftKeys);
    ImmutableBitSet rightKeySet = ImmutableBitSet.of(rightKeys);
    RelCollation nodeCollation;
    RelCollation leftCollation;
    RelCollation rightCollation;
    if (reqKeySet.equals(leftKeySet)) {
        if (joinType == RIGHT) {
            return defaultCollationPair(required, left, right);
        }
        nodeCollation = collation;
        leftCollation = collation;
        rightCollation = collation.apply(buildTransposeMapping(true));
    } else if (containsOrderless(leftKeys, collation)) {
        if (joinType == RIGHT) {
            return defaultCollationPair(required, left, right);
        }
        // if sort keys are subset of left join keys, we can extend collations to make sure all join
        // keys are sorted.
        nodeCollation = collation;
        leftCollation = extendCollation(collation, leftKeys);
        rightCollation = leftCollation.apply(buildTransposeMapping(true));
    } else if (containsOrderless(collation, leftKeys) && reqKeys.stream().allMatch(i -> i < leftInputFieldCount)) {
        if (joinType == RIGHT) {
            return defaultCollationPair(required, left, right);
        }
        // if sort keys are superset of left join keys, and left join keys is prefix of sort keys
        // (order not matter), also sort keys are all from left join input.
        nodeCollation = collation;
        leftCollation = collation;
        rightCollation = leftCollation.apply(buildTransposeMapping(true));
    } else if (reqKeySet.equals(rightKeySet)) {
        if (joinType == LEFT) {
            return defaultCollationPair(required, left, right);
        }
        nodeCollation = collation;
        rightCollation = RelCollations.shift(collation, -leftInputFieldCount);
        leftCollation = rightCollation.apply(buildTransposeMapping(false));
    } else if (containsOrderless(rightKeys, collation)) {
        if (joinType == LEFT) {
            return defaultCollationPair(required, left, right);
        }
        nodeCollation = collation;
        rightCollation = RelCollations.shift(extendCollation(collation, rightKeys), -leftInputFieldCount);
        leftCollation = rightCollation.apply(buildTransposeMapping(false));
    } else {
        return defaultCollationPair(required, left, right);
    }
    return Pair.of(required.replace(nodeCollation), ImmutableList.of(left.replace(leftCollation), right.replace(rightCollation)));
}

Example 74 with ImmutableBitSet

use of org.apache.calcite.util.ImmutableBitSet in project ignite-3 by apache.

the class ExposeIndexRule method onMatch.

/**
 * {@inheritDoc}
 */
@Override
public void onMatch(RelOptRuleCall call) {
    IgniteLogicalTableScan scan = call.rel(0);
    RelOptCluster cluster = scan.getCluster();
    RelOptTable optTable = scan.getTable();
    InternalIgniteTable igniteTable = optTable.unwrap(InternalIgniteTable.class);
    List<RexNode> proj = scan.projects();
    RexNode condition = scan.condition();
    ImmutableBitSet requiredCols = scan.requiredColumns();
    List<IgniteLogicalIndexScan> indexes = igniteTable.indexes().keySet().stream().map(idxName -> igniteTable.toRel(cluster, optTable, idxName, proj, condition, requiredCols)).collect(Collectors.toList());
    if (indexes.isEmpty()) {
        return;
    }
    Map<RelNode, RelNode> equivMap = new HashMap<>(indexes.size());
    for (int i = 1; i < indexes.size(); i++) {
        equivMap.put(indexes.get(i), scan);
    }
    call.transformTo(indexes.get(0), equivMap);
}

Example 75 with ImmutableBitSet

use of org.apache.calcite.util.ImmutableBitSet in project calcite by apache.

the class EnumerableSortedAggregate method passThroughTraits.

@Override
@Nullable
public Pair<RelTraitSet, List<RelTraitSet>> passThroughTraits(final RelTraitSet required) {
    if (!isSimple(this)) {
        return null;
    }
    RelTraitSet inputTraits = getInput().getTraitSet();
    RelCollation collation = requireNonNull(required.getCollation(), () -> "collation trait is null, required traits are " + required);
    ImmutableBitSet requiredKeys = ImmutableBitSet.of(RelCollations.ordinals(collation));
    ImmutableBitSet groupKeys = ImmutableBitSet.range(groupSet.cardinality());
    Mappings.TargetMapping mapping = Mappings.source(groupSet.toList(), input.getRowType().getFieldCount());
    if (requiredKeys.equals(groupKeys)) {
        RelCollation inputCollation = RexUtil.apply(mapping, collation);
        return Pair.of(required, ImmutableList.of(inputTraits.replace(inputCollation)));
    } else if (groupKeys.contains(requiredKeys)) {
        // group by a,b,c order by c,b
        List<RelFieldCollation> list = new ArrayList<>(collation.getFieldCollations());
        groupKeys.except(requiredKeys).forEach(k -> list.add(new RelFieldCollation(k)));
        RelCollation aggCollation = RelCollations.of(list);
        RelCollation inputCollation = RexUtil.apply(mapping, aggCollation);
        return Pair.of(traitSet.replace(aggCollation), ImmutableList.of(inputTraits.replace(inputCollation)));
    }
    // nothing we can do to propagate traits to child nodes.
    return null;
}