Examples with AccumulatorFunction io.ordinate.engine.function.aggregate.AccumulatorFunction used on opensource projects

Example 1 with AccumulatorFunction

use of io.ordinate.engine.function.aggregate.AccumulatorFunction in project Mycat2 by MyCATApache.

the class ExecuteCompiler method agg.

public static PhysicalPlan agg(PhysicalPlan input, AggImpl impl, List<GroupKeys> groupKeyList, List<AccumulatorFunction> groupExprs) {
    PhysicalPlan physicalPlan = null;
    ArrayList<ArrowType> objects = new ArrayList<>();
    for (AccumulatorFunction groupExpr : groupExprs) {
        ArrowType arrowType = groupExpr.getType().getArrowType();
        objects.add(arrowType);
    }
    Schema outputSchema = SchemaBuilder.ofArrowType(objects).toArrow();
    // AggregateVectorExpression[] aggregateVectorExpressions = getAggregateVectorExpressions(groupExprs);
    if (groupKeyList.isEmpty() && groupExprs.stream().allMatch(i -> i.isVector())) {
        physicalPlan = NoKeysAggPlan.create(outputSchema, input, groupExprs.toArray(new AccumulatorFunction[] {}));
    } else {
        switch(impl) {
            case HASH:
                physicalPlan = new GroupByKeyWithAggPlan(input, groupKeyList.toArray(new GroupKeys[] {}), groupExprs.toArray(new AccumulatorFunction[] {}), outputSchema);
                break;
            case MERGE:
                break;
        }
    }
    return physicalPlan;
}

Example 2 with AccumulatorFunction

use of io.ordinate.engine.function.aggregate.AccumulatorFunction in project Mycat2 by MyCATApache.

the class CalciteCompiler method convertAggregate.

public PhysicalPlan convertAggregate(Aggregate aggregate) {
    PhysicalPlan input = convert(aggregate.getInput());
    List<Integer> groupSet = aggregate.getGroupSet().asList();
    List<AggregateCall> aggCallList = aggregate.getAggCallList();
    GroupKeys[] groupSets = aggregate.getGroupSets().stream().map(i -> GroupKeys.of(i.toArray())).toArray(n -> new GroupKeys[n]);
    AccumulatorFunction[] accumulatorFunctions = new AccumulatorFunction[groupSet.size() + aggCallList.size()];
    int index = 0;
    for (Integer integer : groupSet) {
        accumulatorFunctions[index] = ExecuteCompiler.anyValue(input, integer);
        index++;
    }
    index = groupSet.size();
    for (AggregateCall aggregateCall : aggCallList) {
        List<Integer> argList = aggregateCall.getArgList();
        SqlKind kind = aggregateCall.getAggregation().kind;
        AccumulatorFunction accumulatorFunction = null;
        switch(kind) {
            case SUM:
            case SUM0:
                {
                    Integer integer = argList.get(0);
                    accumulatorFunction = executeCompiler.sum(input, integer);
                    break;
                }
            case AVG:
                {
                    accumulatorFunction = executeCompiler.avg(argList.get(0));
                    break;
                }
            case COUNT:
                {
                    boolean distinct = aggregateCall.isDistinct();
                    if (distinct) {
                        // todo check
                        accumulatorFunction = executeCompiler.countDistinct(input, argList.get(0));
                    } else {
                        if (argList.size() == 0) {
                            accumulatorFunction = executeCompiler.count();
                        } else {
                            accumulatorFunction = executeCompiler.count(argList.get(0));
                        }
                    }
                    break;
                }
            case ANY_VALUE:
                {
                    accumulatorFunction = executeCompiler.anyValue(input, argList.get(0));
                    break;
                }
            case MAX:
                {
                    accumulatorFunction = executeCompiler.max(input, argList.get(0));
                    break;
                }
            case MIN:
                {
                    accumulatorFunction = executeCompiler.min(input, argList.get(0));
                    break;
                }
        }
        Objects.requireNonNull(accumulatorFunction);
        accumulatorFunctions[index] = accumulatorFunction;
        ++index;
    }
    int slotOriginalSize = input.schema().getFields().size();
    int slotInc = slotOriginalSize;
    Map<Integer, Map<InnerType, Integer>> indexToTypeMap = new HashMap<>();
    List<Field> fieldList = input.schema().getFields();
    for (AccumulatorFunction accumulatorFunction : accumulatorFunctions) {
        InnerType aggInputType = accumulatorFunction.getInputType();
        InnerType aggOutputType = accumulatorFunction.getOutputType();
        int aggInputIndex = accumulatorFunction.getInputColumnIndex();
        InnerType aggInputSourceType = InnerType.from(fieldList.get(aggInputIndex).getType());
        Function column = ExecuteCompiler.column(aggInputIndex, input.schema());
        Map<InnerType, Integer> indexSlot = indexToTypeMap.computeIfAbsent(aggInputIndex, integer -> {
            HashMap<InnerType, Integer> map = new HashMap<>();
            map.put(aggInputSourceType, integer);
            return map;
        });
        if (aggInputType != null) {
            if (!indexSlot.containsValue(aggInputIndex)) {
                indexSlot.put(aggInputType, aggInputIndex);
            } else {
                accumulatorFunction.setInputColumnIndex(slotInc);
                indexSlot.put(aggInputType, slotInc);
                slotInc++;
            }
        }
    }
    Function[] exprs = new Function[slotInc];
    for (int i = 0; i < slotOriginalSize; i++) {
        exprs[i] = ExecuteCompiler.column(i, input.schema());
    }
    if (slotInc > slotOriginalSize) {
        for (Map.Entry<Integer, Map<InnerType, Integer>> e1 : indexToTypeMap.entrySet()) {
            for (Map.Entry<InnerType, Integer> e2 : e1.getValue().entrySet()) {
                Integer inIndex = e1.getKey();
                Integer outIndex = e2.getValue();
                if (exprs[outIndex] == null) {
                    exprs[outIndex] = (ExecuteCompiler.cast(ExecuteCompiler.column(inIndex, input.schema()), e2.getKey()));
                }
            }
        }
    }
    input = executeCompiler.project(input, exprs);
    return executeCompiler.agg(input, ExecuteCompiler.AggImpl.HASH, Arrays.asList(groupSets), Arrays.asList(accumulatorFunctions));
}

Example 3 with AccumulatorFunction

use of io.ordinate.engine.function.aggregate.AccumulatorFunction in project Mycat2 by MyCATApache.

the class GroupByKeyWithAggPlan method execute.

@Override
public Observable<VectorSchemaRoot> execute(RootContext rootContext) {
    List<Field> fields = schema().getFields();
    InnerType[] innerTypes = schema().getFields().stream().map(i -> InnerType.from(i.getType())).toArray(n -> new InnerType[n]);
    if (groupByKeys.length > 0) {
        ColumnTypes arrayColumnTypes = RecordUtil.getArrayColumnTypes(accumulators);
        Map map = MapFactory.createMap2(innerTypes, arrayColumnTypes);
        RecordSink[] recordSinks = buildRecordSink(fields);
        return physicalPlan.execute(rootContext).reduce(map, (mapKey, input) -> {
            int rowCount = input.getRowCount();
            VectorBatchRecord record = new VectorBatchRecord(input);
            for (RecordSink recordSink : recordSinks) {
                for (int rowId = 0; rowId < rowCount; rowId++) {
                    record.setPosition(rowId);
                    MapKey key = mapKey.withKey();
                    RecordSetter recordSinkSPI = RecordSinkFactory.INSTANCE.getRecordSinkSPI(key);
                    recordSink.copy(record, recordSinkSPI);
                    MapValue value = key.createValue();
                    if (value.isNew()) {
                        for (AccumulatorFunction accumulator : accumulators) {
                            accumulator.computeFirst(value, record);
                        }
                    } else {
                        for (AccumulatorFunction accumulator : accumulators) {
                            accumulator.computeNext(value, record);
                        }
                    }
                }
            }
            physicalPlan.eachFree(input);
            return mapKey;
        }).map(map1 -> {
            int size = (int) map1.size();
            VectorSchemaRoot output = rootContext.getVectorSchemaRoot(schema(), size);
            RecordCursor cursor = map1.getCursor();
            cursor.toTop();
            int index = 0;
            while (cursor.hasNext()) {
                Record record = cursor.getRecord();
                functionSink.copy(accumulators, RecordUtil.wrapAsAggRecord(record), index++, output);
            }
            output.setRowCount(index);
            map1.clear();
            return output;
        }).toObservable().doOnComplete(() -> map.close());
    } else {
        SimpleMapValue mapValue = new SimpleMapValue(RecordUtil.getContextSize(accumulators));
        return physicalPlan.execute(rootContext).reduce(mapValue, new BiFunction<SimpleMapValue, VectorSchemaRoot, SimpleMapValue>() {

            AtomicBoolean first = new AtomicBoolean(true);

            @Override
            public SimpleMapValue apply(SimpleMapValue simpleMapValue, VectorSchemaRoot root) throws Throwable {
                int rowCount = root.getRowCount();
                VectorBatchRecord record = new VectorBatchRecord(root);
                if (first.compareAndSet(true, false)) {
                    record.setPosition(0);
                    for (AccumulatorFunction accumulator : accumulators) {
                        accumulator.computeFirst(mapValue, record);
                    }
                    for (int i = 1; i < rowCount; i++) {
                        record.setPosition(i);
                        for (AccumulatorFunction accumulator : accumulators) {
                            accumulator.computeNext(mapValue, record);
                        }
                    }
                } else {
                    for (int i = 0; i < rowCount; i++) {
                        record.setPosition(i);
                        for (AccumulatorFunction accumulator : accumulators) {
                            accumulator.computeNext(mapValue, record);
                        }
                    }
                }
                root.close();
                return simpleMapValue;
            }
        }).map(new Function<SimpleMapValue, VectorSchemaRoot>() {

            @Override
            public VectorSchemaRoot apply(SimpleMapValue simpleMapValue) throws Throwable {
                VectorSchemaRoot vectorSchemaRoot = rootContext.getVectorSchemaRoot(schema(), 1);
                vectorSchemaRoot.setRowCount(1);
                functionSink.copy(accumulators, RecordUtil.wrapAsAggRecord(simpleMapValue), 0, vectorSchemaRoot);
                return vectorSchemaRoot;
            }
        }).toObservable();
    }
}