Rich-T-kid commented on code in PR #10277:
URL: https://github.com/apache/arrow-rs/pull/10277#discussion_r3604978471
##########
arrow-ipc/src/writer.rs:
##########
@@ -2016,158 +2142,14 @@ pub struct EncodedData {
pub arrow_data: Vec<u8>,
}
-fn encoded_data_to_buffers(
- out: &mut Vec<Buffer>,
- 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 flatbuf_size = encoded.ipc_message.len();
- let prefix_size = if write_options.write_legacy_ipc_format {
- 4
- } else {
- 8
- };
- let aligned_size = (flatbuf_size + prefix_size + alignment_mask) &
!alignment_mask;
- let padding_bytes = aligned_size - flatbuf_size - prefix_size;
-
- push_continuation_buffer(out, write_options, (aligned_size - prefix_size)
as i32)?;
- if flatbuf_size > 0 {
- out.push(Buffer::from(encoded.ipc_message));
- }
- push_padding_buffer(out, padding_bytes);
-
- let body_len = if arrow_data_len > 0 {
- out.push(Buffer::from(encoded.arrow_data));
- arrow_data_len
- } else {
- 0
- };
-
- Ok((aligned_size, body_len))
-}
-
-fn push_continuation_buffer(
- out: &mut Vec<Buffer>,
- write_options: &IpcWriteOptions,
- total_len: i32,
-) -> Result<(), ArrowError> {
- // Continuation bytes are generated stream framing, so they need a small
owned buffer.
- let mut buffer = Vec::with_capacity(8);
- write_continuation(&mut buffer, write_options, total_len)?;
- out.push(Buffer::from(buffer));
- Ok(())
-}
-
-fn push_padding_buffer(out: &mut Vec<Buffer>, len: usize) {
- if len > 0 {
- out.push(Buffer::from(&PADDING[..len]));
- }
-}
-
Review Comment:
🚀
##########
arrow-ipc/src/writer.rs:
##########
@@ -356,12 +515,12 @@ impl IpcDataGenerator {
message.add_bodyLength(0);
message.add_header(schema);
// TODO: custom metadata
- let data = message.finish();
- fbb.finish(data, None);
+ let root = message.finish();
+ fbb.finish(root, None);
- let data = fbb.finished_data();
+ let metadata = fbb.finished_data();
EncodedData {
- ipc_message: data.to_vec(),
+ ipc_message: metadata.to_vec(),
Review Comment:
nice, renamng data to metadata is a nice change
##########
arrow-ipc/src/writer.rs:
##########
@@ -2016,158 +2142,14 @@ pub struct EncodedData {
pub arrow_data: Vec<u8>,
}
-fn encoded_data_to_buffers(
- out: &mut Vec<Buffer>,
- 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 flatbuf_size = encoded.ipc_message.len();
- let prefix_size = if write_options.write_legacy_ipc_format {
- 4
- } else {
- 8
- };
- let aligned_size = (flatbuf_size + prefix_size + alignment_mask) &
!alignment_mask;
- let padding_bytes = aligned_size - flatbuf_size - prefix_size;
-
- push_continuation_buffer(out, write_options, (aligned_size - prefix_size)
as i32)?;
- if flatbuf_size > 0 {
- out.push(Buffer::from(encoded.ipc_message));
- }
- push_padding_buffer(out, padding_bytes);
-
- let body_len = if arrow_data_len > 0 {
- out.push(Buffer::from(encoded.arrow_data));
- arrow_data_len
- } else {
- 0
- };
-
- Ok((aligned_size, body_len))
-}
-
-fn push_continuation_buffer(
- out: &mut Vec<Buffer>,
- write_options: &IpcWriteOptions,
- total_len: i32,
-) -> Result<(), ArrowError> {
- // Continuation bytes are generated stream framing, so they need a small
owned buffer.
- let mut buffer = Vec::with_capacity(8);
- write_continuation(&mut buffer, write_options, total_len)?;
- out.push(Buffer::from(buffer));
- Ok(())
-}
-
-fn push_padding_buffer(out: &mut Vec<Buffer>, len: usize) {
- if len > 0 {
- out.push(Buffer::from(&PADDING[..len]));
- }
-}
-
/// Write a message's IPC data and buffers, returning metadata and buffer data
lengths written
pub fn write_message<W: Write>(
mut writer: W,
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 a = usize::from(write_options.alignment - 1);
- let buffer = encoded.ipc_message;
- let flatbuf_size = buffer.len();
- let prefix_size = if write_options.write_legacy_ipc_format {
- 4
- } else {
- 8
- };
- let aligned_size = (flatbuf_size + prefix_size + a) & !a;
- let padding_bytes = aligned_size - flatbuf_size - prefix_size;
-
- write_continuation(
- &mut writer,
- write_options,
- (aligned_size - prefix_size) as i32,
- )?;
-
- // write the flatbuf
- if flatbuf_size > 0 {
- writer.write_all(&buffer)?;
- }
- // write padding
- writer.write_all(&PADDING[..padding_bytes])?;
-
- // write arrow data
- let body_len = if arrow_data_len > 0 {
- write_body_buffers(&mut writer, &encoded.arrow_data,
write_options.alignment)?
- } else {
- 0
- };
-
- Ok((aligned_size, body_len))
-}
-
-fn write_body_buffers<W: Write>(
- mut writer: W,
- data: &[u8],
- alignment: u8,
-) -> Result<usize, ArrowError> {
- let len = data.len();
- let pad_len = pad_to_alignment(alignment, len);
- let total_len = len + pad_len;
-
- // write body buffer
- writer.write_all(data)?;
- if pad_len > 0 {
- writer.write_all(&PADDING[..pad_len])?;
- }
-
- Ok(total_len)
-}
-
-/// Write a record batch to the writer, writing the message size before the
message
-/// if the record batch is being written to a stream
-fn write_continuation<W: Write>(
- mut writer: W,
- write_options: &IpcWriteOptions,
- total_len: i32,
-) -> Result<usize, ArrowError> {
- let mut written = 8;
-
- // the version of the writer determines whether continuation markers
should be added
- 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 => {
- if !write_options.write_legacy_ipc_format {
- // v0.15.0 format
- writer.write_all(&CONTINUATION_MARKER)?;
- written = 4;
- }
- writer.write_all(&total_len.to_le_bytes()[..])?;
- }
- crate::MetadataVersion::V5 => {
- // write continuation marker and message length
- writer.write_all(&CONTINUATION_MARKER)?;
- writer.write_all(&total_len.to_le_bytes()[..])?;
- }
- z => panic!("Unsupported crate::MetadataVersion {z:?}"),
- };
-
- Ok(written)
Review Comment:
nice!
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,165 @@ 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.
+enum IpcMessageSink<'a> {
+ /// Write bytes directly to a synchronous writer.
+ Writer(&'a mut dyn Write),
+ /// Accumulate ordered buffers for deferred writing.
+ Buffers(&'a mut Vec<Buffer>),
+}
+
+impl IpcMessageSink<'_> {
+ fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
+ if bytes.is_empty() {
+ return Ok(());
+ }
+
+ match self {
+ Self::Writer(writer) => writer.write_all(&bytes)?,
+ Self::Buffers(out) => out.push(Buffer::from(bytes)),
+ }
+ Ok(())
+ }
+
+ fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(),
ArrowError> {
Review Comment:
I think a doc comment above this would be nice, otherwise it makes sense to
me.
##########
arrow-ipc/src/writer.rs:
##########
@@ -135,6 +135,165 @@ 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.
+enum IpcMessageSink<'a> {
+ /// Write bytes directly to a synchronous writer.
+ Writer(&'a mut dyn Write),
+ /// Accumulate ordered buffers for deferred writing.
+ Buffers(&'a mut Vec<Buffer>),
+}
+
+impl IpcMessageSink<'_> {
+ fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
+ if bytes.is_empty() {
+ return Ok(());
+ }
+
+ match self {
+ Self::Writer(writer) => writer.write_all(&bytes)?,
+ Self::Buffers(out) => out.push(Buffer::from(bytes)),
+ }
+ Ok(())
+ }
Review Comment:
nice now callers dont need to handle this at each call site
##########
arrow-ipc/src/writer.rs:
##########
@@ -734,13 +854,31 @@ impl IpcDataGenerator {
write_options: &IpcWriteOptions,
ipc_write_context: &mut IpcWriteContext,
out: &mut Vec<Buffer>,
+ ) -> Result<IpcWriteMetadata, ArrowError> {
+ let mut sink = IpcMessageSink::Buffers(out);
+ self.write_to_sink(
+ batch,
+ dictionary_tracker,
+ write_options,
+ ipc_write_context,
+ &mut sink,
+ )
+ }
Review Comment:
nice, this is a lot cleaner, and removes duplication
--
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]
