Rich-T-kid commented on code in PR #10128:
URL: https://github.com/apache/arrow-rs/pull/10128#discussion_r3839314786
##########
arrow-ipc/src/writer.rs:
##########
@@ -271,31 +269,37 @@ impl<T: IpcMessageSink + ?Sized> IpcMessageSinkExt for T
{}
/// Optional hot-path hook for record batch messages.
trait IpcRecordBatchSink: IpcMessageSinkExt {
- /// Writes a record batch message from its encoded metadata and body
buffers.
+ /// Writes a record batch (or dictionary batch) message from its encoded
+ /// metadata and body buffers.
+ ///
+ /// `metadata` is the raw flatbuffer [`crate::Message`] (without
continuation
+ /// prefix), borrowed from the reused [`FlatBufferBuilder`]. The body
buffers are
+ /// already materialized as [`EncodedBuffer`] segments, allowing buffer
output to
+ /// preserve uncompressed Arrow buffers. They are drained out of
+ /// `encoded_buffers` so the caller can reuse its allocation.
+ ///
+ /// Each body buffer is padded to the alignment as it is written, so the
body
+ /// needs no trailing padding.
///
- /// 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>,
+ metadata: &[u8],
+ encoded_buffers: &mut Vec<EncodedBuffer>,
body_len: usize,
- tail_pad: usize,
write_options: &IpcWriteOptions,
) -> Result<(usize, usize), ArrowError> {
let alignment = write_options.alignment;
let layout = MetadataLayout::new(metadata.len(), write_options);
self.write_continuation(write_options, layout.padded_metadata_len as
i32)?;
- self.write_vec(metadata)?;
+ self.write_slice(metadata)?;
self.write_padding(layout.metadata_padding)?;
- for enc in encoded_buffers {
+ for enc in encoded_buffers.drain(..) {
Review Comment:
is there a specific reason to use `.drain(..)` here instead of the regular
`for x in y` loop?
##########
arrow-ipc/src/writer.rs:
##########
@@ -1,5 +1,3 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
Review Comment:
can we add this back in
##########
arrow-ipc/src/writer.rs:
##########
@@ -949,29 +960,48 @@ impl IpcDataGenerator {
ipc_write_context: &mut IpcWriteContext,
sink: &mut S,
) -> 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(sink.write_encoded_data(dict,
write_options)?);
- }
+ let dictionaries = self.collect_all_dicts(batch, dictionary_tracker,
write_options)?;
+ let mut dictionary_block_sizes =
Vec::with_capacity(dictionaries.len());
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 (metadata, body_len, tail_pad) = self.record_batch_to_bytes(
+
+ for dict in &dictionaries {
+ encoded_buffers.clear();
+
+ let body_len = self.dictionary_batch_to_sink(
+ dict,
+ write_options,
+ ipc_write_context,
+ &mut IpcBodySink::Collect(&mut encoded_buffers),
+ )?;
+
+ dictionary_block_sizes.push(sink.write_record_batch(
+ ipc_write_context.mut_fbb().finished_data(),
+ &mut encoded_buffers,
+ body_len,
+ write_options,
+ )?);
+ }
+
+ encoded_buffers.clear();
+ let body_len = self.record_batch_to_bytes(
batch,
write_options,
ipc_write_context,
&mut IpcBodySink::Collect(&mut encoded_buffers),
)?;
- let (padded_header_len, body_len) =
- sink.write_record_batch(metadata, encoded_buffers, body_len,
tail_pad, write_options)?;
+ let (padded_header_len, body_len) = sink.write_record_batch(
+ ipc_write_context.mut_fbb().finished_data(),
+ &mut encoded_buffers,
+ body_len,
+ write_options,
+ )?;
Review Comment:
This seems to be the only portion of this PR that interacts with the
non-dictionary encoded write path. I dont see anything that should be causing
an 18% regression.
##########
arrow-ipc/src/writer.rs:
##########
@@ -1001,105 +1031,141 @@ impl IpcDataGenerator {
/// Encodes a `RecordBatch` into a flatbuffer IPC message and fills `sink`
with the
/// serialised buffer data.
///
- /// Returns `(metadata, body_len, tail_pad)`: the FlatBuffer
[`crate::Message`] bytes, the
- /// total body length including trailing padding, and the trailing
alignment padding byte count.
+ /// Returns the total body length written to `sink` (including per-buffer
alignment
+ /// padding).
+ ///
+ /// The FlatBuffer [`crate::Message`] is located in `ipc_write_context`'s
+ /// [`FlatBufferBuilder`] finished bytes. A successful Result from this
function
+ /// guarantees the builder is in a finished state to call
+ /// [`FlatBufferBuilder::finished_data`].
fn record_batch_to_bytes(
&self,
batch: &RecordBatch,
write_options: &IpcWriteOptions,
ipc_write_context: &mut IpcWriteContext,
sink: &mut IpcBodySink<'_>,
- ) -> Result<(Vec<u8>, usize, usize), ArrowError> {
- let batch_compression_type = write_options.batch_compression_type;
+ ) -> Result<usize, ArrowError> {
+ // Reset the fbb
+ ipc_write_context.mut_fbb().reset();
- let compression = batch_compression_type.map(|batch_compression_type| {
- let fbb = ipc_write_context.mut_fbb();
- let mut c = crate::BodyCompressionBuilder::new(fbb);
- c.add_method(crate::BodyCompressionMethod::BUFFER);
- c.add_codec(batch_compression_type);
- c.finish()
- });
-
- let batch_compression_level = write_options.batch_compression_level;
- let compression_codec: Option<CompressionCodec> =
batch_compression_type
- .map(|compression_type| match batch_compression_level {
- Some(level) => {
-
CompressionCodec::try_new_with_compression_level(compression_type, level)
- }
- None => compression_type.try_into(),
- })
- .transpose()?;
-
- let alignment = write_options.alignment;
- let mut variadic_buffer_counts = vec![];
- let mut meta = IpcMetadataBuilder::default();
- let mut offset = 0i64;
-
- for array in batch.columns() {
- let array_data = array.to_data();
- offset = write_array_data(
- &array_data,
- &mut meta,
- sink,
- offset,
- compression_codec,
- ipc_write_context,
- write_options,
- )?;
- append_variadic_buffer_counts(&mut variadic_buffer_counts,
&array_data);
- }
-
- let tail_pad = pad_to_alignment(alignment, offset as usize);
- let body_len = offset as usize + tail_pad;
+ let EncodedRecordBatchMeta {
+ fb_offset: record_batch,
+ body_len,
+ } = self.encode_record_batch_data(
+ batch.columns().iter().map(|array| array.to_data()),
+ batch.num_rows() as i64,
+ write_options,
+ ipc_write_context,
+ sink,
Review Comment:
a possible issue is the number of function calls added. @JakeDern could you
try placing an inline statement above `encode_record_batch_data`.
--
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]
