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

Example 31 with BatchIterator

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

the class NodeStatsIteratorTest method testNoRequestIfNotRequired.

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

Example 32 with BatchIterator

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

the class NodeStatsIteratorTest method testRequestsIfRequired.

@Test
public void testRequestsIfRequired() throws InterruptedException, ExecutionException, TimeoutException {
    List<Symbol> toCollect = new ArrayList<>();
    toCollect.add(idRef);
    toCollect.add(hostnameRef);
    when(collectPhase.toCollect()).thenReturn(toCollect);
    BatchIterator iterator = NodeStatsIterator.newInstance(transportNodeStatsAction, collectPhase, nodes, new InputFactory(getFunctions()));
    iterator.loadNextBatch();
    // Hostnames needs to be collected so requests need to be performed
    verify(transportNodeStatsAction).execute(eq("nodeOne"), any(NodeStatsRequest.class), any(ActionListener.class), eq(TimeValue.timeValueMillis(3000L)));
    verify(transportNodeStatsAction).execute(eq("nodeTwo"), any(NodeStatsRequest.class), any(ActionListener.class), eq(TimeValue.timeValueMillis(3000L)));
    verifyNoMoreInteractions(transportNodeStatsAction);
}

Example 33 with BatchIterator

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

the class ShardingUpsertExecutor method apply.

@Override
public CompletableFuture<? extends Iterable<Row>> apply(BatchIterator<Row> batchIterator) {
    final ConcurrencyLimit nodeLimit = nodeLimits.get(localNode);
    long startTime = nodeLimit.startSample();
    var isUsedBytesOverThreshold = new IsUsedBytesOverThreshold(queryCircuitBreaker, nodeLimit);
    var reqBatchIterator = BatchIterators.partition(batchIterator, bulkSize, () -> new ShardedRequests<>(requestFactory, ramAccounting), grouper, bulkShardCreationLimiter.or(isUsedBytesOverThreshold));
    // If IO is involved the source iterator should pause when the target node reaches a concurrent job counter limit.
    // Without IO, we assume that the source iterates over in-memory structures which should be processed as
    // fast as possible to free resources.
    Predicate<ShardedRequests<ShardUpsertRequest, ShardUpsertRequest.Item>> shouldPause = this::shouldPauseOnPartitionCreation;
    if (batchIterator.hasLazyResultSet()) {
        shouldPause = shouldPause.or(this::shouldPauseOnTargetNodeJobsCounter).or(isUsedBytesOverThreshold);
    }
    BatchIteratorBackpressureExecutor<ShardedRequests<ShardUpsertRequest, ShardUpsertRequest.Item>, UpsertResults> executor = new BatchIteratorBackpressureExecutor<>(jobId, scheduler, this.executor, reqBatchIterator, this::execute, resultCollector.combiner(), resultCollector.supplier().get(), shouldPause, earlyTerminationCondition, earlyTerminationExceptionGenerator, this::getMaxLastRttInMs);
    return executor.consumeIteratorAndExecute().thenApply(upsertResults -> resultCollector.finisher().apply(upsertResults)).whenComplete((res, err) -> {
        nodeLimit.onSample(startTime, err != null);
    });
}

Example 34 with BatchIterator

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

the class NodeStatsTest method testRequestsIfRequired.

@Test
public void testRequestsIfRequired() throws Exception {
    List<Symbol> toCollect = new ArrayList<>();
    toCollect.add(idRef);
    toCollect.add(hostnameRef);
    when(collectPhase.toCollect()).thenReturn(toCollect);
    BatchIterator iterator = NodeStats.newInstance(transportNodeStatsAction, collectPhase, nodes, txnCtx, new InputFactory(nodeCtx));
    iterator.loadNextBatch();
    // Hostnames needs to be collected so requests need to be performed
    verify(transportNodeStatsAction).execute(eq("nodeOne"), any(NodeStatsRequest.class), any(ActionListener.class), eq(TimeValue.timeValueMillis(3000L)));
    verify(transportNodeStatsAction).execute(eq("nodeTwo"), any(NodeStatsRequest.class), any(ActionListener.class), eq(TimeValue.timeValueMillis(3000L)));
    verifyNoMoreInteractions(transportNodeStatsAction);
}

Example 35 with BatchIterator

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

the class FileReadingIteratorTest method testIteratorContract_givenDefaultJsonInputFormat_AndCSVExtension_thenWritesAsMap.

@Test
public void testIteratorContract_givenDefaultJsonInputFormat_AndCSVExtension_thenWritesAsMap() throws Exception {
    tempFilePath = createTempFile("tempfile", ".csv");
    tmpFile = tempFilePath.toFile();
    try (OutputStreamWriter writer = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)) {
        writer.write("name,id,age\n");
        writer.write("Arthur,4,38\n");
        writer.write("Trillian,5,33\n");
    }
    fileUri = tempFilePath.toUri().toString();
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> createBatchIterator(Collections.singletonList(fileUri), JSON);
    List<Object[]> expectedResult = Arrays.asList(new Object[] { CSV_AS_MAP_FIRST_LINE }, new Object[] { CSV_AS_MAP_SECOND_LINE });
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(expectedResult);
}