Examples with CollectExpression io.crate.execution.engine.collect.CollectExpression used on opensource projects

Example 6 with CollectExpression

use of io.crate.execution.engine.collect.CollectExpression 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 7 with CollectExpression

use of io.crate.execution.engine.collect.CollectExpression in project crate by crate.

the class ProjectionToProjectorVisitor method visitSourceIndexWriterProjection.

@Override
public Projector visitSourceIndexWriterProjection(SourceIndexWriterProjection projection, Context context) {
    InputFactory.Context<CollectExpression<Row, ?>> ctx = inputFactory.ctxForInputColumns(context.txnCtx);
    List<Input<?>> partitionedByInputs = new ArrayList<>(projection.partitionedBySymbols().size());
    for (Symbol partitionedBySymbol : projection.partitionedBySymbols()) {
        partitionedByInputs.add(ctx.add(partitionedBySymbol));
    }
    Input<?> sourceInput = ctx.add(projection.rawSource());
    Supplier<String> indexNameResolver = IndexNameResolver.create(projection.tableIdent(), projection.partitionIdent(), partitionedByInputs);
    ClusterState state = clusterService.state();
    Settings tableSettings = TableSettingsResolver.get(state.getMetadata(), projection.tableIdent(), !projection.partitionedBySymbols().isEmpty());
    int targetTableNumShards = IndexMetadata.INDEX_NUMBER_OF_SHARDS_SETTING.get(tableSettings);
    int targetTableNumReplicas = NumberOfReplicas.fromSettings(tableSettings, state.getNodes().getSize());
    UpsertResultContext upsertResultContext;
    if (projection instanceof SourceIndexWriterReturnSummaryProjection) {
        upsertResultContext = UpsertResultContext.forReturnSummary(context.txnCtx, (SourceIndexWriterReturnSummaryProjection) projection, clusterService.localNode(), inputFactory);
    } else {
        upsertResultContext = UpsertResultContext.forRowCount();
    }
    return new IndexWriterProjector(clusterService, nodeJobsCounter, circuitBreakerService.getBreaker(HierarchyCircuitBreakerService.QUERY), context.ramAccounting, threadPool.scheduler(), threadPool.executor(ThreadPool.Names.SEARCH), context.txnCtx, nodeCtx, state.metadata().settings(), targetTableNumShards, targetTableNumReplicas, transportActionProvider.transportBulkCreateIndicesAction(), transportActionProvider.transportShardUpsertAction()::execute, indexNameResolver, projection.rawSourceReference(), projection.primaryKeys(), projection.ids(), projection.clusteredBy(), projection.clusteredByIdent(), sourceInput, ctx.expressions(), projection.bulkActions(), projection.includes(), projection.excludes(), projection.autoCreateIndices(), projection.overwriteDuplicates(), context.jobId, upsertResultContext, projection.failFast());
}

Example 8 with CollectExpression

use of io.crate.execution.engine.collect.CollectExpression in project crate by crate.

the class InputFactoryTest method testProcessGroupByProjectionSymbolsAggregation.

@Test
public void testProcessGroupByProjectionSymbolsAggregation() throws Exception {
    // select count(x), x, y * 2 ... group by x, y * 2
    // keys: [ in(0), in(1) + 10 ]
    List<Symbol> keys = Arrays.asList(new InputColumn(0, DataTypes.LONG), add);
    Function countX = (Function) expressions.asSymbol("count(x)");
    // values: [ count(in(0)) ]
    List<Aggregation> values = List.of(new Aggregation(countX.signature(), countX.valueType(), List.of(new InputColumn(0))));
    InputFactory.Context<CollectExpression<Row, ?>> ctx = factory.ctxForAggregations(txnCtx);
    ctx.add(keys);
    // inputs: [ x, add ]
    List<Input<?>> keyInputs = ctx.topLevelInputs();
    ctx.add(values);
    List<AggregationContext> aggregations = ctx.aggregations();
    assertThat(aggregations.size(), is(1));
    // collectExpressions: [ in0, in1 ]
    List<CollectExpression<Row, ?>> expressions = new ArrayList<>(ctx.expressions());
    assertThat(expressions.size(), is(2));
    List<Input<?>> allInputs = ctx.topLevelInputs();
    // only 2 because count is no input
    assertThat(allInputs.size(), is(2));
    RowN row = new RowN(1L, 2L);
    for (CollectExpression<Row, ?> expression : expressions) {
        expression.setNextRow(row);
    }
    assertThat(expressions.get(0).value(), is(1L));
    // raw input value
    assertThat(expressions.get(1).value(), is(2L));
    assertThat(keyInputs.size(), is(2));
    assertThat(keyInputs.get(0).value(), is(1L));
    // 2 + 10
    assertThat(keyInputs.get(1).value(), is(12));
}

Example 9 with CollectExpression

use of io.crate.execution.engine.collect.CollectExpression in project crate by crate.

the class InputFactoryTest method testAggregationSymbolsInputReuse.

@Test
public void testAggregationSymbolsInputReuse() throws Exception {
    Function countX = (Function) expressions.asSymbol("count(x)");
    Function avgX = (Function) expressions.asSymbol("avg(x)");
    List<Symbol> aggregations = Arrays.asList(new Aggregation(countX.signature(), countX.signature().getReturnType().createType(), List.of(new InputColumn(0))), new Aggregation(avgX.signature(), avgX.signature().getReturnType().createType(), List.of(new InputColumn(0))));
    InputFactory.Context<CollectExpression<Row, ?>> ctx = factory.ctxForAggregations(txnCtx);
    ctx.add(aggregations);
    List<AggregationContext> aggregationContexts = ctx.aggregations();
    Input<?> inputCount = aggregationContexts.get(0).inputs()[0];
    Input<?> inputAverage = aggregationContexts.get(1).inputs()[0];
    assertSame(inputCount, inputAverage);
}

Example 10 with CollectExpression

use of io.crate.execution.engine.collect.CollectExpression in project crate by crate.

the class MapRowUsingInputsTest method createInputs.

@Before
public void createInputs() throws Exception {
    InputFactory inputFactory = new InputFactory(createNodeContext());
    InputFactory.Context<CollectExpression<Row, ?>> ctx = inputFactory.ctxForInputColumns(txnCtx);
    var addFunction = new Function(Signature.scalar(ArithmeticFunctions.Names.ADD, DataTypes.LONG.getTypeSignature(), DataTypes.LONG.getTypeSignature(), DataTypes.LONG.getTypeSignature()).withFeatures(Scalar.DETERMINISTIC_AND_COMPARISON_REPLACEMENT), List.of(new InputColumn(0, DataTypes.LONG), Literal.of(2L)), DataTypes.LONG);
    inputs = Collections.singletonList(ctx.add(addFunction));
    expressions = ctx.expressions();
}