address comments from viirya

klion26 · klion26 · commit c76e612205b5 · 2025-07-21T19:28:20.000+08:00
diff --git a/parquet-variant/src/builder.rs b/parquet-variant/src/builder.rs
@@ -63,6 +63,12 @@ fn write_offset(buf: &mut Vec<u8>, value: usize, nbytes: u8) {
     buf.extend_from_slice(&bytes[..nbytes as usize]);
 }
 
+/// Write little-endian integer to buffer at a specific position
+fn write_offset_at_pos(buf: &mut [u8], start_pos: usize, value: usize, nbytes: u8) {
+    let bytes = value.to_le_bytes();
+    buf[start_pos..start_pos + nbytes as usize].copy_from_slice(&bytes[..nbytes as usize]);
+}
+
 /// Wrapper around a `Vec<u8>` that provides methods for appending
 /// primitive values, variant types, and metadata.
 ///
@@ -342,6 +348,63 @@ impl ValueBuffer {
             write_offset(buf, data_size, nbytes);
         }
     }
+
+    /// Writes out the header byte for a variant object or list, from the starting position
+    /// of the buffer, will return the position after this write
+    fn append_header_start_from_buf_pos(
+        &mut self,
+        start_pos: usize, // the start position where the header will be inserted
+        header_byte: u8,
+        is_large: bool,
+        num_fields: usize,
+    ) -> usize {
+        let buffer = self.inner_mut();
+
+        // Write header at the original start position
+        let mut header_pos = start_pos;
+
+        // Write header byte
+        buffer[header_pos] = header_byte;
+        header_pos += 1;
+
+        // 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;
+        }
+
+        header_pos
+    }
+
+    /// Writes out the offsets for an array of offsets, including the final offset (data size).
+    /// from the starting position of the buffer, will return the position after this write
+    fn append_offset_array_start_from_buf_pos(
+        &mut self,
+        start_pos: usize,
+        offsets: impl IntoIterator<Item = usize>,
+        data_size: Option<usize>,
+        nbytes: u8,
+    ) -> usize {
+        let buf = self.inner_mut();
+
+        let mut current_pos = start_pos;
+        for relative_offset in offsets {
+            write_offset_at_pos(buf, current_pos, relative_offset, nbytes);
+            current_pos += nbytes as usize;
+        }
+
+        // Write data_size
+        if let Some(data_size) = data_size {
+            // Write data_size at the end of the offsets
+            write_offset_at_pos(buf, current_pos, data_size, nbytes);
+            current_pos += nbytes as usize;
+        }
+
+        current_pos
+    }
 }
 
 /// Builder for constructing metadata for [`Variant`] values.
@@ -506,7 +569,7 @@ enum ParentState<'a> {
         metadata_builder: &'a mut MetadataBuilder,
         fields: &'a mut IndexMap<u32, usize>,
         field_name: &'a str,
-        object_start_offset: usize,
+        parent_offset_base: usize,
     },
 }
 
@@ -545,7 +608,7 @@ impl ParentState<'_> {
                 metadata_builder,
                 fields,
                 field_name,
-                object_start_offset,
+                parent_offset_base: object_start_offset,
                 ..
             } => {
                 let field_id = metadata_builder.upsert_field_name(field_name);
@@ -576,7 +639,7 @@ impl ParentState<'_> {
         }
     }
 
-    // return the offset of the underlying buffer at the time of calling this method.
+    // Return the offset of the underlying buffer at the time of calling this method.
     fn buffer_current_offset(&self) -> usize {
         match self {
             ParentState::Variant { buffer, .. }
@@ -585,7 +648,7 @@ impl ParentState<'_> {
         }
     }
 
-    // return the current index of the undelying metadata buffer at the time of calling this method.
+    // Return the current index of the undelying metadata buffer at the time of calling this method.
     fn metadata_current_offset(&self) -> usize {
         match self {
             ParentState::Variant {
@@ -1048,8 +1111,6 @@ impl<'a> ListBuilder<'a> {
 
         // Get parent's buffer
         let parent_buffer = self.parent_state.buffer();
-        // as object builder has been reused the parent buffer,
-        // we need to shift the offset by the starting offset of the parent object
         let starting_offset = parent_buffer.offset();
 
         // Write header
@@ -1078,12 +1139,12 @@ impl Drop for ListBuilder<'_> {
 pub struct ObjectBuilder<'a> {
     parent_state: ParentState<'a>,
     fields: IndexMap<u32, usize>, // (field_id, offset)
-    /// the starting offset in the parent's buffer where this object starts
-    object_start_offset: usize,
-    /// the starting offset in the parent's metadata buffer where this object starts
+    /// The starting offset in the parent's buffer where this object starts
+    parent_offset_base: 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,
-    /// whether the object has been finished, the written content of the current object
+    parent_metadata_offset_base: usize,
+    /// Whether the object has been finished, the written content of the current object
     /// will be truncated in `drop` if `has_been_finished` is false
     has_been_finished: bool,
     validate_unique_fields: bool,
@@ -1093,14 +1154,14 @@ pub struct ObjectBuilder<'a> {
 
 impl<'a> ObjectBuilder<'a> {
     fn new(parent_state: ParentState<'a>, validate_unique_fields: bool) -> Self {
-        let start_offset = parent_state.buffer_current_offset();
-        let meta_start_offset = parent_state.metadata_current_offset();
+        let offset_base = parent_state.buffer_current_offset();
+        let meta_offset_base = parent_state.metadata_current_offset();
         Self {
             parent_state,
             fields: IndexMap::new(),
-            object_start_offset: start_offset,
+            parent_offset_base: offset_base,
             has_been_finished: false,
-            object_meta_start_offset: meta_start_offset,
+            parent_metadata_offset_base: meta_offset_base,
             validate_unique_fields,
             duplicate_fields: HashSet::new(),
         }
@@ -1128,7 +1189,7 @@ impl<'a> ObjectBuilder<'a> {
         let (buffer, metadata_builder) = self.parent_state.buffer_and_metadata_builder();
 
         let field_id = metadata_builder.upsert_field_name(key);
-        let field_start = buffer.offset() - self.object_start_offset;
+        let field_start = buffer.offset() - self.parent_offset_base;
 
         if self.fields.insert(field_id, field_start).is_some() && self.validate_unique_fields {
             self.duplicate_fields.insert(field_id);
@@ -1158,7 +1219,7 @@ impl<'a> ObjectBuilder<'a> {
             metadata_builder,
             fields: &mut self.fields,
             field_name: key,
-            object_start_offset: self.object_start_offset,
+            parent_offset_base: self.parent_offset_base,
         };
         (state, validate_unique_fields)
     }
@@ -1207,14 +1268,14 @@ impl<'a> ObjectBuilder<'a> {
 
         // 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.
+        // carefully crafted object to use only the early field ids of a large dictionary.
         let max_id = metadata_builder.field_names.len();
         let id_size = int_size(max_id);
 
         let parent_buffer = self.parent_state.buffer();
         let current_offset = parent_buffer.offset();
-        // current object starts from `object_start_offset`
-        let data_size = current_offset - self.object_start_offset;
+        // Current object starts from `object_start_offset`
+        let data_size = current_offset - self.parent_offset_base;
         let offset_size = int_size(data_size);
 
         let num_fields = self.fields.len();
@@ -1225,7 +1286,7 @@ impl<'a> ObjectBuilder<'a> {
             (num_fields * id_size as usize) + // field IDs
             ((num_fields + 1) * offset_size as usize); // field offsets + data_size
 
-        let starting_offset = self.object_start_offset;
+        let starting_offset = self.parent_offset_base;
 
         // Shift existing data to make room for the header
         let buffer = parent_buffer.inner_mut();
@@ -1239,42 +1300,33 @@ impl<'a> ObjectBuilder<'a> {
 
         // Write header byte
         let header = object_header(is_large, id_size, offset_size);
-        buffer[header_pos] = header;
-        header_pos += 1;
-
-        // 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;
-        }
-
-        // Write field IDs
-        for field_id in self.fields.keys() {
-            let id_bytes = field_id.to_le_bytes();
-            buffer[header_pos..header_pos + id_size as usize]
-                .copy_from_slice(&id_bytes[..id_size as usize]);
-            header_pos += id_size as usize;
-        }
-
-        // Write field offsets (adjusted for header)
-        for relative_offset in self.fields.values() {
-            let offset_bytes = relative_offset.to_le_bytes();
-            buffer[header_pos..header_pos + offset_size as usize]
-                .copy_from_slice(&offset_bytes[..offset_size as usize]);
-            header_pos += offset_size as usize;
-        }
-
-        // Write data_size
-        let data_size_bytes = data_size.to_le_bytes();
-        buffer[header_pos..header_pos + offset_size as usize]
-            .copy_from_slice(&data_size_bytes[..offset_size as usize]);
 
+        header_pos = self
+            .parent_state
+            .buffer()
+            .append_header_start_from_buf_pos(header_pos, header, is_large, num_fields);
+
+        header_pos = self
+            .parent_state
+            .buffer()
+            .append_offset_array_start_from_buf_pos(
+                header_pos,
+                self.fields.keys().copied().map(|id| id as usize),
+                None,
+                id_size,
+            );
+
+        self.parent_state
+            .buffer()
+            .append_offset_array_start_from_buf_pos(
+                header_pos,
+                self.fields.values().copied(),
+                Some(data_size),
+                offset_size,
+            );
         self.parent_state.finish(starting_offset);
 
-        // mark that this object has been finished
+        // Mark that this object has been finished
         self.has_been_finished = true;
 
         Ok(())
@@ -1287,17 +1339,17 @@ impl<'a> ObjectBuilder<'a> {
 /// is finalized.
 impl Drop for ObjectBuilder<'_> {
     fn drop(&mut self) {
-        // truncate the buffer if the `finish` method has not been called.
+        // Truncate the buffer if the `finish` method has not been called.
         if !self.has_been_finished {
             self.parent_state
                 .buffer()
                 .inner_mut()
-                .truncate(self.object_start_offset);
+                .truncate(self.parent_offset_base);
 
             self.parent_state
                 .metadata_builder()
                 .field_names
-                .truncate(self.object_meta_start_offset);
+                .truncate(self.parent_metadata_offset_base);
         }
     }
 }
@@ -2078,7 +2130,7 @@ mod tests {
         assert_eq!(Variant::from(false), second_inner_list_g.get(1).unwrap());
     }
 
-    // this test wants to cover the logic for reuse parent buffer for list builder
+    // This test wants to cover the logic for reuse parent buffer for list builder
     // the builder looks like
     // [ "apple", "false", [{"a": "b", "b": "c"}, {"c":"d", "d":"e"}], [[1, true], ["tree", false]], 1]
     #[test]
@@ -2148,12 +2200,12 @@ mod tests {
 
         assert_eq!(5, outer_list.len());
 
-        // primitive value
+        // Primitive value
         assert_eq!(Variant::from("apple"), outer_list.get(0).unwrap());
         assert_eq!(Variant::from(false), outer_list.get(1).unwrap());
         assert_eq!(Variant::from(1), outer_list.get(4).unwrap());
 
-        // the first inner list [{"a": "b", "b": "c"}, {"c":"d", "d":"e"}]
+        // The first inner list [{"a": "b", "b": "c"}, {"c":"d", "d":"e"}]
         let list1_variant = outer_list.get(2).unwrap();
         let list1 = list1_variant.as_list().unwrap();
         assert_eq!(2, list1.len());
@@ -2166,19 +2218,19 @@ mod tests {
         assert_eq!("b", list1_obj1.field_name(1).unwrap());
         assert_eq!(Variant::from("c"), list1_obj1.field(1).unwrap());
 
-        // the second inner list [[1, true], ["tree", false]]
+        // The second inner list [[1, true], ["tree", false]]
         let list2_variant = outer_list.get(3).unwrap();
         let list2 = list2_variant.as_list().unwrap();
         assert_eq!(2, list2.len());
 
-        // the list [1, true]
+        // The list [1, true]
         let list2_list1_variant = list2.get(0).unwrap();
         let list2_list1 = list2_list1_variant.as_list().unwrap();
         assert_eq!(2, list2_list1.len());
         assert_eq!(Variant::from(1), list2_list1.get(0).unwrap());
         assert_eq!(Variant::from(true), list2_list1.get(1).unwrap());
 
-        // the list ["true", false]
+        // The list ["true", false]
         let list2_list2_variant = list2.get(1).unwrap();
         let list2_list2 = list2_list2_variant.as_list().unwrap();
         assert_eq!(2, list2_list2.len());
@@ -2673,8 +2725,8 @@ mod tests {
         // 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(), 1); // rolled back
-        assert_eq!(&metadata[0], "name");
+        assert_eq!(metadata.len(), 1);
+        assert_eq!(&metadata[0], "name"); // not rolled back
 
         let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();
         assert_eq!(variant, Variant::Int8(2));
