JakeDern commented on code in PR #10128:
URL: https://github.com/apache/arrow-rs/pull/10128#discussion_r3839985479
##########
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:
Sure, done!
--
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]
