Rich-T-kid commented on code in PR #10128:
URL: https://github.com/apache/arrow-rs/pull/10128#discussion_r3666450257
##########
arrow-ipc/src/writer.rs:
##########
@@ -333,6 +325,28 @@ impl Default for IpcWriteOptions {
/// [Arrow IPC Format]:
https://arrow.apache.org/docs/format/Columnar.html#serialization-and-interprocess-communication-ipc
pub struct IpcDataGenerator {}
+/// A dictionary determined to need sending as a part of a record batch.
+struct DictionaryToEncode {
+ id: i64,
+ data: ArrayData,
+ is_delta: bool,
+}
+
+/// The result of encoding a record batch's body and flatbuffer `RecordBatch`
table
+/// via [`IpcDataGenerator::encode_record_batch_data`].
+///
+/// The `RecordBatch` table is left in-progress in the caller's
[`FlatBufferBuilder`];
+/// the caller wraps `record_batch` in a `Message` (directly for a record batch
+/// message, or inside a `DictionaryBatch` for a dictionary message).
+struct EncodedRecordBatch {
Review Comment:
nit: this holds just metadata, mabey `EncodedRecordBatchMetaData` or
`EncodedBatchMetaData` or something of the sort would fit better here
##########
arrow-ipc/src/writer.rs:
##########
@@ -746,105 +749,143 @@ impl IpcDataGenerator {
/// Encodes a `RecordBatch` into a flatbuffer IPC message and fills `sink`
with the
/// serialised buffer data.
///
- /// Returns `(ipc_message, body_len, tail_pad)`: the flatbuffer header
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 message header 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;
-
- 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);
- }
+ ) -> Result<usize, ArrowError> {
+ // Reuse the builder's internal buffer across messages; `reset` keeps
the
+ // allocated capacity and only clears the in-progress state.
+ ipc_write_context.mut_fbb().reset();
- let tail_pad = pad_to_alignment(alignment, offset as usize);
- let body_len = offset as usize + tail_pad;
+ let EncodedRecordBatch {
+ 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,
+ )?;
+ // create an crate::Message
Review Comment:
nit: we can remove this
##########
arrow-ipc/src/writer.rs:
##########
@@ -525,16 +528,17 @@ impl IpcDataGenerator {
Ok(())
}
- #[allow(clippy::too_many_arguments)]
Review Comment:
nice
##########
arrow-ipc/src/writer.rs:
##########
@@ -746,105 +749,143 @@ impl IpcDataGenerator {
/// Encodes a `RecordBatch` into a flatbuffer IPC message and fills `sink`
with the
/// serialised buffer data.
///
- /// Returns `(ipc_message, body_len, tail_pad)`: the flatbuffer header
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 message header 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;
-
- 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);
- }
+ ) -> Result<usize, ArrowError> {
+ // Reuse the builder's internal buffer across messages; `reset` keeps
the
+ // allocated capacity and only clears the in-progress state.
+ ipc_write_context.mut_fbb().reset();
- let tail_pad = pad_to_alignment(alignment, offset as usize);
- let body_len = offset as usize + tail_pad;
+ let EncodedRecordBatch {
+ 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,
+ )?;
+ // create an crate::Message
let fbb = ipc_write_context.mut_fbb();
- let buffers = fbb.create_vector(&meta.buffers);
- let nodes = fbb.create_vector(&meta.nodes);
- let variadic_buffer = if variadic_buffer_counts.is_empty() {
- None
- } else {
- Some(fbb.create_vector(&variadic_buffer_counts))
- };
-
- let root = {
- let mut batch_builder = crate::RecordBatchBuilder::new(fbb);
- batch_builder.add_length(batch.num_rows() as i64);
- batch_builder.add_nodes(nodes);
- batch_builder.add_buffers(buffers);
- if let Some(c) = compression {
- batch_builder.add_compression(c);
- }
- if let Some(v) = variadic_buffer {
- batch_builder.add_variadicBufferCounts(v);
- }
- batch_builder.finish().as_union_value()
- };
let mut message = crate::MessageBuilder::new(fbb);
message.add_version(write_options.metadata_version);
message.add_header_type(crate::MessageHeader::RecordBatch);
message.add_bodyLength(body_len as i64);
- message.add_header(root);
+ message.add_header(record_batch.as_union_value());
let root = message.finish();
fbb.finish(root, None);
- let ipc_message = fbb.finished_data().to_vec();
- fbb.reset();
- Ok((ipc_message, body_len, tail_pad))
+ Ok(body_len)
}
/// Write dictionary values into two sets of bytes, one for the header
(crate::Message) and the
/// other for the data
fn dictionary_batch_to_bytes(
&self,
- dict_id: i64,
- array_data: &ArrayData,
+ dict: DictionaryToEncode,
write_options: &IpcWriteOptions,
- is_delta: bool,
- ipc_write_context: &mut IpcWriteContext,
+ compression_context: &mut IpcWriteContext,
) -> Result<EncodedData, ArrowError> {
let mut arrow_data: Vec<u8> = vec![];
+ self.dictionary_batch_to_sink(
+ &dict,
+ write_options,
+ compression_context,
+ &mut IpcBodySink::Write(&mut arrow_data),
+ )?;
- // get the type of compression
+ Ok(EncodedData {
+ ipc_message:
compression_context.mut_fbb().finished_data().to_vec(),
+ arrow_data,
+ })
+ }
+
+ /// Encodes a dictionary batch's flatbuffer metadata into
`compression_context`'s
+ /// builder (read back via
`compression_context.mut_fbb().finished_data()`) and fills
+ /// `sink` with its body buffers.
+ ///
+ /// Returns the total body length written to `sink` (including per-buffer
alignment
+ /// padding).
+ fn dictionary_batch_to_sink(
+ &self,
+ dict: &DictionaryToEncode,
+ write_options: &IpcWriteOptions,
+ compression_context: &mut IpcWriteContext,
+ sink: &mut IpcBodySink<'_>,
+ ) -> Result<usize, ArrowError> {
+ // Reuse the builder's internal buffer across messages; `reset` keeps
the
+ // allocated capacity and only clears the in-progress state.
Review Comment:
nit: this comment appear a few times in separate places. I think leaving it
on just the struct is fine
##########
arrow-ipc/src/writer.rs:
##########
@@ -1836,38 +1885,66 @@ pub fn write_message<W: Write>(
));
}
- let a = usize::from(write_options.alignment - 1);
- let buffer = encoded.ipc_message;
- let flatbuf_size = buffer.len();
+ let aligned_size = write_message_header(&mut writer, &encoded.ipc_message,
write_options)?;
+
+ // 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))
+}
+
+/// Write the encapsulated-message header to `writer`: the continuation marker
and
+/// metadata length, the flatbuffer metadata (`ipc_message`), and the alignment
+/// padding that follows it.
+///
+/// Returns the padded header length (continuation prefix + metadata +
padding),
+/// which is also the value encoded in the continuation length field.
+fn write_message_header<W: Write>(
+ writer: &mut W,
+ ipc_message: &[u8],
+ write_options: &IpcWriteOptions,
+) -> Result<usize, ArrowError> {
+ let flatbuf_size = ipc_message.len();
+
+ // Formats prior to V5 did not have continuation bytes, so 4 bytes for just
+ // the metadata len for the legacy case and an additional 4 for non legacy
+ // formats.
let prefix_size = if write_options.write_legacy_ipc_format {
4
} else {
8
};
+
+ // Round the total length of the message up to the desired alignment e.g.
+ //
+ // alignment = 8 = 0b0000_1000
+ // a = 8 - 1 = 7 = 0b000_0111 <- Lower 3 bits = 0 implies 8 byte aligned
+ //
+ // Masking with `a` rounds any value down to the next byte boundary, so we
+ // add `a` first to push the value past the next boundary and then round
+ // down to it and therefore round up.
Review Comment:
nice! thank you
##########
arrow-ipc/src/compression.rs:
##########
@@ -84,6 +84,7 @@ impl std::fmt::Debug for IpcWriteContext {
}
/// Deprecated alias for [`IpcWriteContext`].
+#[expect(dead_code)]
Review Comment:
yea I'm not sure. shouldn't be a huge deal
##########
arrow-ipc/src/writer.rs:
##########
@@ -857,34 +898,34 @@ impl IpcDataGenerator {
let batch_compression_level = write_options.batch_compression_level;
let compression_codec: Option<CompressionCodec> =
batch_compression_type
- .map(|batch_compression_type| match batch_compression_level {
+ .map(|compression_type| match batch_compression_level {
Some(level) => {
-
CompressionCodec::try_new_with_compression_level(batch_compression_type, level)
+
CompressionCodec::try_new_with_compression_level(compression_type, level)
}
- None => batch_compression_type.try_into(),
+ None => compression_type.try_into(),
})
.transpose()?;
- let alignment = write_options.alignment;
let mut meta = IpcMetadataBuilder::default();
- let mut sink = IpcBodySink::Write(&mut arrow_data);
- let offset = write_array_data(
- array_data,
- &mut meta,
- &mut sink,
- 0,
- compression_codec,
- ipc_write_context,
- write_options,
- )?;
+ let mut variadic_buffer_counts: Vec<i64> = vec![];
+ let mut offset = 0i64;
- let mut variadic_buffer_counts = vec![];
- append_variadic_buffer_counts(&mut variadic_buffer_counts, array_data);
+ for array_data in columns {
+ 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);
+ }
- // pad the tail of body data
- let tail_pad = pad_to_alignment(alignment, offset as usize);
- let body_len = offset as usize + tail_pad;
- arrow_data.extend_from_slice(&PADDING[..tail_pad]);
+ // Each buffer is padded to the alignment as it is written, so
`offset` is
+ // already a multiple of the alignment -- the body needs no trailing
padding.
+ let body_len = offset as usize;
Review Comment:
It would be nice to add a `debug_assert!` here to ensure the `body_len` is a
multiple of the alignment
--
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]
