Examples with BatchIterator io.crate.data.BatchIterator used on opensource projects

Example 21 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class BatchPagingIteratorTest method testBatchPagingIteratorWithPagedSource.

@Test
public void testBatchPagingIteratorWithPagedSource() throws Exception {
    List<Object[]> expectedResult = StreamSupport.stream(RowGenerator.range(0, 10).spliterator(), false).map(Row::materialize).collect(Collectors.toList());
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> {
        BatchSimulatingIterator<Row> source = new BatchSimulatingIterator<>(TestingBatchIterators.range(0, 10), 2, 5, executor);
        Function<Integer, KillableCompletionStage<? extends Iterable<? extends KeyIterable<Integer, Row>>>> fetchMore = exhausted -> {
            List<Row> rows = new ArrayList<>();
            while (source.moveNext()) {
                rows.add(new RowN(source.currentElement().materialize()));
            }
            if (source.allLoaded()) {
                return KillableCompletionStage.whenKilled(CompletableFuture.completedFuture(singletonList(new KeyIterable<>(1, rows))), t -> {
                });
            }
            // but to simulate multiple pages and fetchMore calls
            try {
                return KillableCompletionStage.whenKilled(source.loadNextBatch().toCompletableFuture().thenApply(ignored -> {
                    while (source.moveNext()) {
                        rows.add(new RowN(source.currentElement().materialize()));
                    }
                    return singleton(new KeyIterable<>(1, rows));
                }), t -> {
                });
            } catch (Exception e) {
                return KillableCompletionStage.failed(e);
            }
        };
        return new BatchPagingIterator<>(PassThroughPagingIterator.repeatable(), fetchMore, source::allLoaded, throwable -> {
            if (throwable == null) {
                source.close();
            } else {
                source.kill(throwable);
            }
        });
    };
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(expectedResult);
}

Example 22 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class NodeStatsTest method testNoRequestIfNotRequired.

@Test
public void testNoRequestIfNotRequired() throws Exception {
    List<Symbol> toCollect = new ArrayList<>();
    toCollect.add(idRef);
    when(collectPhase.toCollect()).thenReturn(toCollect);
    BatchIterator iterator = NodeStats.newInstance(transportNodeStatsAction, collectPhase, nodes, txnCtx, new InputFactory(nodeCtx));
    iterator.loadNextBatch();
    // Because only id and name are selected, transportNodesStatsAction should be never used
    verifyNoMoreInteractions(transportNodeStatsAction);
}

Example 23 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class NodeStatsTest method testRequestsAreIssued.

@Test
public void testRequestsAreIssued() throws Exception {
    List<Symbol> toCollect = new ArrayList<>();
    toCollect.add(idRef);
    toCollect.add(nameRef);
    when(collectPhase.toCollect()).thenReturn(toCollect);
    BatchIterator iterator = NodeStats.newInstance(transportNodeStatsAction, collectPhase, nodes, txnCtx, new InputFactory(nodeCtx));
    iterator.loadNextBatch();
    verifyNoMoreInteractions(transportNodeStatsAction);
}

Example 24 with BatchIterator

use of io.crate.data.BatchIterator 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);
}

Example 25 with BatchIterator

use of io.crate.data.BatchIterator in project crate by crate.

the class FileCollectSource method getCollector.

@Override
public CrateCollector getCollector(CollectPhase collectPhase, BatchConsumer consumer, JobCollectContext jobCollectContext) {
    FileUriCollectPhase fileUriCollectPhase = (FileUriCollectPhase) collectPhase;
    InputFactory.Context<LineCollectorExpression<?>> ctx = inputFactory.ctxForRefs(FileLineReferenceResolver::getImplementation);
    ctx.add(collectPhase.toCollect());
    String[] readers = fileUriCollectPhase.nodeIds().toArray(new String[fileUriCollectPhase.nodeIds().size()]);
    Arrays.sort(readers);
    List<String> fileUris;
    fileUris = targetUriToStringList(fileUriCollectPhase.targetUri());
    BatchIterator fileReadingIterator = FileReadingIterator.newInstance(fileUris, ctx.topLevelInputs(), ctx.expressions(), fileUriCollectPhase.compression(), fileInputFactoryMap, fileUriCollectPhase.sharedStorage(), readers.length, Arrays.binarySearch(readers, clusterService.state().nodes().getLocalNodeId()));
    return BatchIteratorCollectorBridge.newInstance(fileReadingIterator, consumer);
}