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

Example 1 with AggregationFunction

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

the class HyperLogLogDistinctAggregationBenchmark method setUp.

@Setup
public void setUp() throws Exception {
    hash = new MurmurHash3.Hash128();
    final InputCollectExpression inExpr0 = new InputCollectExpression(0);
    Functions functions = new ModulesBuilder().add(new ExtraFunctionsModule()).createInjector().getInstance(Functions.class);
    final HyperLogLogDistinctAggregation hllAggregation = (HyperLogLogDistinctAggregation) functions.getQualified(Signature.aggregate(HyperLogLogDistinctAggregation.NAME, DataTypes.STRING.getTypeSignature(), DataTypes.LONG.getTypeSignature()), List.of(DataTypes.STRING), DataTypes.STRING);
    onHeapMemoryManager = new OnHeapMemoryManager(bytes -> {
    });
    offHeapMemoryManager = new OffHeapMemoryManager();
    hyperLogLogPlusPlus = new HyperLogLogPlusPlus(HyperLogLogPlusPlus.DEFAULT_PRECISION, onHeapMemoryManager::allocate);
    onHeapCollector = new AggregateCollector(Collections.singletonList(inExpr0), RamAccounting.NO_ACCOUNTING, onHeapMemoryManager, Version.CURRENT, AggregateMode.ITER_FINAL, new AggregationFunction[] { hllAggregation }, Version.CURRENT, new Input[][] { { inExpr0 } }, new Input[] { Literal.BOOLEAN_TRUE });
    offHeapCollector = new AggregateCollector(Collections.singletonList(inExpr0), RamAccounting.NO_ACCOUNTING, offHeapMemoryManager, Version.CURRENT, AggregateMode.ITER_FINAL, new AggregationFunction[] { hllAggregation }, Version.CURRENT, new Input[][] { { inExpr0 } }, new Input[] { Literal.BOOLEAN_TRUE });
}

Example 2 with AggregationFunction

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

the class WindowProjector method fromProjection.

public static Projector fromProjection(WindowAggProjection projection, NodeContext nodeCtx, InputFactory inputFactory, TransactionContext txnCtx, RamAccounting ramAccounting, MemoryManager memoryManager, Version minNodeVersion, Version indexVersionCreated, IntSupplier numThreads, Executor executor) {
    var windowFunctionSymbols = projection.windowFunctions();
    var numWindowFunctions = windowFunctionSymbols.size();
    assert numWindowFunctions > 0 : "WindowAggProjection must have at least 1 window function.";
    ArrayList<WindowFunction> windowFunctions = new ArrayList<>(numWindowFunctions);
    ArrayList<CollectExpression<Row, ?>> windowFuncArgsExpressions = new ArrayList<>(numWindowFunctions);
    Input[][] windowFuncArgsInputs = new Input[numWindowFunctions][];
    Boolean[] ignoreNulls = new Boolean[numWindowFunctions];
    for (int idx = 0; idx < numWindowFunctions; idx++) {
        var windowFunctionSymbol = windowFunctionSymbols.get(idx);
        InputFactory.Context<CollectExpression<Row, ?>> ctx = inputFactory.ctxForInputColumns(txnCtx);
        ctx.add(windowFunctionSymbol.arguments());
        FunctionImplementation impl = nodeCtx.functions().getQualified(windowFunctionSymbol, txnCtx.sessionSettings().searchPath());
        assert impl != null : "Function implementation not found using full qualified lookup";
        if (impl instanceof AggregationFunction) {
            var filterInputFactoryCtx = inputFactory.ctxForInputColumns(txnCtx);
            var filterSymbol = windowFunctionSymbol.filter();
            // noinspection unchecked
            Input<Boolean> filterInput = filterSymbol == null ? Literal.BOOLEAN_TRUE : (Input<Boolean>) filterInputFactoryCtx.add(filterSymbol);
            ExpressionsInput<Row, Boolean> filter = new ExpressionsInput<>(filterInput, filterInputFactoryCtx.expressions());
            windowFunctions.add(new AggregateToWindowFunctionAdapter((AggregationFunction) impl, filter, indexVersionCreated, ramAccounting, memoryManager, minNodeVersion));
        } else if (impl instanceof WindowFunction) {
            windowFunctions.add((WindowFunction) impl);
        } else {
            throw new AssertionError("Function needs to be either a window or an aggregate function");
        }
        windowFuncArgsExpressions.addAll(ctx.expressions());
        windowFuncArgsInputs[idx] = ctx.topLevelInputs().toArray(new Input[0]);
        ignoreNulls[idx] = windowFunctionSymbol.ignoreNulls();
    }
    var windowDefinition = projection.windowDefinition();
    var partitions = windowDefinition.partitions();
    Supplier<InputFactory.Context<CollectExpression<Row, ?>>> createInputFactoryContext = () -> inputFactory.ctxForInputColumns(txnCtx);
    int arrayListElementOverHead = 32;
    RowAccountingWithEstimators accounting = new RowAccountingWithEstimators(Symbols.typeView(projection.standalone()), ramAccounting, arrayListElementOverHead);
    Comparator<Object[]> cmpPartitionBy = partitions.isEmpty() ? null : createComparator(createInputFactoryContext, new OrderBy(windowDefinition.partitions()));
    Comparator<Object[]> cmpOrderBy = createComparator(createInputFactoryContext, windowDefinition.orderBy());
    int numCellsInSourceRow = projection.standalone().size();
    ComputeFrameBoundary<Object[]> computeFrameStart = createComputeStartFrameBoundary(numCellsInSourceRow, txnCtx, nodeCtx, windowDefinition, cmpOrderBy);
    ComputeFrameBoundary<Object[]> computeFrameEnd = createComputeEndFrameBoundary(numCellsInSourceRow, txnCtx, nodeCtx, windowDefinition, cmpOrderBy);
    return sourceRows -> WindowFunctionBatchIterator.of(sourceRows, accounting, computeFrameStart, computeFrameEnd, cmpPartitionBy, cmpOrderBy, numCellsInSourceRow, numThreads, executor, windowFunctions, windowFuncArgsExpressions, ignoreNulls, windowFuncArgsInputs);
}

Example 3 with AggregationFunction

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

the class CollectSetAggregationTest method testLongSerialization.

@Test
public void testLongSerialization() throws Exception {
    AggregationFunction impl = (AggregationFunction) nodeCtx.functions().get(null, "collect_set", List.of(Literal.of(DataTypes.LONG, null)), SearchPath.pathWithPGCatalogAndDoc());
    Object state = impl.newState(RAM_ACCOUNTING, Version.CURRENT, Version.CURRENT, memoryManager);
    BytesStreamOutput streamOutput = new BytesStreamOutput();
    impl.partialType().streamer().writeValueTo(streamOutput, state);
    Object newState = impl.partialType().streamer().readValueFrom(streamOutput.bytes().streamInput());
    assertEquals(state, newState);
}

Example 4 with AggregationFunction

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

the class PercentileAggregationTest method testSingleItemFractionsArgumentResultsInArrayResult.

@Test
public void testSingleItemFractionsArgumentResultsInArrayResult() {
    AggregationFunction impl = (AggregationFunction<?, ?>) nodeCtx.functions().getQualified(Signature.aggregate(PercentileAggregation.NAME, DataTypes.LONG.getTypeSignature(), DataTypes.DOUBLE_ARRAY.getTypeSignature(), DataTypes.DOUBLE_ARRAY.getTypeSignature()), List.of(DataTypes.LONG, DataTypes.DOUBLE_ARRAY), DataTypes.DOUBLE_ARRAY);
    RamAccounting ramAccounting = RamAccounting.NO_ACCOUNTING;
    Object state = impl.newState(ramAccounting, Version.CURRENT, Version.CURRENT, memoryManager);
    Literal<List<Double>> fractions = Literal.of(Collections.singletonList(0.95D), DataTypes.DOUBLE_ARRAY);
    impl.iterate(ramAccounting, memoryManager, state, Literal.of(10L), fractions);
    impl.iterate(ramAccounting, memoryManager, state, Literal.of(20L), fractions);
    Object result = impl.terminatePartial(ramAccounting, state);
    assertThat("result must be an array", result, instanceOf(List.class));
}

Example 5 with AggregationFunction

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

the class AbstractWindowFunctionTest method assertEvaluate.

@SuppressWarnings("unchecked")
protected <T> void assertEvaluate(String functionExpression, Matcher<T> expectedValue, List<ColumnIdent> rowsColumnDescription, Object[]... inputRows) throws Throwable {
    performInputSanityChecks(inputRows);
    Symbol normalizedFunctionSymbol = sqlExpressions.normalize(sqlExpressions.asSymbol(functionExpression));
    assertThat(normalizedFunctionSymbol, instanceOf(io.crate.expression.symbol.WindowFunction.class));
    var windowFunctionSymbol = (io.crate.expression.symbol.WindowFunction) normalizedFunctionSymbol;
    ReferenceResolver<InputCollectExpression> referenceResolver = r -> new InputCollectExpression(rowsColumnDescription.indexOf(r.column()));
    var sourceSymbols = Lists2.map(rowsColumnDescription, x -> sqlExpressions.normalize(sqlExpressions.asSymbol(x.sqlFqn())));
    ensureInputRowsHaveCorrectType(sourceSymbols, inputRows);
    var argsCtx = inputFactory.ctxForRefs(txnCtx, referenceResolver);
    argsCtx.add(windowFunctionSymbol.arguments());
    FunctionImplementation impl = sqlExpressions.nodeCtx.functions().getQualified(windowFunctionSymbol, txnCtx.sessionSettings().searchPath());
    assert impl instanceof WindowFunction || impl instanceof AggregationFunction : "Got " + impl + " but expected a window function";
    WindowFunction windowFunctionImpl;
    if (impl instanceof AggregationFunction) {
        windowFunctionImpl = new AggregateToWindowFunctionAdapter((AggregationFunction) impl, new ExpressionsInput<>(Literal.BOOLEAN_TRUE, List.of()), Version.CURRENT, RamAccounting.NO_ACCOUNTING, memoryManager, Version.CURRENT);
    } else {
        windowFunctionImpl = (WindowFunction) impl;
    }
    int numCellsInSourceRows = inputRows[0].length;
    var windowDef = windowFunctionSymbol.windowDefinition();
    var partitionOrderBy = windowDef.partitions().isEmpty() ? null : new OrderBy(windowDef.partitions());
    Comparator<Object[]> cmpOrderBy = createComparator(() -> inputFactory.ctxForRefs(txnCtx, referenceResolver), windowDef.orderBy());
    InputColumns.SourceSymbols inputColSources = new InputColumns.SourceSymbols(sourceSymbols);
    var mappedWindowDef = windowDef.map(s -> InputColumns.create(s, inputColSources));
    BatchIterator<Row> iterator = WindowFunctionBatchIterator.of(InMemoryBatchIterator.of(Arrays.stream(inputRows).map(RowN::new).collect(Collectors.toList()), SENTINEL, true), new IgnoreRowAccounting(), WindowProjector.createComputeStartFrameBoundary(numCellsInSourceRows, txnCtx, sqlExpressions.nodeCtx, mappedWindowDef, cmpOrderBy), WindowProjector.createComputeEndFrameBoundary(numCellsInSourceRows, txnCtx, sqlExpressions.nodeCtx, mappedWindowDef, cmpOrderBy), createComparator(() -> inputFactory.ctxForRefs(txnCtx, referenceResolver), partitionOrderBy), cmpOrderBy, numCellsInSourceRows, () -> 1, Runnable::run, List.of(windowFunctionImpl), argsCtx.expressions(), new Boolean[] { windowFunctionSymbol.ignoreNulls() }, argsCtx.topLevelInputs().toArray(new Input[0]));
    List<Object> actualResult;
    try {
        actualResult = BatchIterators.collect(iterator, Collectors.mapping(row -> row.get(numCellsInSourceRows), Collectors.toList())).get(5, TimeUnit.SECONDS);
    } catch (ExecutionException e) {
        throw e.getCause();
    }
    assertThat((T) actualResult, expectedValue);
}