Jefffrey commented on code in PR #10277:
URL: https://github.com/apache/arrow-rs/pull/10277#discussion_r3644468932
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,225 @@ 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_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(
+ &mut self,
+ write_options: &IpcWriteOptions,
+ metadata_len: i32,
+ ) -> Result<(), ArrowError> {
+ let mut buffer = [0; 8];
+ let len = match write_options.metadata_version {
+ crate::MetadataVersion::V1
+ | crate::MetadataVersion::V2
+ | crate::MetadataVersion::V3 => {
+ unreachable!("Options with the metadata version cannot be
created")
+ }
+ crate::MetadataVersion::V4 => {
+ let metadata_len_bytes = metadata_len.to_le_bytes();
+ if !write_options.write_legacy_ipc_format {
+ // v0.15.0 format
+ buffer[..4].copy_from_slice(&CONTINUATION_MARKER);
+ buffer[4..].copy_from_slice(&metadata_len_bytes);
+ 8
+ } else {
+ buffer[..4].copy_from_slice(&metadata_len_bytes);
+ 4
+ }
+ }
+ crate::MetadataVersion::V5 => {
+ buffer[..4].copy_from_slice(&CONTINUATION_MARKER);
+ buffer[4..].copy_from_slice(&metadata_len.to_le_bytes());
+ 8
+ }
+ z => panic!("Unsupported crate::MetadataVersion {z:?}"),
+ };
+ self.write_slice(&buffer[..len])
+ }
+
+ 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(())
+ }
+}
+
+impl<W> IpcMessageSink for W
+where
+ W: Write,
+{
+ fn write_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError> {
+ if !bytes.is_empty() {
+ self.write_all(bytes)?;
+ }
+ Ok(())
+ }
+
+ fn write_record_batch(
Review Comment:
do we still need this custom impl, or can rely on the default?
--
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]
