Examples with MemorySegment org.apache.flink.core.memory.MemorySegment used on opensource projects

Example 46 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class SpillingBuffer method nextSegment.

@Override
protected MemorySegment nextSegment(MemorySegment current, int positionInCurrent) throws IOException {
    // check if we are still in memory
    if (this.writer == null) {
        this.fullSegments.add(current);
        final MemorySegment nextSeg = this.memorySource.nextSegment();
        if (nextSeg != null) {
            return nextSeg;
        } else {
            // out of memory, need to spill: create a writer
            this.writer = this.ioManager.createBlockChannelWriter(this.ioManager.createChannel());
            // add all segments to the writer
            this.blockCount = this.fullSegments.size();
            this.numMemorySegmentsInWriter = this.blockCount;
            for (int i = 0; i < this.fullSegments.size(); i++) {
                this.writer.writeBlock(this.fullSegments.get(i));
            }
            this.fullSegments.clear();
            final MemorySegment seg = this.writer.getNextReturnedBlock();
            this.numMemorySegmentsInWriter--;
            return seg;
        }
    } else {
        // spilling
        this.writer.writeBlock(current);
        this.blockCount++;
        return this.writer.getNextReturnedBlock();
    }
}

Example 47 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class ChannelWriterOutputView method close.

// --------------------------------------------------------------------------------------------
/**
 * Closes this OutputView, closing the underlying writer and returning all memory segments.
 *
 * @return A list containing all memory segments originally supplied to this view.
 * @throws IOException Thrown, if the underlying writer could not be properly closed.
 */
public List<MemorySegment> close() throws IOException {
    // send off set last segment
    writeSegment(getCurrentSegment(), getCurrentPositionInSegment(), true);
    clear();
    // close the writer and gather all segments
    final LinkedBlockingQueue<MemorySegment> queue = this.writer.getReturnQueue();
    this.writer.close();
    // re-collect all memory segments
    ArrayList<MemorySegment> list = new ArrayList<MemorySegment>(this.numSegments);
    for (int i = 0; i < this.numSegments; i++) {
        final MemorySegment m = queue.poll();
        if (m == null) {
            // properly closed.
            throw new RuntimeException("ChannelWriterOutputView: MemorySegments have been taken from return queue by different actor.");
        }
        list.add(m);
    }
    return list;
}

Example 48 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class BufferDecompressor method decompressToOriginalBuffer.

/**
 * The difference between this method and {@link #decompressToIntermediateBuffer(Buffer)} is
 * that this method copies the decompressed data to the input {@link Buffer} starting from
 * offset 0.
 *
 * <p>The caller must guarantee that the input {@link Buffer} is writable and there's enough
 * space left.
 */
@VisibleForTesting
public Buffer decompressToOriginalBuffer(Buffer buffer) {
    int decompressedLen = decompress(buffer);
    // copy the decompressed data back
    int memorySegmentOffset = buffer.getMemorySegmentOffset();
    MemorySegment segment = buffer.getMemorySegment();
    segment.put(memorySegmentOffset, internalBuffer.array(), 0, decompressedLen);
    return new ReadOnlySlicedNetworkBuffer(buffer.asByteBuf(), 0, decompressedLen, memorySegmentOffset, false);
}

Example 49 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class NonSpanningWrapper method getUnconsumedSegment.

CloseableIterator<Buffer> getUnconsumedSegment() {
    if (!hasRemaining()) {
        return CloseableIterator.empty();
    }
    MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(remaining());
    this.segment.copyTo(position, segment, 0, remaining());
    return singleBufferIterator(segment);
}

Example 50 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class BufferReaderWriterUtil method sliceNextBuffer.

@Nullable
static Buffer sliceNextBuffer(ByteBuffer memory) {
    final int remaining = memory.remaining();
    // buffer underflow exceptions which will cause the read to fail.
    if (remaining == 0) {
        return null;
    }
    final boolean isEvent = memory.getShort() == HEADER_VALUE_IS_EVENT;
    final boolean isCompressed = memory.getShort() == BUFFER_IS_COMPRESSED;
    final int size = memory.getInt();
    memory.limit(memory.position() + size);
    ByteBuffer buf = memory.slice();
    memory.position(memory.limit());
    memory.limit(memory.capacity());
    MemorySegment memorySegment = MemorySegmentFactory.wrapOffHeapMemory(buf);
    Buffer.DataType dataType = isEvent ? Buffer.DataType.EVENT_BUFFER : Buffer.DataType.DATA_BUFFER;
    return new NetworkBuffer(memorySegment, FreeingBufferRecycler.INSTANCE, dataType, isCompressed, size);
}