Examples with ReferenceResolver io.crate.expression.reference.ReferenceResolver used on opensource projects

Example 1 with ReferenceResolver

use of io.crate.expression.reference.ReferenceResolver 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);
}