Examples with PageBucketReceiver io.crate.execution.jobs.PageBucketReceiver used on opensource projects

Example 1 with PageBucketReceiver

use of io.crate.execution.jobs.PageBucketReceiver in project crate by crate.

the class RemoteCollector method createPageDownstreamContext.

private RootTask.Builder createPageDownstreamContext() {
    RootTask.Builder builder = tasksService.newBuilder(jobId, sessionSettings.userName(), localNode, Collections.emptySet());
    PassThroughPagingIterator<Integer, Row> pagingIterator;
    if (scrollRequired) {
        pagingIterator = PassThroughPagingIterator.repeatable();
    } else {
        pagingIterator = PassThroughPagingIterator.oneShot();
    }
    PageBucketReceiver pageBucketReceiver = new CumulativePageBucketReceiver(localNode, RECEIVER_PHASE_ID, executor, DataTypes.getStreamers(collectPhase.outputTypes()), consumer, pagingIterator, 1);
    builder.addTask(new DistResultRXTask(RECEIVER_PHASE_ID, "RemoteCollectPhase", pageBucketReceiver, ramAccounting, 1));
    return builder;
}

Example 2 with PageBucketReceiver

use of io.crate.execution.jobs.PageBucketReceiver in project crate by crate.

the class BucketForwarder method setBuckets.

protected void setBuckets(List<? extends Bucket> buckets) {
    initializationTracker.jobInitialized();
    for (int i = 0; i < pageBucketReceivers.size(); i++) {
        PageBucketReceiver pageBucketReceiver = pageBucketReceivers.get(i);
        Bucket bucket = buckets.get(i);
        assert bucket != null : "buckets must contain a non-null bucket at idx=" + i;
        pageBucketReceiver.setBucket(bucketIdx, bucket, true, PagingUnsupportedResultListener.INSTANCE);
    }
}

Example 3 with PageBucketReceiver

use of io.crate.execution.jobs.PageBucketReceiver in project crate by crate.

the class JobLauncher method getHandlerBucketReceivers.

private List<PageBucketReceiver> getHandlerBucketReceivers(RootTask rootTask, List<Tuple<ExecutionPhase, RowConsumer>> handlerPhases) {
    final List<PageBucketReceiver> pageBucketReceivers = new ArrayList<>(handlerPhases.size());
    for (Tuple<ExecutionPhase, ?> handlerPhase : handlerPhases) {
        Task ctx = rootTask.getTaskOrNull(handlerPhase.v1().phaseId());
        if (ctx instanceof DownstreamRXTask) {
            PageBucketReceiver pageBucketReceiver = ((DownstreamRXTask) ctx).getBucketReceiver((byte) 0);
            pageBucketReceivers.add(pageBucketReceiver);
        }
    }
    return pageBucketReceivers;
}

Example 4 with PageBucketReceiver

use of io.crate.execution.jobs.PageBucketReceiver in project crate by crate.

the class DistributingConsumerTest method createFakeTransport.

private TransportDistributedResultAction createFakeTransport(Streamer<?>[] streamers, DistResultRXTask distResultRXTask) {
    TransportDistributedResultAction distributedResultAction = mock(TransportDistributedResultAction.class);
    doAnswer((InvocationOnMock invocationOnMock) -> {
        Object[] args = invocationOnMock.getArguments();
        DistributedResultRequest resultRequest = (DistributedResultRequest) args[1];
        ActionListener<DistributedResultResponse> listener = (ActionListener<DistributedResultResponse>) args[2];
        Throwable throwable = resultRequest.throwable();
        PageBucketReceiver bucketReceiver = distResultRXTask.getBucketReceiver(resultRequest.executionPhaseInputId());
        assertThat(bucketReceiver, Matchers.notNullValue());
        if (throwable == null) {
            bucketReceiver.setBucket(resultRequest.bucketIdx(), resultRequest.readRows(streamers), resultRequest.isLast(), needMore -> listener.onResponse(new DistributedResultResponse(needMore)));
        } else {
            bucketReceiver.kill(throwable);
        }
        return null;
    }).when(distributedResultAction).pushResult(anyString(), any(), any());
    return distributedResultAction;
}

Example 5 with PageBucketReceiver

use of io.crate.execution.jobs.PageBucketReceiver in project crate by crate.

the class TransportDistributedResultAction method nodeOperation.

private CompletableFuture<DistributedResultResponse> nodeOperation(final DistributedResultRequest request, @Nullable Iterator<TimeValue> retryDelay) {
    RootTask rootTask = tasksService.getTaskOrNull(request.jobId());
    if (rootTask == null) {
        if (tasksService.recentlyFailed(request.jobId())) {
            return CompletableFuture.failedFuture(JobKilledException.of("Received result for job=" + request.jobId() + " but there is no context for this job due to a failure during the setup."));
        } else {
            return retryOrFailureResponse(request, retryDelay);
        }
    }
    DownstreamRXTask rxTask;
    try {
        rxTask = rootTask.getTask(request.executionPhaseId());
    } catch (ClassCastException e) {
        return CompletableFuture.failedFuture(new IllegalStateException(String.format(Locale.ENGLISH, "Found execution rootTask for %d but it's not a downstream rootTask", request.executionPhaseId()), e));
    } catch (Throwable t) {
        return CompletableFuture.failedFuture(t);
    }
    PageBucketReceiver pageBucketReceiver = rxTask.getBucketReceiver(request.executionPhaseInputId());
    if (pageBucketReceiver == null) {
        return CompletableFuture.failedFuture(new IllegalStateException(String.format(Locale.ENGLISH, "Couldn't find BucketReceiver for input %d", request.executionPhaseInputId())));
    }
    Throwable throwable = request.throwable();
    if (throwable == null) {
        SendResponsePageResultListener pageResultListener = new SendResponsePageResultListener();
        pageBucketReceiver.setBucket(request.bucketIdx(), request.readRows(pageBucketReceiver.streamers()), request.isLast(), pageResultListener);
        return pageResultListener.future;
    } else {
        pageBucketReceiver.kill(throwable);
        return CompletableFuture.completedFuture(new DistributedResultResponse(false));
    }
}