Examples with ExchangeSink io.trino.spi.exchange.ExchangeSink used on opensource projects

Example 1 with ExchangeSink

use of io.trino.spi.exchange.ExchangeSink in project trino by trinodb.

the class SpoolingExchangeOutputBuffer method setNoMorePages.

@Override
public void setNoMorePages() {
    if (!stateMachine.noMorePages()) {
        return;
    }
    ExchangeSink sink = exchangeSink;
    if (sink == null) {
        // abort might've released the sink in a meantime
        return;
    }
    sink.finish().whenComplete((value, failure) -> {
        if (failure != null) {
            stateMachine.fail(failure);
        } else {
            stateMachine.finish();
        }
        exchangeSink = null;
        updateMemoryUsage(0);
    });
}

Example 2 with ExchangeSink

use of io.trino.spi.exchange.ExchangeSink in project trino by trinodb.

the class SpoolingExchangeOutputBuffer method enqueue.

@Override
public void enqueue(int partition, List<Slice> pages) {
    requireNonNull(pages, "pages is null");
    // this can happen with a limit query
    if (!stateMachine.getState().canAddPages()) {
        return;
    }
    ExchangeSink sink = exchangeSink;
    checkState(sink != null, "exchangeSink is null");
    for (Slice page : pages) {
        sink.add(partition, page);
        totalRowsAdded.addAndGet(getSerializedPagePositionCount(page));
    }
    updateMemoryUsage(sink.getMemoryUsage());
    totalPagesAdded.addAndGet(pages.size());
}

Example 3 with ExchangeSink

use of io.trino.spi.exchange.ExchangeSink in project trino by trinodb.

the class LazyOutputBuffer method setOutputBuffers.

@Override
public void setOutputBuffers(OutputBuffers newOutputBuffers) {
    Set<OutputBufferId> destroyedBuffers = ImmutableSet.of();
    List<PendingRead> pendingReads = ImmutableList.of();
    OutputBuffer outputBuffer = delegate;
    if (outputBuffer == null) {
        synchronized (this) {
            outputBuffer = delegate;
            if (outputBuffer == null) {
                // ignore set output if buffer was already destroyed or failed
                if (stateMachine.getState().isTerminal()) {
                    return;
                }
                switch(newOutputBuffers.getType()) {
                    case PARTITIONED:
                        outputBuffer = new PartitionedOutputBuffer(taskInstanceId, stateMachine, newOutputBuffers, maxBufferSize, memoryContextSupplier, executor);
                        break;
                    case BROADCAST:
                        outputBuffer = new BroadcastOutputBuffer(taskInstanceId, stateMachine, maxBroadcastBufferSize, memoryContextSupplier, executor, notifyStatusChanged);
                        break;
                    case ARBITRARY:
                        outputBuffer = new ArbitraryOutputBuffer(taskInstanceId, stateMachine, maxBufferSize, memoryContextSupplier, executor);
                        break;
                    case SPOOL:
                        ExchangeSinkInstanceHandle exchangeSinkInstanceHandle = newOutputBuffers.getExchangeSinkInstanceHandle().orElseThrow(() -> new IllegalArgumentException("exchange sink handle is expected to be present for buffer type EXTERNAL"));
                        ExchangeManager exchangeManager = exchangeManagerRegistry.getExchangeManager();
                        ExchangeSink exchangeSink = exchangeManager.createSink(exchangeSinkInstanceHandle, false);
                        outputBuffer = new SpoolingExchangeOutputBuffer(stateMachine, newOutputBuffers, exchangeSink, memoryContextSupplier);
                        break;
                    default:
                        throw new IllegalArgumentException("Unexpected output buffer type: " + newOutputBuffers.getType());
                }
                // process pending aborts and reads outside of synchronized lock
                destroyedBuffers = ImmutableSet.copyOf(this.destroyedBuffers);
                this.destroyedBuffers.clear();
                pendingReads = ImmutableList.copyOf(this.pendingReads);
                this.pendingReads.clear();
                // Must be assigned last to avoid a race condition with unsynchronized readers
                delegate = outputBuffer;
            }
        }
    }
    outputBuffer.setOutputBuffers(newOutputBuffers);
    // process pending aborts and reads outside of synchronized lock
    destroyedBuffers.forEach(outputBuffer::destroy);
    for (PendingRead pendingRead : pendingReads) {
        pendingRead.process(outputBuffer);
    }
}

Example 4 with ExchangeSink

use of io.trino.spi.exchange.ExchangeSink in project trino by trinodb.

the class SpoolingExchangeOutputBuffer method abort.

@Override
public void abort() {
    if (!stateMachine.abort()) {
        return;
    }
    ExchangeSink sink = exchangeSink;
    if (sink == null) {
        return;
    }
    sink.abort().whenComplete((value, failure) -> {
        if (failure != null) {
            log.warn(failure, "Error aborting exchange sink");
        }
        exchangeSink = null;
        updateMemoryUsage(0);
    });
}

Example 5 with ExchangeSink

use of io.trino.spi.exchange.ExchangeSink in project trino by trinodb.

the class AbstractTestExchangeManager method writeData.

private void writeData(ExchangeSinkInstanceHandle handle, Multimap<Integer, String> data, boolean finish) {
    ExchangeSink sink = exchangeManager.createSink(handle, false);
    data.forEach((key, value) -> {
        sink.add(key, Slices.utf8Slice(value));
    });
    if (finish) {
        getFutureValue(sink.finish());
    } else {
        getFutureValue(sink.abort());
    }
}