[Variant] Avoid extra allocation in object builder #7935

klion26 · 2025-07-16T06:10:13Z

Which issue does this PR close?

Closes [Variant] Avoid extra allocation in ObjectBuilder #7899 .

This pr wants to avoid the extra allocation for the object builder and the later buffer copy.

Rationale for this change

Avoid extra allocation in the object builder like the issue descripted.

What changes are included in this PR?

This removes the internal buffer in ObjectBuilder. All data insertion is done directly to the parent buffer wrapped in parent_state.

The corresponding new fields are added to ObjectBuilder.

add object_start_offset in ObjectBuilder, which describes the start offset in the parent buffer for the current object
Add has_been_finished in ObjectBuilder, which describes whether the current object has been finished; it will be used in the Drop function.

This patch modifies the logic of new, finish, parent_state, and drop function according to the change.

In particular, it writes data into the parent buffer directly when adding a field to the object (i.e., insert/try_insert is called). When finalizing (finish is called) the object, as header and field ids are must be put in front of data in the buffer, the builder will shift written data bytes for the necessary space for header and field ids. Then it writes header and field ids.

In drop, if the builder is not finalized before being dropped, it will truncate the written bytes to roll back the parent buffer status.

Are these changes tested?

The logic has been covered by the exist logic.

Are there any user-facing changes?

No

klion26 · 2025-07-16T06:12:07Z

@alamb Please help to review this when you have time, thanks.

klion26 · 2025-07-16T06:22:19Z

parquet-variant/src/builder.rs

-        (state, self.validate_unique_fields)
+        let validate_unique_fields = self.validate_unique_fields;
+
+        match &mut self.parent_state {


Did not find a better solution for this. I can change this if there is a better solution.

I think we can use the pattern in this pR: klion26#1

Fixed, it becomes cleaner now.

klion26 · 2025-07-16T06:57:25Z

~~The test in builder.rs completed successfully, will investigate why CI fails.~~
pushed a fixup to fix the failed CI.

alamb

Thank you for this PR @klion26 -- I plan to review it carefully tomorrow

This commit will reuse the parent buffer for object builder. It can avoid the extra allocation for the object and the later buffer copy.

trigger ci

klion26 · 2025-07-17T03:50:41Z

parquet-variant/src/builder.rs

@@ -1064,20 +1084,58 @@ impl<'a> ObjectBuilder<'a> {
        key: &str,
        value: T,
    ) -> Result<(), ArrowError> {
-        // Get metadata_builder from parent state
-        let metadata_builder = self.parent_state.metadata_builder();
+        match &mut self.parent_state {


same as below

Here is a proposal of how to avoid the duplication:

Introduce buffer_and_metadata_builder to avoid duplication klion26/arrow-rs#1

klion26 · 2025-07-17T04:08:29Z

@alamb thank you! I've rebased on the main branch, and hardened the nestes object test.

alamb

Thank you @klion26 -- this looks quite cool.

I think we should try and avoid the replication when possible -- I left a suggestion on how to do that and get the compiler to be happy

Likewise I think it would be nice to add a few more tests. Let me know if it makes sense

alamb · 2025-07-17T14:09:28Z

parquet-variant/src/builder.rs

+        assert_eq!(inner_inner_object_d.len(), 1);
+        assert_eq!(inner_inner_object_d.field_name(0).unwrap(), "cc");
+        assert_eq!(inner_inner_object_d.field(0).unwrap(), Variant::from("dd"));
+
        assert_eq!(outer_object.field_name(1).unwrap(), "b");
        assert_eq!(outer_object.field(1).unwrap(), Variant::from(true));
    }


I think we should also add tests for the rollback behavior (as in starting an ObjectBuilder but not calling finish)

Similar we should test a list builder rollback too

Okay, I will add more tests for this. Currently, the list inside the object depends on the from_json test, add a unit test for it in builder.rs is better, I'll add this.

Not sure if the tests like test_xx_no_finishi()(such as test_object_builder_to_list_builder_inner_no_finish()) are enough to cover the rollback logic?

We've called drop in the test_xx_no_finish() test, do we need to call drop if we add tests to cover the rollaback logic?

After reviewing the tests, I agree that the existing tests are sufficient

arrow-rs/parquet-variant/src/builder.rs

Lines 2837 to 2918 in 0fecaa0

fn test_object_builder_to_list_builder_outer_no_finish() {

let mut builder = VariantBuilder::new();

let mut object_builder = builder.new_object();

object_builder.insert("first", 1i8);

// Create a nested list builder and finish it

let mut nested_list_builder = object_builder.new_list("nested");

nested_list_builder.append_value("hi");

nested_list_builder.finish();

// Drop the outer object builder without finishing it

drop(object_builder);

builder.append_value(2i8);

// Only the second attempt should appear in the final variant

let (metadata, value) = builder.finish();

let metadata = VariantMetadata::try_new(&metadata).unwrap();

assert_eq!(metadata.len(), 2);

assert_eq!(&metadata[0], "first");

assert_eq!(&metadata[1], "nested"); // not rolled back

let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();

assert_eq!(variant, Variant::Int8(2));

}

#[test]

fn test_object_builder_to_object_builder_inner_no_finish() {

let mut builder = VariantBuilder::new();

let mut object_builder = builder.new_object();

object_builder.insert("first", 1i8);

// Create a nested object builder but never finish it

let mut nested_object_builder = object_builder.new_object("nested");

nested_object_builder.insert("name", "unknown");

drop(nested_object_builder);

object_builder.insert("second", 2i8);

// The parent object should only contain the original fields

object_builder.finish().unwrap();

let (metadata, value) = builder.finish();

let metadata = VariantMetadata::try_new(&metadata).unwrap();

assert_eq!(metadata.len(), 3);

assert_eq!(&metadata[0], "first");

assert_eq!(&metadata[1], "name"); // not rolled back

assert_eq!(&metadata[2], "second");

let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();

let obj = variant.as_object().unwrap();

assert_eq!(obj.len(), 2);

assert_eq!(obj.get("first"), Some(Variant::Int8(1)));

assert_eq!(obj.get("second"), Some(Variant::Int8(2)));

}

#[test]

fn test_object_builder_to_object_builder_outer_no_finish() {

let mut builder = VariantBuilder::new();

let mut object_builder = builder.new_object();

object_builder.insert("first", 1i8);

// Create a nested object builder and finish it

let mut nested_object_builder = object_builder.new_object("nested");

nested_object_builder.insert("name", "unknown");

nested_object_builder.finish().unwrap();

// Drop the outer object builder without finishing it

drop(object_builder);

builder.append_value(2i8);

// Only the second attempt should appear in the final variant

let (metadata, value) = builder.finish();

let metadata = VariantMetadata::try_new(&metadata).unwrap();

assert_eq!(metadata.len(), 3);

assert_eq!(&metadata[0], "first"); // not rolled back

assert_eq!(&metadata[1], "name"); // not rolled back

assert_eq!(&metadata[2], "nested"); // not rolled back

let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();

assert_eq!(variant, Variant::Int8(2));

}

I also verified test coverage with

cargo llvm-cov --html test -p parquet-variant

And it is indeed covered:

alamb · 2025-07-17T14:10:36Z

parquet-variant/src/builder.rs

@@ -999,8 +1000,17 @@ impl<'a> ListBuilder<'a> {
        let offset_size = int_size(data_size);

        // Get parent's buffer
+        let offset_shift = match &self.parent_state {


It might be nice if this was a function in ParentState, something like

Suggested change

let offset_shift = match &self.parent_state {

let offset_shift = self.parent_state.object_start_offset();

alamb · 2025-07-17T14:17:33Z

parquet-variant/src/builder.rs

@@ -1064,20 +1084,58 @@ impl<'a> ObjectBuilder<'a> {
        key: &str,
        value: T,
    ) -> Result<(), ArrowError> {
-        // Get metadata_builder from parent state
-        let metadata_builder = self.parent_state.metadata_builder();
+        match &mut self.parent_state {


Here is a proposal of how to avoid the duplication:

Introduce buffer_and_metadata_builder to avoid duplication klion26/arrow-rs#1

alamb · 2025-07-17T14:18:10Z

parquet-variant/src/builder.rs

+        let start_offset = match &parent_state {
+            ParentState::Variant { buffer, .. } => buffer.offset(),
+            ParentState::List { buffer, .. } => buffer.offset(),
+            ParentState::Object { buffer, .. } => buffer.offset(),
+        };


Suggested change

let start_offset = match &parent_state {

ParentState::Variant { buffer, .. } => buffer.offset(),

ParentState::List { buffer, .. } => buffer.offset(),

ParentState::Object { buffer, .. } => buffer.offset(),

};

let start_offset = parent_state.buffer().offset();

Tried this way, but it needs to change the parent_state to mutable. Added a function to retrieve the current offset of the buffer.

alamb · 2025-07-17T14:18:28Z

parquet-variant/src/builder.rs

-        (state, self.validate_unique_fields)
+        let validate_unique_fields = self.validate_unique_fields;
+
+        match &mut self.parent_state {


I think we can use the pattern in this pR: klion26#1

alamb · 2025-07-17T14:19:01Z

parquet-variant/src/builder.rs

-        let data_size = self.buffer.offset();
-        let num_fields = self.fields.len();
-        let is_large = num_fields > u8::MAX as usize;
+        let metadata_builder = match &self.parent_state {


it would be nice to put this into a method as well rather than an inline match statement

alamb · 2025-07-17T14:19:49Z

parquet-variant/src/builder.rs

+
+        let starting_offset = self.object_start_offset;
+
+        // Shift existing data to make room for the header


As a follow on PR we can consider avoiding this extra splice somehow (by preallocating the size or something). Future work though

filed an issue(#7960) to trace this.

alamb · 2025-07-17T14:20:31Z

parquet-variant/src/builder.rs

+        buffer[header_pos..header_pos + offset_size as usize]
+            .copy_from_slice(&data_size_bytes[..offset_size as usize]);
+
+        let start_offset_shift = match &self.parent_state {


here is another place we could use the method and avoid an inline match

klion26 · 2025-07-18T02:11:26Z

@alamb Thanks for the detailed review, will adressed them soon.

klion26

@alamb Thanks for the detailed review and the suggestion, the code becomes much more cleaner. I've addressed most of the comments(in commit f5b0465), will add the test to cover the rollback logic after confirmation.

klion26 · 2025-07-18T02:30:21Z

parquet-variant/src/builder.rs

+        let start_offset = match &parent_state {
+            ParentState::Variant { buffer, .. } => buffer.offset(),
+            ParentState::List { buffer, .. } => buffer.offset(),
+            ParentState::Object { buffer, .. } => buffer.offset(),
+        };


Tried this way, but it needs to change the parent_state to mutable. Added a function to retrieve the current offset of the buffer.

klion26 · 2025-07-18T05:00:28Z

parquet-variant/src/builder.rs

+
+        let starting_offset = self.object_start_offset;
+
+        // Shift existing data to make room for the header


filed an issue(#7960) to trace this.

klion26 · 2025-07-18T05:04:55Z

parquet-variant/src/builder.rs

+        buffer[header_pos..header_pos + offset_size as usize]
+            .copy_from_slice(&data_size_bytes[..offset_size as usize]);
+
+        let start_offset_shift = match &self.parent_state {


klion26 · 2025-07-18T05:05:04Z

parquet-variant/src/builder.rs

-        let data_size = self.buffer.offset();
-        let num_fields = self.fields.len();
-        let is_large = num_fields > u8::MAX as usize;
+        let metadata_builder = match &self.parent_state {


klion26 · 2025-07-18T05:05:39Z

parquet-variant/src/builder.rs

-        (state, self.validate_unique_fields)
+        let validate_unique_fields = self.validate_unique_fields;
+
+        match &mut self.parent_state {


Fixed, it becomes cleaner now.

klion26 · 2025-07-18T05:05:53Z

parquet-variant/src/builder.rs

@@ -1064,20 +1084,58 @@ impl<'a> ObjectBuilder<'a> {
        key: &str,
        value: T,
    ) -> Result<(), ArrowError> {
-        // Get metadata_builder from parent state
-        let metadata_builder = self.parent_state.metadata_builder();
+        match &mut self.parent_state {


klion26 · 2025-07-18T05:06:09Z

parquet-variant/src/builder.rs

@@ -999,8 +1000,17 @@ impl<'a> ListBuilder<'a> {
        let offset_size = int_size(data_size);

        // Get parent's buffer
+        let offset_shift = match &self.parent_state {


alamb · 2025-07-18T15:48:43Z

🤖 ./gh_compare_arrow.sh Benchmark Script Running
Linux aal-dev 6.11.0-1016-gcp #16~24.04.1-Ubuntu SMP Wed May 28 02:40:52 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing 7899-avoid-extra-allocation-in-object-builder (0fecaa0) to 99eb1bc diff
BENCH_NAME=variant_kernels
BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental --bench variant_kernels
BENCH_FILTER=
BENCH_BRANCH_NAME=7899-avoid-extra-allocation-in-object-builder
Results will be posted here when complete

alamb · 2025-07-18T15:52:48Z

🤖: Benchmark completed

Details

group                                                                7899-avoid-extra-allocation-in-object-builder    main
-----                                                                ---------------------------------------------    ----
batch_json_string_to_variant json_list 8k string                     1.00     27.9±0.10ms        ? ?/sec              1.03     28.7±0.11ms        ? ?/sec
batch_json_string_to_variant random_json(2633 bytes per document)    1.00    364.7±8.44ms        ? ?/sec              1.08    392.4±6.39ms        ? ?/sec
batch_json_string_to_variant repeated_struct 8k string               1.01      8.6±0.02ms        ? ?/sec              1.00      8.6±0.03ms        ? ?/sec
variant_get_primitive                                                1.00  1368.4±16.34µs        ? ?/sec              1.01   1375.4±8.13µs        ? ?/sec

alamb

This is really nice @klion26 -- thank you very much for the contribution.

It is neat to see we can already see it making conersion from JSON 8% faster in some cases (see benchmark here #7935 (comment)):

batch_json_string_to_variant random_json(2633 bytes per document)    1.00    364.7±8.44ms        ? ?/sec              1.08    392.4±6.39ms        ? ?/sec

FYI @scovich @viirya and @friendlymatthew

alamb · 2025-07-18T15:50:37Z

parquet-variant/src/builder.rs

+
+        // Shift existing data to make room for the header
+        let buffer = parent_buffer.inner_mut();
+        buffer.splice(starting_offset..starting_offset, vec![0u8; header_size]);


TIL: VEC::splice

https://doc.rust-lang.org/std/vec/struct.Vec.html#method.splice

Nice

Can avoid even this allocation:

Suggested change

buffer.splice(starting_offset..starting_offset, vec![0u8; header_size]);

buffer.splice(starting_offset..starting_offset, std::iter::repeat_n(0u8, header_size));

Alternatively -- that small allocation apparently isn't hurting performance. What if we create the temp vec with initial capacity of header_size and then populate it with the actual header info, field ids, and offsets before splicing it in. That way, an incorrect header_size calculation would not impact correctness:

let mut bytes_to_splice = Vec::with_capacity(header_size); bytes_to_splice.push(header_byte); if is_large { bytes_to_splice.extend((num_fields as u32).to_le_bytes()); } else { bytes_to_splice.push(num_fields as u8); } for &field_id in self.fields.keys() { bytes_to_splice.extend(field_id.to_le_bytes().into_iter().take(id_size)); } for &offset in self.fields.values() { bytes_to_splice.extend((offset as u32).to_le_bytes().into_iter().take(offset_size)); } bytes_to_splice.extend((data_size as u32).to_le_bytes().into_iter().take(offset_size)); buffer.splice(starting_offset..starting_offset, bytes_to_splice);

aside: I'd be very impressed if the compiler is smart enough to optimize away the allocation in vec![0u8; header_size].into_iter(), which would potentially cause the above to run slower than the current code.

FYI, I think we can improve the performance of appending packed u32 values to a Vec<u8>. The following code:

pub fn append_bytes_brute2(dest: &mut Vec<u8>, src: u32, nbytes: NumBytes) { let n = dest.len() + nbytes as usize; dest.extend(src.to_le_bytes()); dest.truncate(n); }

... unconditionally adds all four bytes and then truncates the vector to the desired length. This works because (a) it's a single machine instruction to copy the four LE bytes of a u32 to a memory location; (b) truncate drops the higher order bytes, when working with LE bytes; and (c) truncate is dirt cheap (conditional move that doesn't break processor pipelines the way a branch would).

Result: Code that should perform the same regardless of the number of bytes we encode the u32 as, with only a single branch to check for vector capacity.

The above compiles to the following assembly code:

playground::append_packed_u32: ;; set up stack frame (not relevant in practice due to inlining) ... movq (%rdi), %rcx ; dest capacity movq 16(%rdi), %rbx ; dest len subq %rbx, %rcx ; calculate available capacity movq %rbx, %rax cmpq $3, %rcx ; if insufficient capacity... jbe .LBB2_1 ; ... then `reserve` more .LBB2_2: addq %rdx, %rbx ; add nbytes to len (truncate arg) movq 8(%rdi), %rcx ; dest data pointer movl %esi, (%rcx,%rax) ; append all four bytes of src to the vector addq $4, %rax ; increase len by 4 (u32 size) cmpq %rax, %rbx ; if new len is bigger than truncate arg... cmovbq %rbx, %rax ; ... then truncate len movq %rax, 16(%rdi) ; write back the updated dest len ;; tear down stack frame ... retq .LBB2_1: ;; call prologue ... callq alloc::raw_vec::RawVecInner<A>::reserve::do_reserve_and_handle ;; call epilogue ... jmp .LBB2_2

The one bummer is, when adding a reserve followed by a loop over a slice of source bytes:

pub fn append_packed_u32(dest: &mut Vec<u8>, src: &[u32], nbytes: usize) { dest.reserve((src.len() + 1) * nbytes); for val in src { let n = dest.len() + nbytes; dest.extend(val.to_le_bytes()); dest.truncate(n); } }

... the compiler isn't smart enough to eliminate the redundant capacity check inside the loop. Oh, well.

alamb · 2025-07-18T15:51:17Z

parquet-variant/src/builder.rs

@@ -1317,7 +1414,15 @@ impl<'a> ObjectBuilder<'a> {
 /// This is to ensure that the object is always finalized before its parent builder
 /// is finalized.
 impl Drop for ObjectBuilder<'_> {
-    fn drop(&mut self) {}
+    fn drop(&mut self) {


alamb · 2025-07-18T15:58:59Z

parquet-variant/src/builder.rs

+        assert_eq!(inner_inner_object_d.len(), 1);
+        assert_eq!(inner_inner_object_d.field_name(0).unwrap(), "cc");
+        assert_eq!(inner_inner_object_d.field(0).unwrap(), Variant::from("dd"));
+
        assert_eq!(outer_object.field_name(1).unwrap(), "b");
        assert_eq!(outer_object.field(1).unwrap(), Variant::from(true));
    }


After reviewing the tests, I agree that the existing tests are sufficient

arrow-rs/parquet-variant/src/builder.rs

Lines 2837 to 2918 in 0fecaa0

fn test_object_builder_to_list_builder_outer_no_finish() {

let mut builder = VariantBuilder::new();

let mut object_builder = builder.new_object();

object_builder.insert("first", 1i8);

// Create a nested list builder and finish it

let mut nested_list_builder = object_builder.new_list("nested");

nested_list_builder.append_value("hi");

nested_list_builder.finish();

// Drop the outer object builder without finishing it

drop(object_builder);

builder.append_value(2i8);

// Only the second attempt should appear in the final variant

let (metadata, value) = builder.finish();

let metadata = VariantMetadata::try_new(&metadata).unwrap();

assert_eq!(metadata.len(), 2);

assert_eq!(&metadata[0], "first");

assert_eq!(&metadata[1], "nested"); // not rolled back

let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();

assert_eq!(variant, Variant::Int8(2));

}

#[test]

fn test_object_builder_to_object_builder_inner_no_finish() {

let mut builder = VariantBuilder::new();

let mut object_builder = builder.new_object();

object_builder.insert("first", 1i8);

// Create a nested object builder but never finish it

let mut nested_object_builder = object_builder.new_object("nested");

nested_object_builder.insert("name", "unknown");

drop(nested_object_builder);

object_builder.insert("second", 2i8);

// The parent object should only contain the original fields

object_builder.finish().unwrap();

let (metadata, value) = builder.finish();

let metadata = VariantMetadata::try_new(&metadata).unwrap();

assert_eq!(metadata.len(), 3);

assert_eq!(&metadata[0], "first");

assert_eq!(&metadata[1], "name"); // not rolled back

assert_eq!(&metadata[2], "second");

let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();

let obj = variant.as_object().unwrap();

assert_eq!(obj.len(), 2);

assert_eq!(obj.get("first"), Some(Variant::Int8(1)));

assert_eq!(obj.get("second"), Some(Variant::Int8(2)));

}

#[test]

fn test_object_builder_to_object_builder_outer_no_finish() {

let mut builder = VariantBuilder::new();

let mut object_builder = builder.new_object();

object_builder.insert("first", 1i8);

// Create a nested object builder and finish it

let mut nested_object_builder = object_builder.new_object("nested");

nested_object_builder.insert("name", "unknown");

nested_object_builder.finish().unwrap();

// Drop the outer object builder without finishing it

drop(object_builder);

builder.append_value(2i8);

// Only the second attempt should appear in the final variant

let (metadata, value) = builder.finish();

let metadata = VariantMetadata::try_new(&metadata).unwrap();

assert_eq!(metadata.len(), 3);

assert_eq!(&metadata[0], "first"); // not rolled back

assert_eq!(&metadata[1], "name"); // not rolled back

assert_eq!(&metadata[2], "nested"); // not rolled back

let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();

assert_eq!(variant, Variant::Int8(2));

}

I also verified test coverage with

cargo llvm-cov --html test -p parquet-variant

And it is indeed covered:

scovich

Initial review, a bit scattered but hopefully helpful.

Impressive that it's already so much faster than the baseline, potentially with room to improve even further!

scovich · 2025-07-18T16:37:27Z

parquet-variant/src/builder.rs

+    // returns the beginning offset of buffer for the parent if it is object builder, else 0.
+    // for object builder will reuse the buffer from the parent, this is needed for `finish`
+    // which needs the relative offset from the current variant.
+    fn object_start_offset(&self) -> usize {


Doesn't variant array also have a starting offset?

This feature that reuse the parent buffer hasn't been implemented for ListBuilder. Will do it with a follow-up pr.

scovich · 2025-07-18T16:39:21Z

parquet-variant/src/builder.rs

+            ParentState::Variant {
+                buffer,
+                metadata_builder,
+            } => (buffer, metadata_builder),
+            ParentState::List {
+                buffer,
+                metadata_builder,
+                ..
+            } => (buffer, metadata_builder),
+            ParentState::Object {
+                buffer,
+                metadata_builder,
+                ..
+            } => (buffer, metadata_builder),


I think rust allows this:

Suggested change

ParentState::Variant {

buffer,

metadata_builder,

} => (buffer, metadata_builder),

ParentState::List {

buffer,

metadata_builder,

..

} => (buffer, metadata_builder),

ParentState::Object {

buffer,

metadata_builder,

..

} => (buffer, metadata_builder),

ParentState::Variant {

buffer,

metadata_builder,

} |

ParentState::List {

buffer,

metadata_builder,

..

} |

ParentState::Object {

buffer,

metadata_builder,

..

} => (buffer, metadata_builder),

Not clear whether that's better or worse than the current code.

I'm also not sure how it would fmt.

scovich · 2025-07-18T16:41:01Z

parquet-variant/src/builder.rs

+        match self {
+            ParentState::Variant { buffer, .. } => buffer.offset(),
+            ParentState::Object { buffer, .. } => buffer.offset(),
+            ParentState::List { buffer, .. } => buffer.offset(),
+        }


As above, can reduce redundancy by

Suggested change

match self {

ParentState::Variant { buffer, .. } => buffer.offset(),

ParentState::Object { buffer, .. } => buffer.offset(),

ParentState::List { buffer, .. } => buffer.offset(),

}

match self {

ParentState::Variant { buffer, .. }

| ParentState::Object { buffer, .. }

| ParentState::List { buffer, .. } => buffer.offset(),

}

Actually... we can just invoke ParentState::buffer?

Suggested change

match self {

ParentState::Variant { buffer, .. } => buffer.offset(),

ParentState::Object { buffer, .. } => buffer.offset(),

ParentState::List { buffer, .. } => buffer.offset(),

}

self.buffer().offset()

This needs the signature changed to buffer_current_offset(&mut self), not sure if this is ok?

changed to the first version

scovich · 2025-07-18T20:33:24Z

parquet-variant/src/builder.rs

+
+        // Shift existing data to make room for the header
+        let buffer = parent_buffer.inner_mut();
+        buffer.splice(starting_offset..starting_offset, vec![0u8; header_size]);


Can avoid even this allocation:

Suggested change

buffer.splice(starting_offset..starting_offset, vec![0u8; header_size]);

buffer.splice(starting_offset..starting_offset, std::iter::repeat_n(0u8, header_size));

scovich · 2025-07-18T21:13:12Z

parquet-variant/src/builder.rs

+        let header_size = 1 + // header byte
+            (if is_large { 4 } else { 1 }) + // num_fields
+            (num_fields * id_size as usize) + // field IDs
+            ((num_fields + 1) * offset_size as usize); // field offsets + data_size


One thing I don't love about this approach is the independent calculation of the splice size vs. the bytes that actually get inserted. If they ever disagreed... badness as the insertions underflow or overflow the splice.

Ideally, we could produce an iterator that emits the desired bytes, and the splice itself can guarantee correct behavior. But for that to work the iterator would need to provide an accurate lower bound, so which rules out std::iter::from_fn. Even if we did craft a custom iterator, computing its lower bound would basically be this same calculation all over again. We could also chain together a bunch of iterators, which preserves the lower bound, but somehow I doubt that would be efficient.

Something like might almost work?

let field_ids = self.fields.keys().flat_map(|field_id| { (field_id as usize).to_le_bytes().into_iter().take(id_size) }); let offsets = self.fields.values().flat_map(|offset| { offset.to_le_bytes().into_iter().take(offset_size) }); let num_fields = num_fields.to_le_bytes().take(if is_large 4 else 1); let header_and_num_fields = std::iter::once(header_byte).chain(num_fields); let field_ids_and_offsets = field_ids.chain(offsets); let bytes_to_splice = header_and_num_fields.chain(field_ids_and_offsets); buffer.splice(starting_offset..starting_offset, bytes_to_splice);

... but unfortunately Iterator::flat_map does not (honestly, cannot) compute an accurate lower bound size hint. So the splice call would end up having to allocate an internal temp buffer.

I just realized -- a custom iterator that computes its size hint is still safer than the current code, because an incorrect size hint won't corrupt any data. It will just require extra allocations and/or byte shifting. There's still the question of performance tho. If it's drastically slower the extra safety probably wouldn't be worth it.

Another possibility is to take a hybrid approach -- rely on Iterator::chain for most of the heavy lifting, but define a custom iterator for emitting offset/field_id arrays:

let num_fields = num_fields.to_le_bytes().take(if is_large 4 else 1); let header_and_num_fields = std::iter::once(header_byte).chain(num_fields); let field_ids = PackedU32Iterator::new(id_size, self.fields.keys().copied()); let offsets = PackedU32Iterator::new(offset_size, self.fields.values().map(|offset| *offset as u32)); let field_ids_and_offsets = field_ids.chain(offsets); let bytes_to_splice = header_and_num_fields.chain(field_ids_and_offsets); buffer.splice(starting_offset..starting_offset, bytes_to_splice);

PackedU32Iterator

struct PackedU32Iterator<T: impl Iterator<Item = [u8; 4]>> { packed_bytes: usize, iterator: T, current_item: [u8; 4], current_byte: usize, // 0..3 } impl<T: impl Iterator<Item = [u8; 4]>> PackedU32Iterator<T> { fn new(packed_bytes: usize, iterator: T) -> Self { // eliminate corner cases in `next` by initializing with a fake already-consumed "first" item Self { packed_bytes, iterator, current_item: [0; 4], current_byte: packed_bytes, } } } impl<T: impl Iterator<Item = [u8; 4]>> Iterator for PackedU32Iterator<T> { fn size_hint(&self) -> (usize, Option<usize>) { let lower = (packed_bytes - current_byte) + packed_bytes * iterator.size_hint().0; (lower, None) } fn next(&mut self) -> Option<u8> { if self.current_byte >= self.packed_bytes { let Some(next_item) = self.iterator.next() else { return None; }; self.current_item = next_item; self.current_byte = 0; } let rval = self.current_item[self.current_byte]; self.current_byte += 1; Some(rval) } }

... but again, I worry that wrapping up all those for-loops inside an iterator for direct splicing will turn out to be a lot slower than the splice-then-overwrite approach the code currently takes.

Do you think it's ok to make this optimization a follow-up pr and add some benchmark for it?

Or try the other approach that populates the tmp vec with non-zero bytes before splicing it into the main buffer? That's a lot simpler than this iterator-based suggestion, and likely performs better too.

ok, will do it in a followup

I also think we can mitigate the danger of independent size calculation with testing.

I like the idea of going with the approach in this PR (which is already pretty large)

scovich · 2025-07-18T21:24:08Z

parquet-variant/src/builder.rs

-        let ids = self.fields.keys().map(|id| *id as usize);
-        parent_buffer.append_offset_array(ids, None, id_size);
+        // Write field IDs
+        for (&field_id, _) in &self.fields {


Suggested change

for (&field_id, _) in &self.fields {

for field_id in self.fields.keys() {

scovich · 2025-07-18T22:08:01Z

parquet-variant/src/builder.rs

+
+        // Shift existing data to make room for the header
+        let buffer = parent_buffer.inner_mut();
+        buffer.splice(starting_offset..starting_offset, vec![0u8; header_size]);


Alternatively -- that small allocation apparently isn't hurting performance. What if we create the temp vec with initial capacity of header_size and then populate it with the actual header info, field ids, and offsets before splicing it in. That way, an incorrect header_size calculation would not impact correctness:

let mut bytes_to_splice = Vec::with_capacity(header_size); bytes_to_splice.push(header_byte); if is_large { bytes_to_splice.extend((num_fields as u32).to_le_bytes()); } else { bytes_to_splice.push(num_fields as u8); } for &field_id in self.fields.keys() { bytes_to_splice.extend(field_id.to_le_bytes().into_iter().take(id_size)); } for &offset in self.fields.values() { bytes_to_splice.extend((offset as u32).to_le_bytes().into_iter().take(offset_size)); } bytes_to_splice.extend((data_size as u32).to_le_bytes().into_iter().take(offset_size)); buffer.splice(starting_offset..starting_offset, bytes_to_splice);

scovich · 2025-07-18T22:09:16Z

parquet-variant/src/builder.rs

+        let start_offset_shift = self.parent_state.object_start_offset();
+        self.parent_state
+            .finish(starting_offset - start_offset_shift);


Why make the caller of ParentState::finish extract the start offset? Seems like parent state can do that more easily and reliably on its own?

Yes, moving this logic to ParentState::finish is clearer. fixed.

scovich · 2025-07-18T22:11:56Z

parquet-variant/src/builder.rs

+        if !self.has_been_finished {
+            self.parent_state
+                .buffer()
+                .inner_mut()
+                .truncate(self.object_start_offset);
+        }


While we're at it, we should truncate the MetadataBuilder back to the size it had when we started, to clean up any new field ids the failed builder might have created.

Seems that in previous pr #7865, some logics want to cover that the metadata hasn't been rolled back in tests like test_list_builder_to_object_builder_inner_no_finish, do we need to roll back the MetadataBuilder here?

// The parent list should only contain the original values list_builder.finish(); let (metadata, value) = builder.finish(); let metadata = VariantMetadata::try_new(&metadata).unwrap(); assert_eq!(metadata.len(), 1); assert_eq!(&metadata[0], "name"); // not rolled back

I updated this into a separate commit, we can keep it or revert it easily.

IIUC, the Metadata can be rolled back because the object has not been written successfully, but not sure if there are any cases I did not follow here.

Yes, existing tests expect the fields to persist because the original code was unable to roll back the changes. With this PR, we should be able to roll back the changes cleanly, and update the unit tests accordingly.

scovich · 2025-07-18T22:42:15Z

parquet-variant/src/builder.rs

+
+        // Shift existing data to make room for the header
+        let buffer = parent_buffer.inner_mut();
+        buffer.splice(starting_offset..starting_offset, vec![0u8; header_size]);


aside: I'd be very impressed if the compiler is smart enough to optimize away the allocation in vec![0u8; header_size].into_iter(), which would potentially cause the above to run slower than the current code.

apache#7774)

klion26

@scovich Thank you very much for the detailed review and suggestion, I've addressed most of them (add comments for the rest), please take another look when you're free.

klion26 · 2025-07-20T11:09:44Z

parquet-variant/src/builder.rs

+    // returns the beginning offset of buffer for the parent if it is object builder, else 0.
+    // for object builder will reuse the buffer from the parent, this is needed for `finish`
+    // which needs the relative offset from the current variant.
+    fn object_start_offset(&self) -> usize {


This feature that reuse the parent buffer hasn't been implemented for ListBuilder. Will do it with a follow-up pr.

klion26 · 2025-07-20T12:08:40Z

parquet-variant/src/builder.rs

+        match self {
+            ParentState::Variant { buffer, .. } => buffer.offset(),
+            ParentState::Object { buffer, .. } => buffer.offset(),
+            ParentState::List { buffer, .. } => buffer.offset(),
+        }


This needs the signature changed to buffer_current_offset(&mut self), not sure if this is ok?

changed to the first version

klion26 · 2025-07-20T12:22:44Z

parquet-variant/src/builder.rs

+        let start_offset_shift = self.parent_state.object_start_offset();
+        self.parent_state
+            .finish(starting_offset - start_offset_shift);


Yes, moving this logic to ParentState::finish is clearer. fixed.

klion26 · 2025-07-20T13:26:12Z

parquet-variant/src/builder.rs

+        if !self.has_been_finished {
+            self.parent_state
+                .buffer()
+                .inner_mut()
+                .truncate(self.object_start_offset);
+        }


Seems that in previous pr #7865, some logics want to cover that the metadata hasn't been rolled back in tests like test_list_builder_to_object_builder_inner_no_finish, do we need to roll back the MetadataBuilder here?

// The parent list should only contain the original values list_builder.finish(); let (metadata, value) = builder.finish(); let metadata = VariantMetadata::try_new(&metadata).unwrap(); assert_eq!(metadata.len(), 1); assert_eq!(&metadata[0], "name"); // not rolled back

I updated this into a separate commit, we can keep it or revert it easily.

IIUC, the Metadata can be rolled back because the object has not been written successfully, but not sure if there are any cases I did not follow here.

klion26 · 2025-07-20T14:40:10Z

parquet-variant/src/builder.rs

+            ParentState::Variant {
+                buffer,
+                metadata_builder,
+            } => (buffer, metadata_builder),
+            ParentState::List {
+                buffer,
+                metadata_builder,
+                ..
+            } => (buffer, metadata_builder),
+            ParentState::Object {
+                buffer,
+                metadata_builder,
+                ..
+            } => (buffer, metadata_builder),


klion26 · 2025-07-20T14:55:32Z

parquet-variant/src/builder.rs

        });
-
        let max_id = self.fields.iter().map(|(i, _)| *i).max().unwrap_or(0);


Using the length of the metadata builder doesn't have a correctness problem, but it will have a tiny problem that will waste some space for the header. added in a separate commit.

klion26 · 2025-07-20T14:59:10Z

parquet-variant/src/builder.rs

+        let header_size = 1 + // header byte
+            (if is_large { 4 } else { 1 }) + // num_fields
+            (num_fields * id_size as usize) + // field IDs
+            ((num_fields + 1) * offset_size as usize); // field offsets + data_size


Do you think it's ok to make this optimization a follow-up pr and add some benchmark for it?

klion26 · 2025-07-20T16:30:01Z

@alamb thank you very much for the review and help, happy to see these improvements.

viirya · 2025-07-20T18:32:51Z

parquet-variant/src/builder.rs

+        // as object builder has been reused the parent buffer,
+        // we need to shift the offset by the starting offset of the parent object


I don't get it why put the comment here, seems ListBuilder.finish hasn't been updated.

Good catch, because if we create a ListBuilder inside an ObjectBuilder, then we need do this shift, but after the last change, the comment here need to be moved to ParentState::finish().

viirya · 2025-07-20T18:33:49Z

parquet-variant/src/builder.rs

@@ -1028,18 +1078,29 @@ impl Drop for ListBuilder<'_> {
 pub struct ObjectBuilder<'a> {
    parent_state: ParentState<'a>,
    fields: IndexMap<u32, usize>, // (field_id, offset)
-    buffer: ValueBuffer,
+    /// the starting offset in the parent's buffer where this object starts


Style suggestion, it'd be better just following existing doc style.

Suggested change

/// the starting offset in the parent's buffer where this object starts

/// The starting offset in the parent's buffer where this object starts

Similar style issues on other comment/doc.

viirya · 2025-07-20T18:45:34Z

parquet-variant/src/builder.rs

        if self.validate_unique_fields && !self.duplicate_fields.is_empty() {
+            let metadata_builder = self.parent_state.metadata_builder();


Hm? Why move metadata_builder inside?

This was intended to avoid the double mutable reference problem, but it's not a problem after the implementation has changed. I can revert this if needed.

I reverted this change in 19bb544 to keep the diff cleaner

viirya · 2025-07-20T18:46:31Z

parquet-variant/src/builder.rs

@@ -506,6 +506,7 @@ enum ParentState<'a> {
        metadata_builder: &'a mut MetadataBuilder,
        fields: &'a mut IndexMap<u32, usize>,
        field_name: &'a str,
+        object_start_offset: usize,


Suggested change

object_start_offset: usize,

parent_offset_base: usize,

Suggested a more understandable name (at least to me).

Fixed, object_start_offset wants to indicate that this is the offset of the current object start.

viirya · 2025-07-20T18:47:14Z

parquet-variant/src/builder.rs

@@ -1028,18 +1078,29 @@ impl Drop for ListBuilder<'_> {
 pub struct ObjectBuilder<'a> {
    parent_state: ParentState<'a>,
    fields: IndexMap<u32, usize>, // (field_id, offset)
-    buffer: ValueBuffer,
+    /// the starting offset in the parent's buffer where this object starts
+    object_start_offset: usize,


Suggested change

object_start_offset: usize,

parent_offset_base: usize,

viirya · 2025-07-20T18:47:36Z

parquet-variant/src/builder.rs

+    object_start_offset: usize,
+    /// the starting offset in the parent's metadata buffer where this object starts
+    /// used to truncate the written fields in `drop` if the current object has not been finished
+    object_meta_start_offset: usize,


Suggested change

object_meta_start_offset: usize,

parent_metadata_offset_base: usize,

viirya · 2025-07-20T18:59:13Z

parquet-variant/src/builder.rs

+        // current object starts from `object_start_offset`
+        let data_size = current_offset - self.object_start_offset;


Hmm, so we assume that no other object was appended at the same time? If we create two object builders and they insert into the parent buffer in an interleaved manner? Won't be the base offset (object start offset) same for them?

Oh, I forgot that parent status has exclusively access to the parent buffer.

viirya · 2025-07-20T19:36:35Z

parquet-variant/src/builder.rs

+        // the length of the metadata's field names is a very cheap to compute the upper bound.
+        // it will almost always be a tight upper bound as well -- it would take a pretty
+        // carefully  crafted object to use only the early field ids of a large dictionary.
+        let max_id = metadata_builder.field_names.len();


Hmm, honestly I didn't get what this comment trying to say. And the max_id is same as the length of field_names exactly?

IIUC, metadata_builder.field_names.len() returns the size of the underlying map, it may be bigger than the actual max field_id, changing to this here is that it can avoid one pass for the field_ids to calucalute the max_id.

But this change can make test more diffict(more difficult to calculate the header size) -- such as the failed case from_json::test::test_json_to_variant_object_very_large (test code here/failed ci)

@alamb @viirya @scovich Maybe we can revert this to the original code(travel one pass from filed_ids) what do you think about this? thanks. -- added a separate commit which reverted this logic to see if any other tests failed.

viirya · 2025-07-20T19:44:11Z

parquet-variant/src/builder.rs

+        // Write header byte
        let header = object_header(is_large, id_size, offset_size);
-        parent_buffer.append_header(header, is_large, num_fields);
+        buffer[header_pos] = header;
+        header_pos += 1;

-        // Write field IDs (sorted order)
-        let ids = self.fields.keys().map(|id| *id as usize);
-        parent_buffer.append_offset_array(ids, None, id_size);
+        // Write number of fields
+        if is_large {
+            buffer[header_pos..header_pos + 4].copy_from_slice(&(num_fields as u32).to_le_bytes());
+            header_pos += 4;
+        } else {
+            buffer[header_pos] = num_fields as u8;
+            header_pos += 1;
+        }


This is basically append_header did before, right? Maybe we can also extract to a function instead of inline here.

viirya · 2025-07-20T20:07:26Z

parquet-variant/src/builder.rs

-        assert_eq!(metadata.len(), 1);
-        assert_eq!(&metadata[0], "name"); // not rolled back
+        assert_eq!(metadata.len(), 1); // rolled back
+        assert_eq!(&metadata[0], "name");


Hmm, why update this to "rolled back"? The metadata is same as before.

Fixed, my bad, this was unintentionally modified.

viirya

Thanks for this contribution. It looks reasonable and the benchmark looks good. I left a few comments. I also updated the PR description to more clearly describe what the changes are to help it understandable to others interested in this.

klion26

@viirya Thanks for the detailed review, I've update the code, please take another look when you're free, thanks.

klion26 · 2025-07-21T01:17:16Z

parquet-variant/src/builder.rs

@@ -506,6 +506,7 @@ enum ParentState<'a> {
        metadata_builder: &'a mut MetadataBuilder,
        fields: &'a mut IndexMap<u32, usize>,
        field_name: &'a str,
+        object_start_offset: usize,


Fixed, object_start_offset wants to indicate that this is the offset of the current object start.

klion26 · 2025-07-21T01:22:43Z

parquet-variant/src/builder.rs

        if self.validate_unique_fields && !self.duplicate_fields.is_empty() {
+            let metadata_builder = self.parent_state.metadata_builder();


This was intended to avoid the double mutable reference problem, but it's not a problem after the implementation has changed. I can revert this if needed.

klion26 · 2025-07-21T01:26:38Z

parquet-variant/src/builder.rs

-        assert_eq!(metadata.len(), 1);
-        assert_eq!(&metadata[0], "name"); // not rolled back
+        assert_eq!(metadata.len(), 1); // rolled back
+        assert_eq!(&metadata[0], "name");


Fixed, my bad, this was unintentionally modified.

klion26 · 2025-07-21T03:16:24Z

parquet-variant/src/builder.rs

@@ -1028,18 +1078,29 @@ impl Drop for ListBuilder<'_> {
 pub struct ObjectBuilder<'a> {
    parent_state: ParentState<'a>,
    fields: IndexMap<u32, usize>, // (field_id, offset)
-    buffer: ValueBuffer,
+    /// the starting offset in the parent's buffer where this object starts


klion26 · 2025-07-21T03:16:32Z

parquet-variant/src/builder.rs

@@ -1028,18 +1078,29 @@ impl Drop for ListBuilder<'_> {
 pub struct ObjectBuilder<'a> {
    parent_state: ParentState<'a>,
    fields: IndexMap<u32, usize>, // (field_id, offset)
-    buffer: ValueBuffer,
+    /// the starting offset in the parent's buffer where this object starts
+    object_start_offset: usize,


klion26 · 2025-07-21T03:16:38Z

parquet-variant/src/builder.rs

+    object_start_offset: usize,
+    /// the starting offset in the parent's metadata buffer where this object starts
+    /// used to truncate the written fields in `drop` if the current object has not been finished
+    object_meta_start_offset: usize,


klion26 · 2025-07-21T03:26:29Z

parquet-variant/src/builder.rs

+        // the length of the metadata's field names is a very cheap to compute the upper bound.
+        // it will almost always be a tight upper bound as well -- it would take a pretty
+        // carefully  crafted object to use only the early field ids of a large dictionary.
+        let max_id = metadata_builder.field_names.len();


IIUC, metadata_builder.field_names.len() returns the size of the underlying map, it may be bigger than the actual max field_id, changing to this here is that it can avoid one pass for the field_ids to calucalute the max_id.

But this change can make test more diffict(more difficult to calculate the header size) -- such as the failed case from_json::test::test_json_to_variant_object_very_large (test code here/failed ci)

@alamb @viirya @scovich Maybe we can revert this to the original code(travel one pass from filed_ids) what do you think about this? thanks. -- added a separate commit which reverted this logic to see if any other tests failed.

…revert if needed

scovich · 2025-07-21T13:05:49Z

@scovich Thank you very much for the detailed review and suggestion, I've addressed most of them (add comments for the rest), please take another look when you're free.

I'm excited to take a look ASAP, but I'll be on the road a lot this week. Please don't take my absenteeism as a lack of interest!

scovich

This looks correct and is already faster than the baseline.

I'd love to see a follow-up that tries populating that temp vec with the actual header bytes instead of zeros -- hopefully it would give another nice speed boost on top of this one.

scovich · 2025-07-21T19:40:42Z

parquet-variant/src/builder.rs

        });
-
        let max_id = self.fields.iter().map(|(i, _)| *i).max().unwrap_or(0);


Why would it waste space? With high probability, the largest field id in the object will take the same number of bytes to encode as self.fields.len() - 1 (the highest field id in the dictionary), so using the latter shouldn't change anything?

scovich · 2025-07-21T19:43:00Z

parquet-variant/src/builder.rs

+        let header_size = 1 + // header byte
+            (if is_large { 4 } else { 1 }) + // num_fields
+            (num_fields * id_size as usize) + // field IDs
+            ((num_fields + 1) * offset_size as usize); // field offsets + data_size


Or try the other approach that populates the tmp vec with non-zero bytes before splicing it into the main buffer? That's a lot simpler than this iterator-based suggestion, and likely performs better too.

scovich · 2025-07-21T19:44:32Z

parquet-variant/src/builder.rs

+            self.parent_state
+                .buffer()
+                .inner_mut()
+                .truncate(self.parent_offset_base);


now that we have two offset bases, should we rename this one as parent_value_offset_base for clarity?

scovich · 2025-07-21T19:45:35Z

parquet-variant/src/builder.rs

-        assert_eq!(metadata.len(), 1);
-        assert_eq!(&metadata[0], "name"); // not rolled back
+        assert!(metadata.is_empty()); // rolled back


alamb · 2025-07-21T23:00:53Z

🤖 ./gh_compare_arrow.sh Benchmark Script Running
Linux aal-dev 6.11.0-1016-gcp #16~24.04.1-Ubuntu SMP Wed May 28 02:40:52 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing 7899-avoid-extra-allocation-in-object-builder (f0d35de) to 03a837e diff
BENCH_NAME=variant_kernels
BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental --bench variant_kernels
BENCH_FILTER=
BENCH_BRANCH_NAME=7899-avoid-extra-allocation-in-object-builder
Results will be posted here when complete

klion26

@scovich thanks for the review, parent_offset_base has been renamed to parent_value_offset_base, and fill the tmp vec with header before splice improvements will be done with a follow-up pr.

klion26 · 2025-07-22T05:42:10Z

parquet-variant/src/builder.rs

        });
-
        let max_id = self.fields.iter().map(|(i, _)| *i).max().unwrap_or(0);


~~There will be no waste here, this is where I made a mistake before~~

klion26 · 2025-07-22T05:45:41Z

parquet-variant/src/builder.rs

+        let header_size = 1 + // header byte
+            (if is_large { 4 } else { 1 }) + // num_fields
+            (num_fields * id_size as usize) + // field IDs
+            ((num_fields + 1) * offset_size as usize); // field offsets + data_size


ok, will do it in a followup

klion26 · 2025-07-22T05:46:03Z

parquet-variant/src/builder.rs

+            self.parent_state
+                .buffer()
+                .inner_mut()
+                .truncate(self.parent_offset_base);


…cation-in-object-builder

alamb

Thanks again everyone -- this change is looking really nice now.

I went through this PR again and it looks to me like all the comments have been addressed, but I am not 100% sure

@scovich has approved the PR so I assume he is happy with this approach.
@viirya are you happy with this PR now or shall we wait for another round of reivew?

alamb · 2025-07-22T18:34:37Z

🤖 ./gh_compare_arrow.sh Benchmark Script Running
Linux aal-dev 6.11.0-1016-gcp #16~24.04.1-Ubuntu SMP Wed May 28 02:40:52 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing 7899-avoid-extra-allocation-in-object-builder (19bb544) to 291e6e5 diff
BENCH_NAME=variant_kernels
BENCH_COMMAND=cargo bench --features=arrow,async,test_common,experimental --bench variant_kernels
BENCH_FILTER=
BENCH_BRANCH_NAME=7899-avoid-extra-allocation-in-object-builder
Results will be posted here when complete

alamb · 2025-07-22T18:39:14Z

🤖: Benchmark completed

Details

group                                                                7899-avoid-extra-allocation-in-object-builder    main
-----                                                                ---------------------------------------------    ----
batch_json_string_to_variant json_list 8k string                     1.00     28.4±0.11ms        ? ?/sec              1.00     28.6±0.15ms        ? ?/sec
batch_json_string_to_variant random_json(2633 bytes per document)    1.00    368.2±4.36ms        ? ?/sec              1.06    389.3±6.43ms        ? ?/sec
batch_json_string_to_variant repeated_struct 8k string               1.00      8.4±0.02ms        ? ?/sec              1.02      8.5±0.02ms        ? ?/sec
variant_get_primitive                                                1.00   1381.5±3.14µs        ? ?/sec              1.00   1380.5±3.31µs        ? ?/sec

alamb · 2025-07-22T21:58:53Z

Ok, looks good to me. Thanks again. Let's address any additional comments as a follow on PR

klion26 · 2025-07-23T05:04:56Z

@alamb @scovich @viirya Thank you for the review; I learned a lot from this PR. I've created issues #7977 and #7978 to track the follow-up

github-actions bot added the parquet Changes to the parquet crate label Jul 16, 2025

klion26 commented Jul 16, 2025

View reviewed changes

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from 91cfb73 to 1f2bcc3 Compare July 16, 2025 10:29

alamb reviewed Jul 16, 2025

View reviewed changes

klion26 added 2 commits July 17, 2025 09:59

[Variant] Avoid extra allocation in object builder

6dbd9e7

This commit will reuse the parent buffer for object builder. It can avoid the extra allocation for the object and the later buffer copy.

fixup! [Variant] Avoid extra allocation in object builder

55def1d

trigger ci

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from 1f2bcc3 to 6096566 Compare July 17, 2025 03:48

klion26 commented Jul 17, 2025

View reviewed changes

fixup! [Variant] Avoid extra allocation in object builder

442c935

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from 6096566 to 442c935 Compare July 17, 2025 05:55

alamb mentioned this pull request Jul 17, 2025

Introduce buffer_and_metadata_builder to avoid duplication klion26/arrow-rs#1

Open

alamb reviewed Jul 17, 2025

View reviewed changes

klion26 mentioned this pull request Jul 18, 2025

[Variant] Avoiding extra splice in ObjectBuilder::finish if possible #7960

Open

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from e24d8e3 to deb0782 Compare July 18, 2025 05:11

klion26 commented Jul 18, 2025

View reviewed changes

fixup! [Variant] Avoid extra allocation in object builder

f5b0465

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from deb0782 to f5b0465 Compare July 18, 2025 05:20

alamb mentioned this pull request Jul 18, 2025

Convert JSON to VariantArray without copying (8 - 32% faster) #7911

Merged

alamb approved these changes Jul 18, 2025

View reviewed changes

scovich reviewed Jul 18, 2025

View reviewed changes

klion26 added 2 commits July 21, 2025 00:07

fixup! Add tests for BatchCoalescer::push_batch_with_filter, fix bug (

2b1edde

apache#7774)

truncate metadata when drop object builder

690dc35

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from 0fecaa0 to c4d26db Compare July 20, 2025 16:22

klion26 commented Jul 20, 2025

View reviewed changes

use metadata lenght as the max_id for field id

bdf1f2d

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from c4d26db to bdf1f2d Compare July 20, 2025 16:32

viirya reviewed Jul 20, 2025

View reviewed changes

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from d11441d to 5991cc7 Compare July 21, 2025 06:39

klion26 commented Jul 21, 2025

View reviewed changes

klion26 added 2 commits July 21, 2025 19:28

address comments from viirya

c76e612

calculate max_id with for-loop to see if we can pass all tests, will …

f0d35de

…revert if needed

klion26 force-pushed the 7899-avoid-extra-allocation-in-object-builder branch from b6af58d to f0d35de Compare July 21, 2025 11:28

alamb mentioned this pull request Jul 21, 2025

[Variant] WIP Tests for variant_get of shredded variants #7965

Closed

scovich approved these changes Jul 21, 2025

View reviewed changes

rename parent_offset_base to parent_value_offset_base

d9a61d4

klion26 commented Jul 22, 2025

View reviewed changes

alamb added 2 commits July 22, 2025 06:46

Merge remote-tracking branch 'apache/main' into 7899-avoid-extra-allo…

123d6f3

…cation-in-object-builder

revert unecessary refactor

19bb544

alamb approved these changes Jul 22, 2025

View reviewed changes

alamb merged commit 6874ffa into apache:main Jul 22, 2025
12 checks passed

carpecodeum mentioned this pull request Jul 23, 2025

[VARIANT] Path-based Field Extraction for VariantArray #7946

Draft

klion26 mentioned this pull request Jul 23, 2025

[Variant] Optimize the object header generation logic in ObjectBuilder::finish #7978

Open

	fn test_object_builder_to_list_builder_outer_no_finish() {
	let mut builder = VariantBuilder::new();
	let mut object_builder = builder.new_object();
	object_builder.insert("first", 1i8);

	// Create a nested list builder and finish it
	let mut nested_list_builder = object_builder.new_list("nested");
	nested_list_builder.append_value("hi");
	nested_list_builder.finish();

	// Drop the outer object builder without finishing it
	drop(object_builder);

	builder.append_value(2i8);

	// Only the second attempt should appear in the final variant
	let (metadata, value) = builder.finish();
	let metadata = VariantMetadata::try_new(&metadata).unwrap();
	assert_eq!(metadata.len(), 2);
	assert_eq!(&metadata[0], "first");
	assert_eq!(&metadata[1], "nested"); // not rolled back

	let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();
	assert_eq!(variant, Variant::Int8(2));
	}

	#[test]
	fn test_object_builder_to_object_builder_inner_no_finish() {
	let mut builder = VariantBuilder::new();
	let mut object_builder = builder.new_object();
	object_builder.insert("first", 1i8);

	// Create a nested object builder but never finish it
	let mut nested_object_builder = object_builder.new_object("nested");
	nested_object_builder.insert("name", "unknown");
	drop(nested_object_builder);

	object_builder.insert("second", 2i8);

	// The parent object should only contain the original fields
	object_builder.finish().unwrap();
	let (metadata, value) = builder.finish();
	let metadata = VariantMetadata::try_new(&metadata).unwrap();
	assert_eq!(metadata.len(), 3);
	assert_eq!(&metadata[0], "first");
	assert_eq!(&metadata[1], "name"); // not rolled back
	assert_eq!(&metadata[2], "second");

	let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();
	let obj = variant.as_object().unwrap();
	assert_eq!(obj.len(), 2);
	assert_eq!(obj.get("first"), Some(Variant::Int8(1)));
	assert_eq!(obj.get("second"), Some(Variant::Int8(2)));
	}

	#[test]
	fn test_object_builder_to_object_builder_outer_no_finish() {
	let mut builder = VariantBuilder::new();
	let mut object_builder = builder.new_object();
	object_builder.insert("first", 1i8);

	// Create a nested object builder and finish it
	let mut nested_object_builder = object_builder.new_object("nested");
	nested_object_builder.insert("name", "unknown");
	nested_object_builder.finish().unwrap();

	// Drop the outer object builder without finishing it
	drop(object_builder);

	builder.append_value(2i8);

	// Only the second attempt should appear in the final variant
	let (metadata, value) = builder.finish();
	let metadata = VariantMetadata::try_new(&metadata).unwrap();
	assert_eq!(metadata.len(), 3);
	assert_eq!(&metadata[0], "first"); // not rolled back
	assert_eq!(&metadata[1], "name"); // not rolled back
	assert_eq!(&metadata[2], "nested"); // not rolled back

	let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();
	assert_eq!(variant, Variant::Int8(2));
	}

	let offset_shift = match &self.parent_state {
	let offset_shift = self.parent_state.object_start_offset();


		let starting_offset = self.object_start_offset;

		// Shift existing data to make room for the header

	buffer.splice(starting_offset..starting_offset, vec![0u8; header_size]);
	buffer.splice(starting_offset..starting_offset, std::iter::repeat_n(0u8, header_size));

	for (&field_id, _) in &self.fields {
	for field_id in self.fields.keys() {

		});

		let max_id = self.fields.iter().map(\|(i, _)\| *i).max().unwrap_or(0);

		// as object builder has been reused the parent buffer,
		// we need to shift the offset by the starting offset of the parent object

	/// the starting offset in the parent's buffer where this object starts
	/// The starting offset in the parent's buffer where this object starts

		if self.validate_unique_fields && !self.duplicate_fields.is_empty() {
		let metadata_builder = self.parent_state.metadata_builder();

	object_meta_start_offset: usize,
	parent_metadata_offset_base: usize,

		// current object starts from `object_start_offset`
		let data_size = current_offset - self.object_start_offset;

[Variant] Avoid extra allocation in object builder #7935

[Variant] Avoid extra allocation in object builder #7935

Uh oh!

Conversation

klion26 commented Jul 16, 2025 • edited by viirya Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

klion26 commented Jul 16, 2025

Uh oh!

klion26 Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

klion26 commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

klion26 commented Jul 17, 2025

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

klion26 Jul 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

klion26 commented Jul 18, 2025

Uh oh!

klion26 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

klion26 commented Jul 16, 2025 •

edited by viirya

Loading

klion26 Jul 16, 2025 •

edited

Loading

klion26 commented Jul 16, 2025 •

edited

Loading

klion26 Jul 18, 2025 •

edited

Loading

klion26 left a comment •

edited

Loading

scovich Jul 19, 2025 •

edited

Loading