Rich-T-kid commented on code in PR #11374:
URL: https://github.com/apache/arrow-rs/pull/11374#discussion_r4394298179
##########
arrow-ipc/src/reader/stream.rs:
##########
@@ -309,6 +309,186 @@ impl StreamDecoder {
_ => Err(ArrowError::IpcError("Unexpected End of
Stream".to_string())),
}
}
+
+ /// Returns `true` once the stream's EOS marker has been consumed.
+ pub fn is_finished(&self) -> bool {
+ matches!(self.state, DecoderState::Finished)
+ }
+}
+
+/// Pull-based IPC stream reader over an owned in-memory [`Buffer`], with
+/// zero-copy decoding where possible.
+///
+/// Where [`StreamReader`](crate::reader::StreamReader) reads from an arbitrary
+/// `R: Read` and copies each message body into a fresh allocation, this reader
+/// takes a fully materialized [`Buffer`] and decodes message bodies directly
+/// from it. Use this when the entire IPC stream is already in memory, such as
+/// an HTTP body, a memory-mapped file, a shared-memory segment, or
+/// [`bytes::Bytes`] from object storage.
+///
+/// For a streaming push-based interface that handles partial buffers (e.g. a
+/// chunked byte stream from object storage), use [`StreamDecoder`] directly.
+///
+/// # Buffer sharing
+///
+/// Uncompressed, properly aligned column buffers can alias the input
allocation.
+/// Misaligned buffers are copied into aligned allocations, and compressed
buffers
+/// require allocations for decompression. Zero-copy decoding therefore depends
+/// on both the input alignment and how the stream was encoded.
+///
+/// Decoding still allocates schema and array metadata and validates array
data;
+/// sharing buffers does not skip validation. Keeping a decoded batch with
shared
+/// buffers alive also retains the input allocation, even after this reader is
dropped.
+///
+/// # Errors
+///
+/// Errors encountered while decoding the schema, leading dictionaries, or the
+/// first record batch may be returned by the constructor. Later decoding
errors
+/// are returned as iterator items. Stop reading after an error; the reader
does
+/// not guarantee that it can resume decoding or that the next call returns
`None`.
+///
+/// # Example
+///
+/// ```
+/// # use std::sync::Arc;
+/// # use arrow_array::{Int32Array, RecordBatch};
+/// # use arrow_buffer::Buffer;
+/// # use arrow_ipc::reader::BufferStreamReader;
+/// # use arrow_ipc::writer::StreamWriter;
+/// # use arrow_schema::{DataType, Field, Schema};
+/// let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32,
false)]));
+/// let batch = RecordBatch::try_new(
+/// schema.clone(),
+/// vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+/// ).unwrap();
+///
+/// let mut bytes = Vec::new();
+/// let mut w = StreamWriter::try_new(&mut bytes, schema.as_ref()).unwrap();
+/// w.write(&batch).unwrap();
+/// w.finish().unwrap();
+/// drop(w);
+///
+/// let mut reader = BufferStreamReader::try_new(Buffer::from(bytes)).unwrap();
+/// assert_eq!(reader.schema().as_ref(), schema.as_ref());
+/// let read = reader.next().unwrap().unwrap();
+/// assert_eq!(read, batch);
+/// assert!(reader.next().is_none());
+/// ```
+#[derive(Debug)]
+pub struct BufferStreamReader {
+ decoder: StreamDecoder,
+ buffer: Buffer,
+ schema: SchemaRef,
+ /// `StreamDecoder::decode` may parse the schema and the first
+ /// record batch in a single call (its inner loop only exits on a
+ /// `RecordBatch` header or empty input). When that happens we stash
+ /// the eager batch here so it's returned by the first `next()`.
+ pending: Option<RecordBatch>,
+ /// Set once the iterator has returned its final item.
+ finished: bool,
+}
+
+impl BufferStreamReader {
+ /// Create a new reader, eagerly draining the schema message (and any
+ /// leading dictionary messages) so [`schema`](Self::schema) is cheap
+ /// and infallible afterwards.
+ ///
+ /// This may also decode the first record batch, which is retained and
+ /// returned by the first call to [`next`](Iterator::next).
+ /// See [buffer sharing](Self#buffer-sharing) for when decoding can avoid
copies.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if `buffer` does not start with a valid IPC stream
+ /// schema message, or if decoding the leading dictionaries or first record
+ /// batch fails during construction.
+ pub fn try_new(buffer: Buffer) -> Result<Self, ArrowError> {
Review Comment:
for performance reasons we should add a way to avoid the saftey checks that
`StreamDecoder` does.
see
https://github.com/apache/arrow-rs/blob/5db01b1fa7ba6f70d5f63abf31c48fabed40929b/arrow-ipc/src/reader/stream.rs#L125
see related
https://github.com/apache/arrow-rs/blob/5db01b1fa7ba6f70d5f63abf31c48fabed40929b/arrow-ipc/src/reader/stream.rs#L106
##########
arrow-ipc/src/reader/stream.rs:
##########
@@ -309,6 +309,186 @@ impl StreamDecoder {
_ => Err(ArrowError::IpcError("Unexpected End of
Stream".to_string())),
}
}
+
+ /// Returns `true` once the stream's EOS marker has been consumed.
+ pub fn is_finished(&self) -> bool {
+ matches!(self.state, DecoderState::Finished)
+ }
+}
+
+/// Pull-based IPC stream reader over an owned in-memory [`Buffer`], with
+/// zero-copy decoding where possible.
+///
+/// Where [`StreamReader`](crate::reader::StreamReader) reads from an arbitrary
+/// `R: Read` and copies each message body into a fresh allocation, this reader
+/// takes a fully materialized [`Buffer`] and decodes message bodies directly
+/// from it. Use this when the entire IPC stream is already in memory, such as
+/// an HTTP body, a memory-mapped file, a shared-memory segment, or
+/// [`bytes::Bytes`] from object storage.
+///
+/// For a streaming push-based interface that handles partial buffers (e.g. a
+/// chunked byte stream from object storage), use [`StreamDecoder`] directly.
+///
+/// # Buffer sharing
+///
+/// Uncompressed, properly aligned column buffers can alias the input
allocation.
+/// Misaligned buffers are copied into aligned allocations, and compressed
buffers
+/// require allocations for decompression. Zero-copy decoding therefore depends
+/// on both the input alignment and how the stream was encoded.
Review Comment:
🤔 have you experimented with doing 1 bulk copy if the initual buffer is
mis-aligned?
##########
arrow-ipc/src/reader/stream.rs:
##########
@@ -309,6 +309,186 @@ impl StreamDecoder {
_ => Err(ArrowError::IpcError("Unexpected End of
Stream".to_string())),
}
}
+
+ /// Returns `true` once the stream's EOS marker has been consumed.
+ pub fn is_finished(&self) -> bool {
+ matches!(self.state, DecoderState::Finished)
+ }
+}
+
+/// Pull-based IPC stream reader over an owned in-memory [`Buffer`], with
+/// zero-copy decoding where possible.
+///
+/// Where [`StreamReader`](crate::reader::StreamReader) reads from an arbitrary
+/// `R: Read` and copies each message body into a fresh allocation, this reader
+/// takes a fully materialized [`Buffer`] and decodes message bodies directly
+/// from it. Use this when the entire IPC stream is already in memory, such as
+/// an HTTP body, a memory-mapped file, a shared-memory segment, or
+/// [`bytes::Bytes`] from object storage.
+///
+/// For a streaming push-based interface that handles partial buffers (e.g. a
+/// chunked byte stream from object storage), use [`StreamDecoder`] directly.
Review Comment:
this is a neat idea
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
