[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?

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.
Modify the logic of new, finish, parent_state, and drop function according to the change.

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?

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 deb0782), 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 {


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 this extra splice in object builder finish #7960

Open

fixup! [Variant] Avoid extra allocation in object builder

deb0782

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

	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

[Variant] Avoid extra allocation in object builder #7935

Are you sure you want to change the base?

[Variant] Avoid extra allocation in object builder #7935

Uh oh!

Conversation

klion26 commented Jul 16, 2025

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!

klion26 commented Jul 18, 2025

Uh oh!

klion26 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 Jul 16, 2025 •

edited

Loading

klion26 commented Jul 16, 2025 •

edited

Loading

klion26 Jul 18, 2025 •

edited

Loading