Examples with MessageMetadataResult org.apache.arrow.vector.ipc.message.MessageMetadataResult used on opensource projects

Example 1 with MessageMetadataResult

use of org.apache.arrow.vector.ipc.message.MessageMetadataResult in project flink by apache.

the class ArrowUtils method readNextBatch.

private static byte[] readNextBatch(ReadableByteChannel channel) throws IOException {
    MessageMetadataResult metadata = MessageSerializer.readMessage(new ReadChannel(channel));
    if (metadata == null) {
        return null;
    }
    long bodyLength = metadata.getMessageBodyLength();
    // Only care about RecordBatch messages and skip the other kind of messages
    if (metadata.getMessage().headerType() == MessageHeader.RecordBatch) {
        // Buffer backed output large enough to hold 8-byte length + complete serialized message
        ByteArrayOutputStreamWithPos baos = new ByteArrayOutputStreamWithPos((int) (8 + metadata.getMessageLength() + bodyLength));
        // Write message metadata to ByteBuffer output stream
        MessageSerializer.writeMessageBuffer(new WriteChannel(Channels.newChannel(baos)), metadata.getMessageLength(), metadata.getMessageBuffer());
        baos.close();
        ByteBuffer result = ByteBuffer.wrap(baos.getBuf());
        result.position(baos.getPosition());
        result.limit(result.capacity());
        readFully(channel, result);
        return result.array();
    } else {
        if (bodyLength > 0) {
            // Skip message body if not a RecordBatch
            Channels.newInputStream(channel).skip(bodyLength);
        }
        // Proceed to next message
        return readNextBatch(channel);
    }
}