Examples with AggregationFunction io.crate.execution.engine.aggregation.AggregationFunction used on opensource projects

Example 6 with AggregationFunction

use of io.crate.execution.engine.aggregation.AggregationFunction in project crate by crate.

the class CollectSetAggregationTest method test_value_adding_and_removal.

@Test
public void test_value_adding_and_removal() {
    AggregationFunction impl = (AggregationFunction) nodeCtx.functions().get(null, "collect_set", List.of(Literal.of(DataTypes.LONG, null)), SearchPath.pathWithPGCatalogAndDoc());
    AggregationFunction aggregationFunction = impl.optimizeForExecutionAsWindowFunction();
    Object state = aggregationFunction.newState(RAM_ACCOUNTING, Version.CURRENT, Version.CURRENT, memoryManager);
    state = aggregationFunction.iterate(RAM_ACCOUNTING, memoryManager, state, Literal.of(10));
    state = aggregationFunction.iterate(RAM_ACCOUNTING, memoryManager, state, Literal.of(10));
    aggregationFunction.removeFromAggregatedState(RAM_ACCOUNTING, state, new Input[] { Literal.of(10) });
    aggregationFunction.removeFromAggregatedState(RAM_ACCOUNTING, state, new Input[] { Literal.of(10) });
    Object values = aggregationFunction.terminatePartial(RAM_ACCOUNTING, state);
    assertThat((List<Object>) values, Matchers.empty());
}

Example 7 with AggregationFunction

use of io.crate.execution.engine.aggregation.AggregationFunction in project crate by crate.

the class GroupByOptimizedIterator method initStates.

private static Object[] initStates(List<AggregationContext> aggregations, RamAccounting ramAccounting, MemoryManager memoryManager, Version minNodeVersion) {
    Object[] states = new Object[aggregations.size()];
    for (int i = 0; i < aggregations.size(); i++) {
        AggregationContext aggregation = aggregations.get(i);
        AggregationFunction function = aggregation.function();
        var newState = function.newState(ramAccounting, Version.CURRENT, minNodeVersion, memoryManager);
        if (InputCondition.matches(aggregation.filter())) {
            // noinspection unchecked
            states[i] = function.iterate(ramAccounting, memoryManager, newState, aggregation.inputs());
        } else {
            states[i] = newState;
        }
    }
    return states;
}

Example 8 with AggregationFunction

use of io.crate.execution.engine.aggregation.AggregationFunction in project crate by crate.

the class ProjectionBuilder method getAggregations.

private ArrayList<Aggregation> getAggregations(Collection<Function> functions, AggregateMode mode, InputColumns.SourceSymbols sourceSymbols, SearchPath searchPath, java.util.function.Function<Symbol, Symbol> subQueryAndParamBinder) {
    ArrayList<Aggregation> aggregations = new ArrayList<>(functions.size());
    for (Function function : functions) {
        assert function.type() == FunctionType.AGGREGATE : "function type must be " + FunctionType.AGGREGATE;
        List<Symbol> aggregationInputs;
        Symbol filterInput;
        switch(mode) {
            case ITER_FINAL:
            case ITER_PARTIAL:
                // ITER means that there is no aggregation part upfront, therefore the input
                // symbols need to be in arguments
                aggregationInputs = InputColumns.create(function.arguments(), sourceSymbols);
                Symbol filter = function.filter();
                if (filter != null) {
                    filterInput = InputColumns.create(filter, sourceSymbols);
                } else {
                    filterInput = Literal.BOOLEAN_TRUE;
                }
                break;
            case PARTIAL_FINAL:
                aggregationInputs = List.of(sourceSymbols.getICForSource(function));
                filterInput = Literal.BOOLEAN_TRUE;
                break;
            default:
                throw new AssertionError("Invalid mode: " + mode.name());
        }
        AggregationFunction<?, ?> aggregationFunction = (AggregationFunction<?, ?>) nodeCtx.functions().getQualified(function, searchPath);
        assert aggregationFunction != null : "Aggregation function implementation not found using full qualified lookup: " + function;
        var valueType = mode.returnType(aggregationFunction);
        var functionInfo = FunctionInfo.of(aggregationFunction.signature(), aggregationFunction.boundSignature().getArgumentDataTypes(), valueType);
        Aggregation aggregation = new Aggregation(aggregationFunction.signature(), functionInfo, aggregationFunction.boundSignature().getReturnType().createType(), valueType, Lists2.map(aggregationInputs, subQueryAndParamBinder), subQueryAndParamBinder.apply(filterInput));
        aggregations.add(aggregation);
    }
    return aggregations;
}