Examples with MycatHashAggregate io.mycat.calcite.physical.MycatHashAggregate used on opensource projects

Example 1 with MycatHashAggregate

use of io.mycat.calcite.physical.MycatHashAggregate in project Mycat2 by MyCATApache.

the class SQLRBORewriter method aggregate.

public static Optional<RelNode> aggregate(RelNode original, Aggregate aggregate) {
    RelNode input = original;
    Distribution dataNodeInfo = null;
    MycatView view = null;
    ImmutableList<RelHint> hints = aggregate.getHints();
    if (input instanceof MycatView) {
        dataNodeInfo = ((MycatView) input).getDistribution();
        view = (MycatView) original;
        if (view.banPushdown()) {
            return Optional.empty();
        }
        input = ((MycatView) input).getRelNode();
    }
    if (dataNodeInfo == null) {
        return Optional.empty();
    }
    if (dataNodeInfo.type() == Distribution.Type.PHY || dataNodeInfo.type() == Distribution.Type.BROADCAST) {
        input = aggregate.copy(aggregate.getTraitSet(), ImmutableList.of(input));
        return Optional.of(view.changeTo(input, dataNodeInfo));
    } else {
        ImmutableBitSet groupSet = aggregate.getGroupSet();
        RelMetadataQuery metadataQuery = aggregate.getCluster().getMetadataQuery();
        IdentityHashMap<MycatLogicTable, Set<String>> shardingKeysMap = new IdentityHashMap<>();
        for (Integer integer : groupSet) {
            RelColumnOrigin columnOrigin = metadataQuery.getColumnOrigin(input, integer);
            if (columnOrigin == null || columnOrigin.isDerived()) {
                continue;
            }
            MycatLogicTable mycatLogicTable = columnOrigin.getOriginTable().unwrap(MycatLogicTable.class);
            if (mycatLogicTable == null || !mycatLogicTable.isSharding()) {
                continue;
            }
            ShardingTableHandler tableHandler = (ShardingTableHandler) mycatLogicTable.getTable();
            SimpleColumnInfo simpleColumnInfo = tableHandler.getColumns().get(columnOrigin.getOriginColumnOrdinal());
            if (simpleColumnInfo.isShardingKey()) {
                Set<String> shardingKeySet = shardingKeysMap.computeIfAbsent(mycatLogicTable, s -> new HashSet<>());
                shardingKeySet.add(simpleColumnInfo.getColumnName());
                if (tableHandler.function().requireShardingKeys(shardingKeySet)) {
                    input = aggregate.copy(aggregate.getTraitSet(), ImmutableList.of(input));
                    return Optional.of(view.changeTo(input, dataNodeInfo));
                }
            }
        }
        RelNode backup = input;
        if (!(input instanceof Union)) {
            input = LogicalUnion.create(ImmutableList.of(input, input), true);
            input = LogicalAggregate.create(input, aggregate.getHints(), aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList());
        }
        HepProgramBuilder hepProgram = new HepProgramBuilder();
        hepProgram.addMatchLimit(512);
        hepProgram.addRuleInstance(MycatAggregateUnionTransposeRule.Config.DEFAULT.toRule());
        HepPlanner planner = new HepPlanner(hepProgram.build());
        planner.setRoot(input);
        RelNode bestExp = planner.findBestExp();
        if (bestExp instanceof Aggregate) {
            Aggregate mergeAgg = (Aggregate) bestExp;
            if (mergeAgg.getInput() instanceof Union && mergeAgg.getInput(0).getInput(0) instanceof Aggregate) {
                MycatView multiView = view.changeTo(mergeAgg.getInput(0).getInput(0), dataNodeInfo);
                MycatHashAggregate mycatHashAggregate = MycatHashAggregate.create(mergeAgg.getTraitSet(), hints, multiView, mergeAgg.getGroupSet(), mergeAgg.getGroupSets(), mergeAgg.getAggCallList());
                return Optional.of(mycatHashAggregate);
            }
        }
        {
            return splitAggregate(view, aggregate);
        }
    }
}