[Variant] Reduce variant-related struct sizes (#7888)

# Which issue does this PR close?

- Closes #7831

# Rationale for this change

Variants naturally work with `u32` offsets, field ids, etc. Widening
them artificially to `usize` on 64-bit architectures causes several
problems:
1. A majority of developers will be using 64-bit architectures, and are
unlikely to think about integer overflow issues when working with
`usize`. But it's actually quite easy for malicious data or buggy code
to overflow the `u32` values that variant actually relies on. Worse, it
becomes difficult, if not impossible, to validate the code's resistance
to 32-bit integer overflow, when manipulating `usize` values on 64-bit
hardware.
2. Related to 1/, casting from `usize` to `u32` can clip the value on
64-bit hardware, which makes it harder to reason about the code's
correctness (always wondering whether the value _might_ be larger than
32-bits can hold). In contrast, casting from `u32` to `usize` is safe in
spite of being fallible in rust (assumes we do _not_ need to support
16-bit architectures).
3. The variant-related data structures occupy significantly more space
than they need to, when storing (64-bit) `usize` offsets instead of
`u32`.

# What changes are included in this PR?

Store all variant-related offsets as `u32` instead of `usize`. The
`VariantMetadata`, `VariantObject` and `VariantList` structs shrink to
32/64/64 bytes (previously 40/88/80 bytes).

Also, rename `OffsetSizeBytes::unpack_usize[_at_offset]` methods to
`unpack_u32[_at_offset]`, to more accurately reflect what they actually
do now.

# Are these changes tested?

Existing unit tests cover the use of these values; new static assertions
will catch any future size changes.

# Are there any user-facing changes?

`VariantMetadata` is no longer `Copy`, reflecting the fact that this PR
still leaves it 2x larger than a fat pointer.

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: scovich

parquet-variant-json/src/from_json.rs

@@ -165,7 +165,7 @@ mod test {
         expected: Variant<'a, 'a>,
     }
 
-    impl<'a> JsonToVariantTest<'a> {
+    impl JsonToVariantTest<'_> {
         fn run(self) -> Result<(), ArrowError> {
             let mut variant_builder = VariantBuilder::new();
             json_to_variant(self.json, &mut variant_builder)?;

parquet-variant/src/builder.rs

@@ -1932,7 +1932,6 @@ mod tests {
         assert!(metadata.is_empty());
 
         let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();
-        assert!(metadata.is_empty());
         assert_eq!(variant, Variant::Int8(42));
     }
 

parquet-variant/src/decoder.rs

@@ -22,8 +22,6 @@ use crate::ShortString;
 use arrow_schema::ArrowError;
 use chrono::{DateTime, Duration, NaiveDate, NaiveDateTime, Utc};
 
-use std::num::TryFromIntError;
-
 /// The basic type of a [`Variant`] value, encoded in the first two bits of the
 /// header byte.
 ///
@@ -147,11 +145,9 @@ impl OffsetSizeBytes {
     /// * `bytes` – the byte buffer to index
     /// * `index` – 0-based index into the buffer
     ///
-    /// Each value is `self as usize` bytes wide (1, 2, 3 or 4).
-    /// Three-byte values are zero-extended to 32 bits before the final
-    /// fallible cast to `usize`.
-    pub(crate) fn unpack_usize(&self, bytes: &[u8], index: usize) -> Result<usize, ArrowError> {
-        self.unpack_usize_at_offset(bytes, 0, index)
+    /// Each value is `self as u32` bytes wide (1, 2, 3 or 4), zero-extended to 32 bits as needed.
+    pub(crate) fn unpack_u32(&self, bytes: &[u8], index: usize) -> Result<u32, ArrowError> {
+        self.unpack_u32_at_offset(bytes, 0, index)
     }
 
     /// Return one unsigned little-endian value from `bytes`.
@@ -162,15 +158,13 @@ impl OffsetSizeBytes {
     /// * `offset_index` – 0-based index **after** the skipped bytes
     ///   (`0` is the first value, `1` the next, …).
     ///
-    /// Each value is `self as usize` bytes wide (1, 2, 3 or 4).
-    /// Three-byte values are zero-extended to 32 bits before the final
-    /// fallible cast to `usize`.
-    pub(crate) fn unpack_usize_at_offset(
+    /// Each value is `self as u32` bytes wide (1, 2, 3 or 4), zero-extended to 32 bits as needed.
+    pub(crate) fn unpack_u32_at_offset(
         &self,
         bytes: &[u8],
         byte_offset: usize,  // how many bytes to skip
         offset_index: usize, // which offset in an array of offsets
-    ) -> Result<usize, ArrowError> {
+    ) -> Result<u32, ArrowError> {
         use OffsetSizeBytes::*;
 
         // Index into the byte array:
@@ -179,7 +173,7 @@ impl OffsetSizeBytes {
             .checked_mul(*self as usize)
             .and_then(|n| n.checked_add(byte_offset))
             .ok_or_else(|| overflow_error("unpacking offset array value"))?;
-        let result = match self {
+        let value = match self {
             One => u8::from_le_bytes(array_from_slice(bytes, offset)?).into(),
             Two => u16::from_le_bytes(array_from_slice(bytes, offset)?).into(),
             Three => {
@@ -192,11 +186,7 @@ impl OffsetSizeBytes {
             }
             Four => u32::from_le_bytes(array_from_slice(bytes, offset)?),
         };
-
-        // Convert the u32 we extracted to usize (should always succeed on 32- and 64-bit arch)
-        result
-            .try_into()
-            .map_err(|e: TryFromIntError| ArrowError::InvalidArgumentError(e.to_string()))
+        Ok(value)
     }
 }
 
@@ -518,57 +508,51 @@ mod tests {
     }
 
     #[test]
-    fn unpack_usize_all_widths() {
+    fn unpack_u32_all_widths() {
         // One-byte offsets
         let buf_one = [0x01u8, 0xAB, 0xCD];
-        assert_eq!(
-            OffsetSizeBytes::One.unpack_usize(&buf_one, 0).unwrap(),
-            0x01
-        );
-        assert_eq!(
-            OffsetSizeBytes::One.unpack_usize(&buf_one, 2).unwrap(),
-            0xCD
-        );
+        assert_eq!(OffsetSizeBytes::One.unpack_u32(&buf_one, 0).unwrap(), 0x01);
+        assert_eq!(OffsetSizeBytes::One.unpack_u32(&buf_one, 2).unwrap(), 0xCD);
 
         // Two-byte offsets (little-endian 0x1234, 0x5678)
         let buf_two = [0x34, 0x12, 0x78, 0x56];
         assert_eq!(
-            OffsetSizeBytes::Two.unpack_usize(&buf_two, 0).unwrap(),
+            OffsetSizeBytes::Two.unpack_u32(&buf_two, 0).unwrap(),
             0x1234
         );
         assert_eq!(
-            OffsetSizeBytes::Two.unpack_usize(&buf_two, 1).unwrap(),
+            OffsetSizeBytes::Two.unpack_u32(&buf_two, 1).unwrap(),
             0x5678
         );
 
         // Three-byte offsets (0x030201 and 0x0000FF)
         let buf_three = [0x01, 0x02, 0x03, 0xFF, 0x00, 0x00];
         assert_eq!(
-            OffsetSizeBytes::Three.unpack_usize(&buf_three, 0).unwrap(),
+            OffsetSizeBytes::Three.unpack_u32(&buf_three, 0).unwrap(),
             0x030201
         );
         assert_eq!(
-            OffsetSizeBytes::Three.unpack_usize(&buf_three, 1).unwrap(),
+            OffsetSizeBytes::Three.unpack_u32(&buf_three, 1).unwrap(),
             0x0000FF
         );
 
         // Four-byte offsets (0x12345678, 0x90ABCDEF)
         let buf_four = [0x78, 0x56, 0x34, 0x12, 0xEF, 0xCD, 0xAB, 0x90];
         assert_eq!(
-            OffsetSizeBytes::Four.unpack_usize(&buf_four, 0).unwrap(),
+            OffsetSizeBytes::Four.unpack_u32(&buf_four, 0).unwrap(),
             0x1234_5678
         );
         assert_eq!(
-            OffsetSizeBytes::Four.unpack_usize(&buf_four, 1).unwrap(),
+            OffsetSizeBytes::Four.unpack_u32(&buf_four, 1).unwrap(),
             0x90AB_CDEF
         );
     }
 
     #[test]
-    fn unpack_usize_out_of_bounds() {
+    fn unpack_u32_out_of_bounds() {
         let tiny = [0x00u8]; // deliberately too short
-        assert!(OffsetSizeBytes::Two.unpack_usize(&tiny, 0).is_err());
-        assert!(OffsetSizeBytes::Three.unpack_usize(&tiny, 0).is_err());
+        assert!(OffsetSizeBytes::Two.unpack_u32(&tiny, 0).is_err());
+        assert!(OffsetSizeBytes::Three.unpack_u32(&tiny, 0).is_err());
     }
 
     #[test]
@@ -584,20 +568,20 @@ mod tests {
         let width = OffsetSizeBytes::Two;
 
         // dictionary_size starts immediately after the header byte
-        let dict_size = width.unpack_usize_at_offset(&buf, 1, 0).unwrap();
+        let dict_size = width.unpack_u32_at_offset(&buf, 1, 0).unwrap();
         assert_eq!(dict_size, 2);
 
         // offset array immediately follows the dictionary size
-        let first = width.unpack_usize_at_offset(&buf, 1, 1).unwrap();
+        let first = width.unpack_u32_at_offset(&buf, 1, 1).unwrap();
         assert_eq!(first, 0);
 
-        let second = width.unpack_usize_at_offset(&buf, 1, 2).unwrap();
+        let second = width.unpack_u32_at_offset(&buf, 1, 2).unwrap();
         assert_eq!(second, 5);
 
-        let third = width.unpack_usize_at_offset(&buf, 1, 3).unwrap();
+        let third = width.unpack_u32_at_offset(&buf, 1, 3).unwrap();
         assert_eq!(third, 9);
 
-        let err = width.unpack_usize_at_offset(&buf, 1, 4);
+        let err = width.unpack_u32_at_offset(&buf, 1, 4);
         assert!(err.is_err())
     }
 }

parquet-variant/src/utils.rs

@@ -122,3 +122,12 @@ where
 
     Some(Err(start))
 }
+
+/// Verifies the expected size of type T, for a type that should only grow if absolutely necessary.
+#[allow(unused)]
+pub(crate) const fn expect_size_of<T>(expected: usize) {
+    let size = std::mem::size_of::<T>();
+    if size != expected {
+        let _ = [""; 0][size];
+    }
+}

parquet-variant/src/variant.rs

@@ -256,6 +256,9 @@ pub enum Variant<'m, 'v> {
     List(VariantList<'m, 'v>),
 }
 
+// We don't want this to grow because it could hurt performance of a frequently-created type.
+const _: () = crate::utils::expect_size_of::<Variant>(80);
+
 impl<'m, 'v> Variant<'m, 'v> {
     /// Attempts to interpret a metadata and value buffer pair as a new `Variant`.
     ///

parquet-variant/src/variant/list.rs

@@ -23,7 +23,7 @@ use crate::variant::{Variant, VariantMetadata};
 use arrow_schema::ArrowError;
 
 // The value header occupies one byte; use a named constant for readability
-const NUM_HEADER_BYTES: usize = 1;
+const NUM_HEADER_BYTES: u32 = 1;
 
 /// A parsed version of the variant array value header byte.
 #[derive(Debug, Clone, PartialEq)]
@@ -34,15 +34,15 @@ pub(crate) struct VariantListHeader {
 
 impl VariantListHeader {
     // Hide the ugly casting
-    const fn num_elements_size(&self) -> usize {
+    const fn num_elements_size(&self) -> u32 {
         self.num_elements_size as _
     }
-    const fn offset_size(&self) -> usize {
+    const fn offset_size(&self) -> u32 {
         self.offset_size as _
     }
 
     // Avoid materializing this offset, since it's cheaply and safely computable
-    const fn first_offset_byte(&self) -> usize {
+    const fn first_offset_byte(&self) -> u32 {
         NUM_HEADER_BYTES + self.num_elements_size()
     }
 
@@ -122,11 +122,14 @@ pub struct VariantList<'m, 'v> {
     pub metadata: VariantMetadata<'m>,
     pub value: &'v [u8],
     header: VariantListHeader,
-    num_elements: usize,
-    first_value_byte: usize,
+    num_elements: u32,
+    first_value_byte: u32,
     validated: bool,
 }
 
+// We don't want this to grow because it could increase the size of `Variant` and hurt performance.
+const _: () = crate::utils::expect_size_of::<VariantList>(64);
+
 impl<'m, 'v> VariantList<'m, 'v> {
     /// Attempts to interpret `value` as a variant array value.
     ///
@@ -157,7 +160,7 @@ impl<'m, 'v> VariantList<'m, 'v> {
         let num_elements =
             header
                 .num_elements_size
-                .unpack_usize_at_offset(value, NUM_HEADER_BYTES, 0)?;
+                .unpack_u32_at_offset(value, NUM_HEADER_BYTES as _, 0)?;
 
         // (num_elements + 1) * offset_size + first_offset_byte
         let first_value_byte = num_elements
@@ -185,10 +188,10 @@ impl<'m, 'v> VariantList<'m, 'v> {
 
         // Use the last offset to upper-bound the value buffer
         let last_offset = new_self
-            .get_offset(num_elements)?
+            .get_offset(num_elements as _)?
             .checked_add(first_value_byte)
             .ok_or_else(|| overflow_error("variant array size"))?;
-        new_self.value = slice_from_slice(value, ..last_offset)?;
+        new_self.value = slice_from_slice(value, ..last_offset as _)?;
         Ok(new_self)
     }
 
@@ -210,7 +213,7 @@ impl<'m, 'v> VariantList<'m, 'v> {
 
             let offset_buffer = slice_from_slice(
                 self.value,
-                self.header.first_offset_byte()..self.first_value_byte,
+                self.header.first_offset_byte() as _..self.first_value_byte as _,
             )?;
 
             let offsets =
@@ -226,15 +229,15 @@ impl<'m, 'v> VariantList<'m, 'v> {
                 ));
             }
 
-            let value_buffer = slice_from_slice(self.value, self.first_value_byte..)?;
+            let value_buffer = slice_from_slice(self.value, self.first_value_byte as _..)?;
 
             // Validate whether values are valid variant objects
             for i in 1..offsets.len() {
                 let start_offset = offsets[i - 1];
                 let end_offset = offsets[i];
 
                 let value_bytes = slice_from_slice(value_buffer, start_offset..end_offset)?;
-                Variant::try_new_with_metadata(self.metadata, value_bytes)?;
+                Variant::try_new_with_metadata(self.metadata.clone(), value_bytes)?;
             }
 
             self.validated = true;
@@ -244,7 +247,7 @@ impl<'m, 'v> VariantList<'m, 'v> {
 
     /// Return the length of this array
     pub fn len(&self) -> usize {
-        self.num_elements
+        self.num_elements as _
     }
 
     /// Is the array of zero length
@@ -256,7 +259,7 @@ impl<'m, 'v> VariantList<'m, 'v> {
     ///
     /// [invalid]: Self#Validation
     pub fn get(&self, index: usize) -> Option<Variant<'m, 'v>> {
-        (index < self.num_elements).then(|| {
+        (index < self.len()).then(|| {
             self.try_get_with_shallow_validation(index)
                 .expect("Invalid variant array element")
         })
@@ -272,10 +275,10 @@ impl<'m, 'v> VariantList<'m, 'v> {
     fn try_get_with_shallow_validation(&self, index: usize) -> Result<Variant<'m, 'v>, ArrowError> {
         // Fetch the value bytes between the two offsets for this index, from the value array region
         // of the byte buffer
-        let byte_range = self.get_offset(index)?..self.get_offset(index + 1)?;
+        let byte_range = self.get_offset(index)? as _..self.get_offset(index + 1)? as _;
         let value_bytes =
-            slice_from_slice_at_offset(self.value, self.first_value_byte, byte_range)?;
-        Variant::try_new_with_metadata_and_shallow_validation(self.metadata, value_bytes)
+            slice_from_slice_at_offset(self.value, self.first_value_byte as _, byte_range)?;
+        Variant::try_new_with_metadata_and_shallow_validation(self.metadata.clone(), value_bytes)
     }
 
     /// Iterates over the values of this list. When working with [unvalidated] input, consider
@@ -297,14 +300,14 @@ impl<'m, 'v> VariantList<'m, 'v> {
     fn iter_try_with_shallow_validation(
         &self,
     ) -> impl Iterator<Item = Result<Variant<'m, 'v>, ArrowError>> + '_ {
-        (0..self.len()).map(move |i| self.try_get_with_shallow_validation(i))
+        (0..self.len()).map(|i| self.try_get_with_shallow_validation(i))
     }
 
     // Attempts to retrieve the ith offset from the offset array region of the byte buffer.
-    fn get_offset(&self, index: usize) -> Result<usize, ArrowError> {
-        let byte_range = self.header.first_offset_byte()..self.first_value_byte;
+    fn get_offset(&self, index: usize) -> Result<u32, ArrowError> {
+        let byte_range = self.header.first_offset_byte() as _..self.first_value_byte as _;
         let offset_bytes = slice_from_slice(self.value, byte_range)?;
-        self.header.offset_size.unpack_usize(offset_bytes, index)
+        self.header.offset_size.unpack_u32(offset_bytes, index)
     }
 }
 
@@ -623,7 +626,7 @@ mod tests {
             expected_num_element_size,
             variant_list.header.num_elements_size
         );
-        assert_eq!(list_size, variant_list.num_elements);
+        assert_eq!(list_size, variant_list.num_elements as usize);
 
         // verify the data in the variant
         assert_eq!(list_size, variant_list.len());