[Variant] Optimizing the ObjectBuilder::finish using iterator

Author: klion26

parquet-variant/src/builder.rs

@@ -64,19 +64,55 @@ 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]);
-}
-
 /// Append `value_size` bytes of given `value` into `dest`.
 fn append_packed_u32(dest: &mut Vec<u8>, value: u32, value_size: usize) {
     let n = dest.len() + value_size;
     dest.extend(value.to_le_bytes());
     dest.truncate(n);
 }
 
+/// An iterator that yields the bytes of a packed u32 iterator.
+/// Will yield the first `packed_bytes` bytes of each item in the iterator.
+struct PackedU32Iterator<T: Iterator<Item = [u8; 4]>> {
+    packed_bytes: usize,
+    iterator: T,
+    current_item: [u8; 4],
+    current_byte: usize, // 0..3
+}
+
+impl<T: 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: Iterator<Item = [u8; 4]>> Iterator for PackedU32Iterator<T> {
+    type Item = u8;
+
+    fn next(&mut self) -> Option<u8> {
+        if self.current_byte >= self.packed_bytes {
+            self.current_item = self.iterator.next()?;
+            self.current_byte = 0;
+        }
+
+        let rval = self.current_item[self.current_byte];
+        self.current_byte += 1;
+        Some(rval)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let lower = (self.packed_bytes - self.current_byte)
+            + self.packed_bytes * self.iterator.size_hint().0;
+        (lower, None)
+    }
+}
+
 /// Wrapper around a `Vec<u8>` that provides methods for appending
 /// primitive values, variant types, and metadata.
 ///
@@ -368,63 +404,6 @@ impl ValueBuffer {
 
         Ok(())
     }
-
-    /// 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.
@@ -1462,49 +1441,36 @@ impl<'a> ObjectBuilder<'a> {
         let num_fields = self.fields.len();
         let is_large = num_fields > u8::MAX as usize;
 
-        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
+        let num_fileds_size = if is_large { 4 } else { 1 }; // is_large: 4 bytes, else 1 byte.
 
-        let starting_offset = self.parent_value_offset_base;
+        let num_fields_bytes = num_fields.to_le_bytes();
+        let num_elements_bytes = num_fields_bytes.iter().take(num_fileds_size).copied();
 
-        // Shift existing data to make room for the header
-        let buffer = parent_buffer.inner_mut();
-        buffer.splice(
-            starting_offset..starting_offset,
-            std::iter::repeat_n(0u8, header_size),
+        let fields = PackedU32Iterator::new(
+            id_size as usize,
+            self.fields.keys().map(|offset| offset.to_le_bytes()),
+        );
+        let offsets = PackedU32Iterator::new(
+            offset_size as usize,
+            self.fields
+                .values()
+                .map(|offset| (*offset as u32).to_le_bytes()),
         );
 
-        // Write header at the original start position
-        let mut header_pos = starting_offset;
-
-        // Write header byte
+        let data_size_bytes = (data_size as u32).to_le_bytes();
+        let data_size_bytes_iter = data_size_bytes.iter().take(offset_size as usize).copied();
         let header = object_header(is_large, id_size, offset_size);
+        let bytess_to_splice = std::iter::once(header)
+            .chain(num_elements_bytes)
+            .chain(fields)
+            .chain(offsets)
+            .chain(data_size_bytes_iter);
+
+        let starting_offset = self.parent_value_offset_base;
+        // Shift existing data to make room for the header
+        let buffer = parent_buffer.inner_mut();
+        buffer.splice(starting_offset..starting_offset, bytess_to_splice);
 
-        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