/** * Constructs a streaming reader from a ReadableByteChannel input. Non-blocking. * * @param in ReadableByteChannel to read messages from * @param allocator to allocate new buffers */ public ArrowStreamReader(ReadableByteChannel in, BufferAllocator allocator) { this(new MessageChannelReader(new ReadChannel(in), allocator), allocator); }
/** * Get the number of bytes read from the ReadChannel. * * @return number of bytes */ public long bytesRead() { return in.bytesRead(); }
/** * Close the ReadChannel. * * @throws IOException on error */ @Override public void close() throws IOException { in.close(); } }
/** * Read a Message body from the in channel into an ArrowBuf. * * @param in ReadChannel to read message body from * @param bodyLength Length in bytes of the message body to read * @param allocator Allocate the ArrowBuf to contain message body data * @return an ArrowBuf containing the message body data * @throws IOException on error */ public static ArrowBuf readMessageBody(ReadChannel in, int bodyLength, BufferAllocator allocator) throws IOException { ArrowBuf bodyBuffer = allocator.buffer(bodyLength); if (in.readFully(bodyBuffer, bodyLength) != bodyLength) { throw new IOException("Unexpected end of input trying to read batch."); } return bodyBuffer; } }
/** * Reads up to len into buffer. Returns bytes read. * * @param buffer the buffer to read to * @param l the amount of bytes to read * @return the number of bytes read * @throws IOException if nit enough bytes left to read */ public int readFully(ArrowBuf buffer, int l) throws IOException { int n = readFully(buffer.nioBuffer(buffer.writerIndex(), l)); buffer.writerIndex(n); return n; }
/** * Deserializes an ArrowRecordBatch knowing the size of the entire message up front. This * minimizes the number of reads to the underlying stream. * * @param in the channel to deserialize from * @param block the object to deserialize to * @param alloc to allocate buffers * @return the deserialized ArrowRecordBatch * @throws IOException if something went wrong */ public static ArrowRecordBatch deserializeRecordBatch(ReadChannel in, ArrowBlock block, BufferAllocator alloc) throws IOException { // Metadata length contains integer prefix plus byte padding long totalLen = block.getMetadataLength() + block.getBodyLength(); if (totalLen > Integer.MAX_VALUE) { throw new IOException("Cannot currently deserialize record batches over 2GB"); } ArrowBuf buffer = alloc.buffer((int) totalLen); if (in.readFully(buffer, (int) totalLen) != totalLen) { throw new IOException("Unexpected end of input trying to read batch."); } ArrowBuf metadataBuffer = buffer.slice(4, block.getMetadataLength() - 4); Message messageFB = Message.getRootAsMessage(metadataBuffer.nioBuffer().asReadOnlyBuffer()); RecordBatch recordBatchFB = (RecordBatch) messageFB.header(new RecordBatch()); // Now read the body final ArrowBuf body = buffer.slice(block.getMetadataLength(), (int) totalLen - block.getMetadataLength()); return deserializeRecordBatch(recordBatchFB, body); }
if (in.readFully(buffer, (int) totalLen) != totalLen) { throw new IOException("Unexpected end of input trying to read batch.");