Jefffrey commented on code in PR #10277:
URL: https://github.com/apache/arrow-rs/pull/10277#discussion_r3643470216
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,269 @@ impl<'a> IpcBodySink<'a> {
}
}
+/// Destination for a complete framed IPC message.
+///
+/// This emits the stream/file framing around the serialized FlatBuffer
+/// [`crate::Message`] metadata plus its optional body buffers.
+trait IpcMessageSink {
+ fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError>;
+
+ fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(),
ArrowError>;
+
+ fn write_padding(&mut self, len: usize) -> Result<(), ArrowError>;
+
+ /// Writes the IPC continuation marker and metadata length prefix.
+ fn write_continuation(
+ &mut self,
+ write_options: &IpcWriteOptions,
+ metadata_len: i32,
+ ) -> Result<(), ArrowError>;
+
+ fn write_body_data(&mut self, data: Vec<u8>, alignment: u8) ->
Result<usize, ArrowError> {
+ let len = data.len();
+ let pad_len = pad_to_alignment(alignment, len);
+ self.write_vec(data)?;
+ self.write_padding(pad_len)?;
+ Ok(len + pad_len)
+ }
+
+ /// Writes an already encoded IPC message with optional contiguous body
data.
+ ///
+ /// This is used for schema and dictionary messages represented by
[`EncodedData`].
+ /// Returns the padded metadata length and body length written.
+ fn write_encoded_data(
+ &mut self,
+ encoded: EncodedData,
+ write_options: &IpcWriteOptions,
+ ) -> Result<(usize, usize), ArrowError> {
+ let arrow_data_len = encoded.arrow_data.len();
+ if arrow_data_len % usize::from(write_options.alignment) != 0 {
+ return Err(ArrowError::MemoryError(
+ "Arrow data not aligned".to_string(),
+ ));
+ }
+
+ let alignment_mask = usize::from(write_options.alignment - 1);
+ let metadata = encoded.ipc_message;
+ let metadata_len = metadata.len();
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let padded_header_len = (metadata_len + prefix_size + alignment_mask)
& !alignment_mask;
+ let padded_metadata_len = padded_header_len - prefix_size;
+ let metadata_padding = padded_metadata_len - metadata_len;
+
+ self.write_continuation(write_options, padded_metadata_len as i32)?;
+ self.write_vec(metadata)?;
+ self.write_padding(metadata_padding)?;
+
+ let body_len = if arrow_data_len > 0 {
+ self.write_body_data(encoded.arrow_data, write_options.alignment)?
+ } else {
+ 0
+ };
+
+ Ok((padded_header_len, body_len))
+ }
+
+ /// Writes a record batch message from its encoded metadata and body
buffers.
+ ///
+ /// The body buffers are already materialized as [`EncodedBuffer`]
segments,
+ /// allowing buffer output to preserve uncompressed Arrow buffers.
+ /// Returns the padded metadata length and body length written.
+ fn write_record_batch(
+ &mut self,
+ metadata: Vec<u8>,
+ encoded_buffers: Vec<EncodedBuffer>,
+ body_len: usize,
+ tail_pad: usize,
+ write_options: &IpcWriteOptions,
+ ) -> Result<(usize, usize), ArrowError> {
+ let alignment = write_options.alignment;
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let alignment_mask = usize::from(alignment - 1);
+ let padded_header_len = (metadata.len() + prefix_size +
alignment_mask) & !alignment_mask;
+ let padded_metadata_len = padded_header_len - prefix_size;
+ let metadata_padding = padded_metadata_len - metadata.len();
+
+ self.write_continuation(write_options, padded_metadata_len as i32)?;
+ self.write_vec(metadata)?;
+ self.write_padding(metadata_padding)?;
+ for enc in encoded_buffers {
+ let len = enc.len();
+ self.write_encoded_buffer(enc)?;
+ self.write_padding(pad_to_alignment(alignment, len))?;
+ }
+ self.write_padding(tail_pad)?;
+
+ Ok((padded_header_len, body_len))
+ }
+
+ /// Writes the IPC end-of-stream marker.
+ fn write_eos(&mut self, write_options: &IpcWriteOptions) -> Result<(),
ArrowError> {
+ self.write_continuation(write_options, 0)?;
+ Ok(())
+ }
+}
+
+/// Writes complete framed IPC messages to a synchronous writer.
+struct Writer<'a, W: Write> {
+ writer: &'a mut W,
+}
+
+impl<W: Write> IpcMessageSink for Writer<'_, W> {
Review Comment:
we could do something like this
```rust
impl<W> IpcMessageSink for W
where
W: Write,
{
// ...rest of impl
```
to avoid need for a wrapper `Writer` struct
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,269 @@ impl<'a> IpcBodySink<'a> {
}
}
+/// Destination for a complete framed IPC message.
+///
+/// This emits the stream/file framing around the serialized FlatBuffer
+/// [`crate::Message`] metadata plus its optional body buffers.
+trait IpcMessageSink {
+ fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError>;
+
+ fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(),
ArrowError>;
+
+ fn write_padding(&mut self, len: usize) -> Result<(), ArrowError>;
+
+ /// Writes the IPC continuation marker and metadata length prefix.
+ fn write_continuation(
+ &mut self,
+ write_options: &IpcWriteOptions,
+ metadata_len: i32,
+ ) -> Result<(), ArrowError>;
+
+ fn write_body_data(&mut self, data: Vec<u8>, alignment: u8) ->
Result<usize, ArrowError> {
+ let len = data.len();
+ let pad_len = pad_to_alignment(alignment, len);
+ self.write_vec(data)?;
+ self.write_padding(pad_len)?;
+ Ok(len + pad_len)
+ }
+
+ /// Writes an already encoded IPC message with optional contiguous body
data.
+ ///
+ /// This is used for schema and dictionary messages represented by
[`EncodedData`].
+ /// Returns the padded metadata length and body length written.
+ fn write_encoded_data(
+ &mut self,
+ encoded: EncodedData,
+ write_options: &IpcWriteOptions,
+ ) -> Result<(usize, usize), ArrowError> {
+ let arrow_data_len = encoded.arrow_data.len();
+ if arrow_data_len % usize::from(write_options.alignment) != 0 {
+ return Err(ArrowError::MemoryError(
+ "Arrow data not aligned".to_string(),
+ ));
+ }
+
+ let alignment_mask = usize::from(write_options.alignment - 1);
+ let metadata = encoded.ipc_message;
+ let metadata_len = metadata.len();
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let padded_header_len = (metadata_len + prefix_size + alignment_mask)
& !alignment_mask;
+ let padded_metadata_len = padded_header_len - prefix_size;
+ let metadata_padding = padded_metadata_len - metadata_len;
+
+ self.write_continuation(write_options, padded_metadata_len as i32)?;
+ self.write_vec(metadata)?;
+ self.write_padding(metadata_padding)?;
+
+ let body_len = if arrow_data_len > 0 {
+ self.write_body_data(encoded.arrow_data, write_options.alignment)?
+ } else {
+ 0
+ };
+
+ Ok((padded_header_len, body_len))
+ }
+
+ /// Writes a record batch message from its encoded metadata and body
buffers.
+ ///
+ /// The body buffers are already materialized as [`EncodedBuffer`]
segments,
+ /// allowing buffer output to preserve uncompressed Arrow buffers.
+ /// Returns the padded metadata length and body length written.
+ fn write_record_batch(
+ &mut self,
+ metadata: Vec<u8>,
+ encoded_buffers: Vec<EncodedBuffer>,
+ body_len: usize,
+ tail_pad: usize,
+ write_options: &IpcWriteOptions,
+ ) -> Result<(usize, usize), ArrowError> {
+ let alignment = write_options.alignment;
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let alignment_mask = usize::from(alignment - 1);
+ let padded_header_len = (metadata.len() + prefix_size +
alignment_mask) & !alignment_mask;
+ let padded_metadata_len = padded_header_len - prefix_size;
+ let metadata_padding = padded_metadata_len - metadata.len();
+
+ self.write_continuation(write_options, padded_metadata_len as i32)?;
+ self.write_vec(metadata)?;
+ self.write_padding(metadata_padding)?;
+ for enc in encoded_buffers {
+ let len = enc.len();
+ self.write_encoded_buffer(enc)?;
+ self.write_padding(pad_to_alignment(alignment, len))?;
+ }
+ self.write_padding(tail_pad)?;
+
+ Ok((padded_header_len, body_len))
+ }
+
+ /// Writes the IPC end-of-stream marker.
+ fn write_eos(&mut self, write_options: &IpcWriteOptions) -> Result<(),
ArrowError> {
+ self.write_continuation(write_options, 0)?;
+ Ok(())
+ }
+}
+
+/// Writes complete framed IPC messages to a synchronous writer.
+struct Writer<'a, W: Write> {
+ writer: &'a mut W,
+}
+
+impl<W: Write> IpcMessageSink for Writer<'_, W> {
+ fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
+ if !bytes.is_empty() {
+ self.writer.write_all(&bytes)?;
+ }
+ Ok(())
+ }
+
+ fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(),
ArrowError> {
+ self.writer.write_all(buffer.as_slice())?;
+ Ok(())
+ }
+
+ fn write_padding(&mut self, len: usize) -> Result<(), ArrowError> {
+ if len > 0 {
+ self.writer.write_all(&PADDING[..len])?;
+ }
+ Ok(())
+ }
+
+ fn write_continuation(
Review Comment:
this `write_continuation()` logic seems like it can be hoisted into the
trait impl as it looks duplicated with the buffers version
##########
arrow-ipc/src/writer.rs:
##########
@@ -1564,6 +1852,139 @@ impl<W: Write> RecordBatchWriter for FileWriter<W> {
}
}
+/// Arrow IPC stream encoder.
+///
+/// Encodes Arrow [`RecordBatch`]es to byte buffers using the [IPC Streaming
Format],
+/// without performing any IO.
+///
+/// The returned [`Buffer`]s are ordered and should be written to the
destination
+/// stream in order. Uncompressed record batch body buffers can share the
original
+/// Arrow buffers instead of being copied into an intermediate contiguous
buffer.
+///
+/// # Example
+/// ```
+/// # use arrow_array::record_batch;
+/// # use arrow_ipc::writer::StreamEncoder;
+/// # use arrow_schema::ArrowError;
+/// # fn main() -> Result<(), ArrowError> {
+/// let batch = record_batch!(("a", Int32, [1, 2, 3]))?;
+///
+/// let mut encoder = StreamEncoder::try_new(&batch.schema())?;
+/// let mut stream = vec![];
+/// for buffer in encoder.encode(&batch)? {
+/// stream.extend_from_slice(buffer.as_slice());
+/// }
+/// for buffer in encoder.finish()? {
+/// stream.extend_from_slice(buffer.as_slice());
+/// }
+/// # Ok(())
+/// # }
+/// ```
+pub struct StreamEncoder {
+ schema: Schema,
+ /// IPC write options
+ write_options: IpcWriteOptions,
+ /// Whether the stream schema has been encoded
+ schema_encoded: bool,
+ /// Whether the end-of-stream marker has been encoded
+ finished: bool,
+ /// Keeps track of dictionaries that have been encoded
+ dictionary_tracker: DictionaryTracker,
+ data_gen: IpcDataGenerator,
+ ipc_write_context: IpcWriteContext,
+}
+
+impl StreamEncoder {
+ /// Try to create a new stream encoder.
+ pub fn try_new(schema: &Schema) -> Result<Self, ArrowError> {
+ let write_options = IpcWriteOptions::default();
+ Self::try_new_with_options(schema, write_options)
+ }
+
+ /// Try to create a new stream encoder with [`IpcWriteOptions`].
+ pub fn try_new_with_options(
+ schema: &Schema,
+ write_options: IpcWriteOptions,
+ ) -> Result<Self, ArrowError> {
+ ensure_supported_ipc_schema(schema)?;
+
+ Ok(Self {
+ schema: schema.clone(),
+ write_options,
+ schema_encoded: false,
+ finished: false,
+ dictionary_tracker: DictionaryTracker::new(false),
+ data_gen: IpcDataGenerator::default(),
+ ipc_write_context: IpcWriteContext::default(),
+ })
+ }
+
+ /// Encode a [`RecordBatch`] into buffers.
+ ///
+ /// The first call also includes the IPC stream schema message before the
+ /// record batch message. Later calls only include dictionary and record
+ /// batch messages.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the encoder is already finished or encoding fails.
+ pub fn encode(&mut self, batch: &RecordBatch) -> Result<Vec<Buffer>,
ArrowError> {
+ if self.finished {
+ return Err(ArrowError::IpcError(
+ "Cannot encode record batch to stream encoder as it is
closed".to_string(),
+ ));
+ }
+
+ let mut out = vec![];
+ self.encode_schema(&mut out)?;
+ self.data_gen.encode_to_buffers(
+ batch,
+ &mut self.dictionary_tracker,
+ &self.write_options,
+ &mut self.ipc_write_context,
+ &mut out,
+ )?;
+ Ok(out)
+ }
+
+ /// Encode the end-of-stream marker and mark this encoder as finished.
+ ///
+ /// If no batches have been encoded, this also emits the IPC stream schema
+ /// message so the returned buffers form a valid empty IPC stream.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the encoder is already finished.
+ pub fn finish(&mut self) -> Result<Vec<Buffer>, ArrowError> {
Review Comment:
```suggestion
pub fn finish(self) -> Result<Vec<Buffer>, ArrowError> {
```
if we consume self i think this guarantees no more usage (i.e. no need for
runtime `self.finished` check)
--
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]
