Examples with GroupKey org.apache.calcite.tools.RelBuilder.GroupKey used on opensource projects

Example 1 with GroupKey

use of org.apache.calcite.tools.RelBuilder.GroupKey in project hive by apache.

the class HiveSemiJoinRule method recreateAggregateOperator.

protected static RelNode recreateAggregateOperator(RelBuilder builder, int[] adjustments, Aggregate topAggregate, RelNode newInputOperator) {
    builder.push(newInputOperator);
    ImmutableBitSet.Builder newGroupSet = ImmutableBitSet.builder();
    for (int pos : topAggregate.getGroupSet()) {
        newGroupSet.set(pos + adjustments[pos]);
    }
    GroupKey groupKey;
    if (topAggregate.getGroupType() == Group.SIMPLE) {
        groupKey = builder.groupKey(newGroupSet.build());
    } else {
        List<ImmutableBitSet> newGroupSets = new ArrayList<>();
        for (ImmutableBitSet groupingSet : topAggregate.getGroupSets()) {
            ImmutableBitSet.Builder newGroupingSet = ImmutableBitSet.builder();
            for (int pos : groupingSet) {
                newGroupingSet.set(pos + adjustments[pos]);
            }
            newGroupSets.add(newGroupingSet.build());
        }
        groupKey = builder.groupKey(newGroupSet.build(), newGroupSets);
    }
    List<AggregateCall> newAggCallList = new ArrayList<>();
    for (AggregateCall aggregateCall : topAggregate.getAggCallList()) {
        List<Integer> newArgList = aggregateCall.getArgList().stream().map(pos -> pos + adjustments[pos]).collect(Collectors.toList());
        int newFilterArg = aggregateCall.filterArg != -1 ? aggregateCall.filterArg + adjustments[aggregateCall.filterArg] : -1;
        RelCollation newCollation = aggregateCall.getCollation() != null ? RelCollations.of(aggregateCall.getCollation().getFieldCollations().stream().map(fc -> fc.withFieldIndex(fc.getFieldIndex() + adjustments[fc.getFieldIndex()])).collect(Collectors.toList())) : null;
        newAggCallList.add(aggregateCall.copy(newArgList, newFilterArg, newCollation));
    }
    return builder.push(newInputOperator).aggregate(groupKey, newAggCallList).build();
}