Phoenix500526 commented on code in PR #11374:
URL: https://github.com/apache/arrow-rs/pull/11374#discussion_r4398197290
##########
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> {
+ let mut decoder = StreamDecoder::new();
+ let mut working = buffer;
+ let mut pending: Option<RecordBatch> = None;
+ // Drive the decoder until the schema is parsed. `StreamDecoder`
+ // may consume the schema *and* the first record batch in a
+ // single `decode()` call (its inner loop only exits on a
+ // RecordBatch or on empty input); stash that eager batch.
+ while decoder.schema().is_none() {
+ if working.is_empty() {
+ return Err(ArrowError::IpcError(
+ "Expected schema message, found empty stream.".to_string(),
+ ));
+ }
+ if let Some(batch) = decoder.decode(&mut working)? {
+ if decoder.schema().is_none() {
+ return Err(ArrowError::IpcError(
+ "Expected schema as first IPC message, got record
batch".to_string(),
+ ));
+ }
+ pending = Some(batch);
+ break;
+ }
+ }
+ let schema = decoder.schema().expect("schema decoded above");
+ Ok(Self {
+ decoder,
+ buffer: working,
+ schema,
+ pending,
+ finished: false,
+ })
+ }
+
+ /// Create a new reader from an owned [`bytes::Bytes`].
+ ///
+ /// Conversion to [`Buffer`] shares the existing allocation without
copying.
+ /// Decoding may still allocate for alignment or decompression, as
described
+ /// in [buffer sharing](Self#buffer-sharing).
+ ///
+ /// # Errors
+ ///
+ /// Returns the same errors as [`try_new`](Self::try_new).
+ pub fn try_new_from_bytes(bytes: bytes::Bytes) -> Result<Self, ArrowError>
{
+ Self::try_new(Buffer::from(bytes))
+ }
+
+ /// The schema of the stream.
+ pub fn schema(&self) -> &SchemaRef {
+ &self.schema
+ }
+
+ /// Returns `true` once the reader has consumed the EOS marker or observed
+ /// that the input has been exhausted.
+ ///
+ /// This can also be `true` after reporting a truncated message, so it does
+ /// not imply that decoding succeeded. A decoding error does not always set
+ /// this flag; callers should stop reading after any error.
+ pub fn is_finished(&self) -> bool {
+ self.finished || self.decoder.is_finished()
+ }
+}
+
+impl Iterator for BufferStreamReader {
+ type Item = Result<RecordBatch, ArrowError>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ if let Some(batch) = self.pending.take() {
+ return Some(Ok(batch));
+ }
+ loop {
+ if self.finished || self.decoder.is_finished() {
+ return None;
+ }
+ if self.buffer.is_empty() {
+ self.finished = true;
+ // Like `StreamReader`, a stream that ends on a message
boundary
+ // without an EOS marker is accepted, but a truncated message
is an error
+ return self.decoder.finish().err().map(Err);
+ }
+ match self.decoder.decode(&mut self.buffer) {
+ Ok(Some(batch)) => return Some(Ok(batch)),
+ Ok(None) => {}
+ Err(e) => return Some(Err(e)),
Review Comment:
We should set `self.finished = true` here like what we do when
`self.buffer.is_empty()` is true.
##########
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> {
+ let mut decoder = StreamDecoder::new();
+ let mut working = buffer;
+ let mut pending: Option<RecordBatch> = None;
+ // Drive the decoder until the schema is parsed. `StreamDecoder`
+ // may consume the schema *and* the first record batch in a
+ // single `decode()` call (its inner loop only exits on a
+ // RecordBatch or on empty input); stash that eager batch.
+ while decoder.schema().is_none() {
+ if working.is_empty() {
+ return Err(ArrowError::IpcError(
+ "Expected schema message, found empty stream.".to_string(),
+ ));
+ }
+ if let Some(batch) = decoder.decode(&mut working)? {
+ if decoder.schema().is_none() {
+ return Err(ArrowError::IpcError(
+ "Expected schema as first IPC message, got record
batch".to_string(),
+ ));
+ }
+ pending = Some(batch);
+ break;
+ }
+ }
+ let schema = decoder.schema().expect("schema decoded above");
+ Ok(Self {
+ decoder,
+ buffer: working,
+ schema,
+ pending,
+ finished: false,
+ })
+ }
+
+ /// Create a new reader from an owned [`bytes::Bytes`].
+ ///
+ /// Conversion to [`Buffer`] shares the existing allocation without
copying.
+ /// Decoding may still allocate for alignment or decompression, as
described
+ /// in [buffer sharing](Self#buffer-sharing).
+ ///
+ /// # Errors
+ ///
+ /// Returns the same errors as [`try_new`](Self::try_new).
+ pub fn try_new_from_bytes(bytes: bytes::Bytes) -> Result<Self, ArrowError>
{
+ Self::try_new(Buffer::from(bytes))
+ }
+
+ /// The schema of the stream.
Review Comment:
How about implementing the `RecordBatchReader` trait instead of this
`schema` method?
##########
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:
Could we add a test with deliberately misaligned input to cover the copy
fallback described here? The existing aliases_input test checks buffer sharing,
but doesn't exercise that path.
--
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]
