Phoenix500526 commented on code in PR #10277:
URL: https://github.com/apache/arrow-rs/pull/10277#discussion_r3600765700
##########
arrow-ipc/src/writer.rs:
##########
@@ -725,6 +724,67 @@ impl IpcDataGenerator {
})
}
+ /// Encode dictionary batches and the record batch to output buffers,
skipping the
+ /// intermediate body `Vec<u8>` allocations for uncompressed record batch
buffers.
+ fn encode_to_buffers(
+ &self,
+ batch: &RecordBatch,
+ dictionary_tracker: &mut DictionaryTracker,
+ write_options: &IpcWriteOptions,
+ ipc_write_context: &mut IpcWriteContext,
+ out: &mut Vec<Buffer>,
+ ) -> Result<IpcWriteMetadata, ArrowError> {
+ let encoded_dictionaries =
+ self.encode_all_dicts(batch, dictionary_tracker, write_options,
ipc_write_context)?;
+
+ let mut dictionary_block_sizes =
Vec::with_capacity(encoded_dictionaries.len());
+ for dict in encoded_dictionaries {
+ dictionary_block_sizes.push(encoded_data_to_buffers(out, dict,
write_options)?);
+ }
+
+ let capacity = batch
+ .columns()
+ .iter()
+ .map(|a| estimate_encoded_buffer_count(a.data_type()))
+ .sum();
+ let mut encoded_buffers: Vec<EncodedBuffer> =
Vec::with_capacity(capacity);
+ let (ipc_message, body_len, tail_pad) = self.record_batch_to_bytes(
+ batch,
+ write_options,
+ ipc_write_context,
+ &mut IpcBodySink::Collect(&mut encoded_buffers),
+ )?;
+
+ let alignment = write_options.alignment;
+ let alignment_mask = usize::from(alignment - 1);
+ let prefix_size = if write_options.write_legacy_ipc_format {
+ 4
+ } else {
+ 8
+ };
+ let ipc_message_len = ipc_message.len();
+ let aligned_size = (ipc_message_len + prefix_size + alignment_mask) &
!alignment_mask;
+ push_continuation_buffer(out, write_options, (aligned_size -
prefix_size) as i32)?;
+ out.push(Buffer::from(ipc_message));
+ push_padding_buffer(out, aligned_size - ipc_message_len - prefix_size);
+ for enc in encoded_buffers {
+ let len = enc.len();
+ let buffer = match enc {
+ EncodedBuffer::Raw(buffer) => buffer,
+ EncodedBuffer::Compressed(bytes) => Buffer::from(bytes),
+ };
+ out.push(buffer);
+ push_padding_buffer(out, pad_to_alignment(alignment, len));
+ }
Review Comment:
Good catch. I refactored this to use a shared internal `IpcMessageSink` for
the final framed IPC message emission. I kep `IpcBodySink` focused on
materializing the record batch body buffers, and added `IpcMessageSink` one
layer above it to handle where the complete framed IPC message goes: either a
`Write` sink or ordered `Buffer`s.
With that, `IpcDataGenerator::write` and `encode_to_buffers` now share the
same `write_to_sink` path instead of duplicating dictionary/message framing
logic.
--
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]
