alamb commented on code in PR #10277:
URL: https://github.com/apache/arrow-rs/pull/10277#discussion_r3651422483
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,226 @@ impl<'a> IpcBodySink<'a> {
}
}
+struct MetadataLayout {
+ padded_header_len: usize,
+ padded_metadata_len: usize,
+ metadata_padding: usize,
+}
+
+#[inline]
+fn metadata_layout(metadata_len: usize, write_options: &IpcWriteOptions) ->
MetadataLayout {
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let alignment_mask = usize::from(write_options.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;
+
+ MetadataLayout {
+ padded_header_len,
+ padded_metadata_len,
+ metadata_padding,
+ }
+}
+
+/// 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_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError>;
+
+ fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
+ self.write_slice(&bytes)
+ }
+
+ fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(),
ArrowError> {
+ self.write_slice(buffer.as_slice())
+ }
+
+ fn write_padding(&mut self, len: usize) -> Result<(), ArrowError> {
+ self.write_slice(&PADDING[..len])
+ }
+
+ /// Writes the IPC continuation marker and metadata length prefix.
+ fn write_continuation(
Review Comment:
I found the mixing of methods you are supposed to override (like
`write_slice` and `write_vec`) and methods that you aren't (like
`write_continuation`) to be somewhat confusing.
It would be easier to read this code I think if we separated out the two
types of methods
Perhaps something like this
```rust
/// The only things a sink actually needs to customize.
trait IpcMessageSink {
fn write_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError>;
fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
self.write_slice(&bytes)
}
fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) ->
Result<(), ArrowError> {
self.write_slice(buffer.as_slice())
}
}
/// The shared framing code
trait IpcMessageSinkExt: IpcMessageSink {
fn write_padding(&mut self, len: usize) -> Result<(), ArrowError> {
self.write_slice(&PADDING[..len])
}
fn write_continuation(&mut self, write_options: &IpcWriteOptions,
metadata_len: i32) -> Result<(), ArrowError> {
// ... unchanged body ...
}
fn write_encoded_data(&mut self, encoded: EncodedData, write_options:
&IpcWriteOptions) -> Result<(usize, usize), ArrowError> {
// ... unchanged body, calling self.write_vec / self.write_padding
...
}
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> {
// ... unchanged body ...
}
fn write_eos(&mut self, write_options: &IpcWriteOptions) -> Result<(),
ArrowError> {
self.write_continuation(write_options, 0)
}
}
/// Blanket-implemented, so no impl of `IpcMessageSink` can ever override
these.
impl<T: IpcMessageSink + ?Sized> IpcMessageSinkExt for T {}
```
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,226 @@ impl<'a> IpcBodySink<'a> {
}
}
+struct MetadataLayout {
+ padded_header_len: usize,
+ padded_metadata_len: usize,
+ metadata_padding: usize,
+}
+
+#[inline]
+fn metadata_layout(metadata_len: usize, write_options: &IpcWriteOptions) ->
MetadataLayout {
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let alignment_mask = usize::from(write_options.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;
+
+ MetadataLayout {
+ padded_header_len,
+ padded_metadata_len,
+ metadata_padding,
+ }
+}
+
+/// 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.
Review Comment:
It might also be worth mentioning to point out that the default
implementation will copy data
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,226 @@ impl<'a> IpcBodySink<'a> {
}
}
+struct MetadataLayout {
+ padded_header_len: usize,
+ padded_metadata_len: usize,
+ metadata_padding: usize,
+}
+
+#[inline]
+fn metadata_layout(metadata_len: usize, write_options: &IpcWriteOptions) ->
MetadataLayout {
Review Comment:
As a follow on, this might naturally be encapsualted as
`MetadataLayout::new` type constructor
##########
arrow-ipc/src/writer.rs:
##########
@@ -2523,6 +2795,116 @@ mod tests {
stream_reader.next().unwrap().unwrap()
}
+ /// Encodes record batches with [`StreamEncoder`] into one contiguous byte
vector.
+ ///
+ /// This mirrors callers that need IPC stream bytes without a synchronous
+ /// [`Write`] implementation, such as encoding buffers before forwarding
them
+ /// to an async sink.
+ fn encode_stream(
+ schema: &Schema,
+ batches: &[RecordBatch],
+ options: IpcWriteOptions,
+ ) -> Vec<u8> {
+ let mut encoder = StreamEncoder::try_new_with_options(schema,
options).unwrap();
+ let mut bytes = Vec::new();
Review Comment:
since `Vec` implements `Write` why not just use the Write impl directly?
##########
arrow-ipc/src/writer.rs:
##########
@@ -1564,6 +1805,123 @@ 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,
+ /// 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,
+ 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 encoding fails.
+ pub fn encode(&mut self, batch: &RecordBatch) -> Result<Vec<Buffer>,
ArrowError> {
Review Comment:
I think this API is quite good and supports a usecase like `async` writer
quite well.
My only suggestion / feedback is that this API basically requires buffering
an entire record batch in RAM and now requies copying into several intermediate
Vecs (e.g. for the metadata header, etc)
Another API we might consider would be to make the SteamEncoder templated on
`IpcMessageSink` so it can bass the buffers directly
Something like
```rust
struct StreamEncoder<IpcMessageSink>` {
...
/// encode the batches, calling methods like `append_vec()`,
`append_buffer`, append_slice, etc
pub fn encode(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
}
```
However, this does still feel like there is "IO" in the encoder 🤔 -- on the
other hand you could implement an `IpcMessageSync` that just buffers the data
to get the same effect.
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,226 @@ impl<'a> IpcBodySink<'a> {
}
}
+struct MetadataLayout {
+ padded_header_len: usize,
+ padded_metadata_len: usize,
+ metadata_padding: usize,
+}
+
+#[inline]
+fn metadata_layout(metadata_len: usize, write_options: &IpcWriteOptions) ->
MetadataLayout {
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let alignment_mask = usize::from(write_options.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;
+
+ MetadataLayout {
+ padded_header_len,
+ padded_metadata_len,
+ metadata_padding,
+ }
+}
+
+/// 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_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError>;
+
+ fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
Review Comment:
+1 for taking owned `Vec`
--
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]
