This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new ea71b0a223 [arrow-flight encode path]re-use flatbufferbuilder (#10220)
ea71b0a223 is described below
commit ea71b0a2238725376d710a23131e7e62f15cd3c6
Author: RIchard Baah <[email protected]>
AuthorDate: Mon Jun 29 10:52:53 2026 -0400
[arrow-flight encode path]re-use flatbufferbuilder (#10220)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- smaller piece of #10137
- see comment
https://github.com/apache/arrow-rs/pull/10137#pullrequestreview-4581920672
# Rationale for this change
flat buffer builder was being allocated repeatedly when it could be
created once and reset using `fbb.reset`
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
provides methods for setting and getting fbb. This avoids needing to
re-allocate it on every call
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
arrow-ipc/src/compression.rs | 9 ++++++++-
arrow-ipc/src/writer.rs | 33 +++++++++++++++++----------------
2 files changed, 25 insertions(+), 17 deletions(-)
diff --git a/arrow-ipc/src/compression.rs b/arrow-ipc/src/compression.rs
index 45558adb44..fd8469c2ba 100644
--- a/arrow-ipc/src/compression.rs
+++ b/arrow-ipc/src/compression.rs
@@ -18,6 +18,7 @@
use crate::CompressionType;
use arrow_buffer::Buffer;
use arrow_schema::ArrowError;
+use flatbuffers::FlatBufferBuilder;
const LENGTH_NO_COMPRESSED_DATA: i64 = -1;
const LENGTH_OF_PREFIX_DATA: i64 = 8;
@@ -27,14 +28,20 @@ const DEFAULT_ZSTD_COMPRESSION_LEVEL: i32 = 3;
///
/// In the case of zstd, this will contain the zstd context, which can be
reused between subsequent
/// compression calls to avoid the performance overhead of initialising a new
context for every
-/// compression.
+/// compression. Also holds a [`FlatBufferBuilder`] that is reused across IPC
writes.
#[derive(Default)]
pub struct CompressionContext {
+ fbb: FlatBufferBuilder<'static>,
#[cfg(feature = "zstd")]
compressor: Option<zstd::bulk::Compressor<'static>>,
}
impl CompressionContext {
+ /// Get a mutable reference to the [`FlatBufferBuilder`] that is reused
across IPC writes.
+ pub(crate) fn mut_fbb(&mut self) -> &mut FlatBufferBuilder<'static> {
+ &mut self.fbb
+ }
+
#[cfg(feature = "zstd")]
fn zstd_compressor(&mut self, level: i32) -> &mut
zstd::bulk::Compressor<'static> {
self.compressor.get_or_insert_with(|| {
diff --git a/arrow-ipc/src/writer.rs b/arrow-ipc/src/writer.rs
index 9ab86807d4..b88abdec7e 100644
--- a/arrow-ipc/src/writer.rs
+++ b/arrow-ipc/src/writer.rs
@@ -754,12 +754,11 @@ impl IpcDataGenerator {
compression_context: &mut CompressionContext,
sink: &mut IpcBodySink<'_>,
) -> Result<(Vec<u8>, usize, usize), ArrowError> {
- let mut fbb = FlatBufferBuilder::new();
-
let batch_compression_type = write_options.batch_compression_type;
let compression = batch_compression_type.map(|batch_compression_type| {
- let mut c = crate::BodyCompressionBuilder::new(&mut fbb);
+ let fbb = compression_context.mut_fbb();
+ let mut c = crate::BodyCompressionBuilder::new(fbb);
c.add_method(crate::BodyCompressionMethod::BUFFER);
c.add_codec(batch_compression_type);
c.finish()
@@ -797,6 +796,7 @@ impl IpcDataGenerator {
let tail_pad = pad_to_alignment(alignment, offset as usize);
let body_len = offset as usize + tail_pad;
+ let fbb = compression_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() {
@@ -806,7 +806,7 @@ impl IpcDataGenerator {
};
let root = {
- let mut batch_builder = crate::RecordBatchBuilder::new(&mut fbb);
+ 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);
@@ -818,8 +818,7 @@ impl IpcDataGenerator {
}
batch_builder.finish().as_union_value()
};
- // create an crate::Message
- let mut message = crate::MessageBuilder::new(&mut fbb);
+ 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);
@@ -827,7 +826,9 @@ impl IpcDataGenerator {
let root = message.finish();
fbb.finish(root, None);
- Ok((fbb.finished_data().to_vec(), body_len, tail_pad))
+ let ipc_message = fbb.finished_data().to_vec();
+ fbb.reset();
+ Ok((ipc_message, body_len, tail_pad))
}
/// Write dictionary values into two sets of bytes, one for the header
(crate::Message) and the
@@ -840,15 +841,14 @@ impl IpcDataGenerator {
is_delta: bool,
compression_context: &mut CompressionContext,
) -> Result<EncodedData, ArrowError> {
- let mut fbb = FlatBufferBuilder::new();
-
let mut arrow_data: Vec<u8> = vec![];
// get the type of compression
let batch_compression_type = write_options.batch_compression_type;
let compression = batch_compression_type.map(|batch_compression_type| {
- let mut c = crate::BodyCompressionBuilder::new(&mut fbb);
+ let fbb = compression_context.mut_fbb();
+ let mut c = crate::BodyCompressionBuilder::new(fbb);
c.add_method(crate::BodyCompressionMethod::BUFFER);
c.add_codec(batch_compression_type);
c.finish()
@@ -885,7 +885,7 @@ impl IpcDataGenerator {
let body_len = offset as usize + tail_pad;
arrow_data.extend_from_slice(&PADDING[..tail_pad]);
- // write data
+ let fbb = compression_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() {
@@ -895,7 +895,7 @@ impl IpcDataGenerator {
};
let root = {
- let mut batch_builder = crate::RecordBatchBuilder::new(&mut fbb);
+ let mut batch_builder = crate::RecordBatchBuilder::new(fbb);
batch_builder.add_length(array_data.len() as i64);
batch_builder.add_nodes(nodes);
batch_builder.add_buffers(buffers);
@@ -909,7 +909,7 @@ impl IpcDataGenerator {
};
let root = {
- let mut batch_builder = crate::DictionaryBatchBuilder::new(&mut
fbb);
+ let mut batch_builder = crate::DictionaryBatchBuilder::new(fbb);
batch_builder.add_id(dict_id);
batch_builder.add_data(root);
batch_builder.add_isDelta(is_delta);
@@ -917,7 +917,7 @@ impl IpcDataGenerator {
};
let root = {
- let mut message_builder = crate::MessageBuilder::new(&mut fbb);
+ let mut message_builder = crate::MessageBuilder::new(fbb);
message_builder.add_version(write_options.metadata_version);
message_builder.add_header_type(crate::MessageHeader::DictionaryBatch);
message_builder.add_bodyLength(body_len as i64);
@@ -926,10 +926,11 @@ impl IpcDataGenerator {
};
fbb.finish(root, None);
- let finished_data = fbb.finished_data();
+ let ipc_message = fbb.finished_data().to_vec();
+ fbb.reset();
Ok(EncodedData {
- ipc_message: finished_data.to_vec(),
+ ipc_message,
arrow_data,
})
}
