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

Example 41 with BatchIterator

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

the class RemoteCollectorFactory method createCollector.

/**
 * create a RemoteCollector
 * The RemoteCollector will collect data from another node using a wormhole as if it was collecting on this node.
 * <p>
 * This should only be used if a shard is not available on the current node due to a relocation
 */
public CompletableFuture<BatchIterator<Row>> createCollector(ShardId shardId, RoutedCollectPhase collectPhase, CollectTask collectTask, ShardCollectorProviderFactory shardCollectorProviderFactory) {
    ShardStateObserver shardStateObserver = new ShardStateObserver(clusterService);
    CompletableFuture<ShardRouting> shardBecameActive = shardStateObserver.waitForActiveShard(shardId);
    Runnable onClose = () -> {
    };
    Consumer<Throwable> kill = killReason -> {
        shardBecameActive.cancel(true);
        shardBecameActive.completeExceptionally(killReason);
    };
    return shardBecameActive.thenApply(activePrimaryRouting -> CollectingBatchIterator.newInstance(onClose, kill, () -> retrieveRows(activePrimaryRouting, collectPhase, collectTask, shardCollectorProviderFactory), true));
}

Example 42 with BatchIterator

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

the class SystemCollectSource method getIterator.

@Override
public CompletableFuture<BatchIterator<Row>> getIterator(TransactionContext txnCtx, CollectPhase phase, CollectTask collectTask, boolean supportMoveToStart) {
    RoutedCollectPhase collectPhase = (RoutedCollectPhase) phase;
    Map<String, Map<String, IntIndexedContainer>> locations = collectPhase.routing().locations();
    String table = Iterables.getOnlyElement(locations.get(clusterService.localNode().getId()).keySet());
    RelationName relationName = RelationName.fromIndexName(table);
    StaticTableDefinition<?> tableDefinition = tableDefinition(relationName);
    User user = requireNonNull(userLookup.findUser(txnCtx.sessionSettings().userName()), "User who invoked a statement must exist");
    return CompletableFuture.completedFuture(CollectingBatchIterator.newInstance(() -> {
    }, // If data is already local, then `CollectingBatchIterator` takes care of kill handling.
    t -> {
    }, () -> tableDefinition.retrieveRecords(txnCtx, user).thenApply(records -> recordsToRows(collectPhase, collectTask.txnCtx(), tableDefinition.getReferenceResolver(), supportMoveToStart, records)), tableDefinition.involvesIO()));
}

Example 43 with BatchIterator

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

the class BatchIteratorBackpressureExecutorTest method testPauseOnFirstBatch.

@Test
public void testPauseOnFirstBatch() throws Exception {
    BatchIterator<Integer> numbersBi = InMemoryBatchIterator.of(() -> IntStream.range(0, 5).iterator(), -1, true);
    BatchSimulatingIterator<Integer> it = new BatchSimulatingIterator<>(numbersBi, 2, 5, executor);
    AtomicInteger numRows = new AtomicInteger(0);
    AtomicInteger numPauses = new AtomicInteger(0);
    Predicate<Integer> shouldPause = i -> {
        if (i == 0 && numPauses.get() == 0) {
            numPauses.incrementAndGet();
            return true;
        }
        return false;
    };
    BatchIteratorBackpressureExecutor<Integer, Integer> executor = new BatchIteratorBackpressureExecutor<>(UUID.randomUUID(), scheduler, this.executor, it, i -> CompletableFuture.supplyAsync(numRows::incrementAndGet, this.executor), (a, b) -> a + b, 0, shouldPause, null, null, ignored -> 1L);
    CompletableFuture<Integer> result = executor.consumeIteratorAndExecute();
    result.get(10, TimeUnit.SECONDS);
    assertThat(numPauses.get(), Matchers.is(1));
}

Example 44 with BatchIterator

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

the class IndexWriterProjectorTest method testIndexWriter.

@Test
public void testIndexWriter() throws Throwable {
    execute("create table bulk_import (id int primary key, name string) with (number_of_replicas=0)");
    ensureGreen();
    InputCollectExpression sourceInput = new InputCollectExpression(1);
    List<CollectExpression<Row, ?>> collectExpressions = Collections.<CollectExpression<Row, ?>>singletonList(sourceInput);
    RelationName bulkImportIdent = new RelationName(sqlExecutor.getCurrentSchema(), "bulk_import");
    ClusterState state = clusterService().state();
    Settings tableSettings = TableSettingsResolver.get(state.getMetadata(), bulkImportIdent, false);
    ThreadPool threadPool = internalCluster().getInstance(ThreadPool.class);
    IndexWriterProjector writerProjector = new IndexWriterProjector(clusterService(), new NodeLimits(new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)), new NoopCircuitBreaker("dummy"), RamAccounting.NO_ACCOUNTING, threadPool.scheduler(), threadPool.executor(ThreadPool.Names.SEARCH), CoordinatorTxnCtx.systemTransactionContext(), new NodeContext(internalCluster().getInstance(Functions.class)), Settings.EMPTY, IndexMetadata.INDEX_NUMBER_OF_SHARDS_SETTING.get(tableSettings), NumberOfReplicas.fromSettings(tableSettings, state.getNodes().getSize()), internalCluster().getInstance(TransportCreatePartitionsAction.class), internalCluster().getInstance(TransportShardUpsertAction.class)::execute, IndexNameResolver.forTable(bulkImportIdent), new Reference(new ReferenceIdent(bulkImportIdent, DocSysColumns.RAW), RowGranularity.DOC, DataTypes.STRING, 0, null), Collections.singletonList(ID_IDENT), Collections.<Symbol>singletonList(new InputColumn(0)), null, null, sourceInput, collectExpressions, 20, null, null, false, false, UUID.randomUUID(), UpsertResultContext.forRowCount(), false);
    BatchIterator rowsIterator = InMemoryBatchIterator.of(IntStream.range(0, 100).mapToObj(i -> new RowN(new Object[] { i, "{\"id\": " + i + ", \"name\": \"Arthur\"}" })).collect(Collectors.toList()), SENTINEL, true);
    TestingRowConsumer consumer = new TestingRowConsumer();
    consumer.accept(writerProjector.apply(rowsIterator), null);
    Bucket objects = consumer.getBucket();
    assertThat(objects, contains(isRow(100L)));
    execute("refresh table bulk_import");
    execute("select count(*) from bulk_import");
    assertThat(response.rowCount(), is(1L));
    assertThat(response.rows()[0][0], is(100L));
}

Example 45 with BatchIterator

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

the class HashInnerJoinBatchIteratorTest method testInnerHashJoinWithBlockSizeSmallerThanDataSet.

@Test
public void testInnerHashJoinWithBlockSizeSmallerThanDataSet() throws Exception {
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> new HashInnerJoinBatchIterator(leftIterator.get(), rightIterator.get(), mock(RowAccounting.class), new CombinedRow(1, 1), getCol0EqCol1JoinCondition(), getHashForLeft(), getHashForRight(), () -> 1);
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(expectedResult);
}