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

Example 1 with CollectTask

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

the class RootTaskTest method testFailureClosesAllSubContexts.

@Test
public void testFailureClosesAllSubContexts() throws Throwable {
    String localNodeId = "localNodeId";
    RoutedCollectPhase collectPhase = Mockito.mock(RoutedCollectPhase.class);
    Routing routing = Mockito.mock(Routing.class);
    when(routing.containsShards(localNodeId)).thenReturn(false);
    when(collectPhase.phaseId()).thenReturn(1);
    when(collectPhase.routing()).thenReturn(routing);
    when(collectPhase.maxRowGranularity()).thenReturn(RowGranularity.DOC);
    RootTask.Builder builder = new RootTask.Builder(logger, UUID.randomUUID(), "dummy-user", coordinatorNode, Collections.emptySet(), mock(JobsLogs.class));
    CollectTask collectChildTask = new CollectTask(collectPhase, CoordinatorTxnCtx.systemTransactionContext(), mock(MapSideDataCollectOperation.class), RamAccounting.NO_ACCOUNTING, ramAccounting -> new OnHeapMemoryManager(ramAccounting::addBytes), new TestingRowConsumer(), mock(SharedShardContexts.class), Version.CURRENT, 4096);
    TestingRowConsumer batchConsumer = new TestingRowConsumer();
    PageBucketReceiver pageBucketReceiver = new CumulativePageBucketReceiver("n1", 2, Runnable::run, new Streamer[] { IntegerType.INSTANCE.streamer() }, batchConsumer, PassThroughPagingIterator.oneShot(), 1);
    DistResultRXTask distResultRXTask = spy(new DistResultRXTask(2, "dummy", pageBucketReceiver, RamAccounting.NO_ACCOUNTING, 1));
    builder.addTask(collectChildTask);
    builder.addTask(distResultRXTask);
    RootTask rootTask = builder.build();
    Exception failure = new Exception("failure!");
    collectChildTask.kill(failure);
    // other contexts must be killed with same failure
    verify(distResultRXTask, times(1)).kill(failure);
    assertThat(rootTask.getTask(1).completionFuture().isDone(), is(true));
    assertThat(rootTask.getTask(2).completionFuture().isDone(), is(true));
}

Example 2 with CollectTask

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

the class ShardCollectSource method getIterator.

@Override
public CompletableFuture<BatchIterator<Row>> getIterator(TransactionContext txnCtx, CollectPhase phase, CollectTask collectTask, boolean supportMoveToStart) {
    RoutedCollectPhase collectPhase = (RoutedCollectPhase) phase;
    String localNodeId = clusterService.localNode().getId();
    Projectors projectors = new Projectors(collectPhase.projections(), collectPhase.jobId(), collectTask.txnCtx(), collectTask.getRamAccounting(), collectTask.memoryManager(), sharedProjectorFactory);
    boolean requireMoveToStartSupport = supportMoveToStart && !projectors.providesIndependentScroll();
    if (collectPhase.maxRowGranularity() == RowGranularity.SHARD) {
        return CompletableFuture.completedFuture(projectors.wrap(InMemoryBatchIterator.of(getShardsIterator(collectTask.txnCtx(), collectPhase, localNodeId), SentinelRow.SENTINEL, true)));
    }
    OrderBy orderBy = collectPhase.orderBy();
    if (collectPhase.maxRowGranularity() == RowGranularity.DOC && orderBy != null) {
        return createMultiShardScoreDocCollector(collectPhase, requireMoveToStartSupport, collectTask, localNodeId).thenApply(projectors::wrap);
    }
    boolean hasShardProjections = Projections.hasAnyShardProjections(collectPhase.projections());
    Map<String, IntIndexedContainer> indexShards = collectPhase.routing().locations().get(localNodeId);
    List<CompletableFuture<BatchIterator<Row>>> iterators = indexShards == null ? Collections.emptyList() : getIterators(collectTask, collectPhase, requireMoveToStartSupport, indexShards);
    final CompletableFuture<BatchIterator<Row>> result;
    switch(iterators.size()) {
        case 0:
            result = CompletableFuture.completedFuture(InMemoryBatchIterator.empty(SentinelRow.SENTINEL));
            break;
        case 1:
            result = iterators.get(0);
            break;
        default:
            if (hasShardProjections) {
                // use AsyncCompositeBatchIterator for multi-threaded loadNextBatch
                // in order to process shard-based projections concurrently
                // noinspection unchecked
                result = CompletableFutures.allAsList(iterators).thenApply(its -> CompositeBatchIterator.asyncComposite(executor, availableThreads, its.toArray(new BatchIterator[0])));
            } else {
                // noinspection unchecked
                result = CompletableFutures.allAsList(iterators).thenApply(its -> CompositeBatchIterator.seqComposite(its.toArray(new BatchIterator[0])));
            }
    }
    return result.thenApply(it -> projectors.wrap(it));
}

Example 3 with CollectTask

use of io.crate.execution.engine.collect.CollectTask 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()));
}