add variant object tests with different sizes

fix tests

fix tests

Author: odysa

parquet-variant/src/builder.rs

@@ -990,7 +990,7 @@ impl<'m, 'v> VariantBuilderExt<'m, 'v> for VariantBuilder {
 
 #[cfg(test)]
 mod tests {
-    use crate::VariantMetadata;
+    use crate::{utils::first_byte_from_slice, VariantMetadata};
 
     use super::*;
 
@@ -2171,4 +2171,116 @@ mod tests {
         let variant = Variant::try_new_with_metadata(metadata, &value).unwrap();
         assert_eq!(variant, Variant::Int8(2));
     }
+
+    fn test_variant_object_with_count(count: i32, expected_field_id_size_minus_one: u8) {
+        let mut builder = VariantBuilder::new();
+        let mut obj = builder.new_object();
+
+        for val in 0..count {
+            let key = format!("id_{}", val);
+            obj.insert(&key, val);
+        }
+
+        obj.finish().unwrap();
+        let (metadata, value) = builder.finish();
+        let variant = Variant::try_new(&metadata, &value).unwrap();
+
+        if let Variant::Object(obj) = variant {
+            assert_eq!(obj.len(), count as usize);
+            assert_eq!(obj.get(&format!("id_{}", 0)).unwrap(), Variant::Int32(0));
+            assert_eq!(
+                obj.get(&format!("id_{}", count - 1)).unwrap(),
+                Variant::Int32(count - 1)
+            );
+
+            let header_byte = first_byte_from_slice(&value).unwrap();
+            let value_header = header_byte >> 2;
+            let field_id_size_minus_one = (value_header >> 2) & 0x03;
+            assert_eq!(
+                field_id_size_minus_one, expected_field_id_size_minus_one,
+                "Expected {}-byte field IDs",
+                expected_field_id_size_minus_one
+            );
+        } else {
+            panic!("Expected object variant");
+        }
+    }
+
+    #[test]
+    fn test_variant_object_257_elements() {
+        test_variant_object_with_count(2_i32.pow(8) + 1, 1); // 2^8 + 1, expected 2-byte field IDs
+    }
+
+    #[test]
+    fn test_variant_object_65537_elements() {
+        test_variant_object_with_count(2_i32.pow(16) + 1, 2); // 2^16 + 1, expected 3-byte field IDs
+    }
+
+    #[test]
+    #[ignore]
+    fn test_variant_object_16777217_elements() {
+        test_variant_object_with_count(2_i32.pow(24) + 1, 3); // 2^24 + 1, expected 4-byte field IDs
+    }
+
+    #[test]
+    fn test_variant_object_small_sizes_255_elements() {
+        // should pass
+        test_variant_object_with_count(255, 0);
+    }
+
+    fn test_variant_object_with_large_data(
+        data_size_per_field: usize,
+        expected_field_offset_size_minus_one: u8,
+    ) {
+        let num_fields = 20;
+        let mut builder = VariantBuilder::new();
+        let mut obj = builder.new_object();
+
+        let str_val = "a".repeat(data_size_per_field);
+
+        for val in 0..num_fields {
+            let key = format!("id_{}", val);
+            obj.insert(&key, str_val.as_str());
+        }
+
+        obj.finish().unwrap();
+        let (metadata, value) = builder.finish();
+        let variant = Variant::try_new(&metadata, &value).unwrap();
+
+        if let Variant::Object(obj) = variant {
+            assert_eq!(obj.len(), num_fields);
+            let header_byte = first_byte_from_slice(&value).unwrap();
+            let value_header = header_byte >> 2;
+            let field_offset_size_minus_one = value_header & 0x03;
+            assert_eq!(
+                field_offset_size_minus_one,
+                expected_field_offset_size_minus_one,
+                "Expected {}-byte field offsets, got {}-byte field offsets",
+                expected_field_offset_size_minus_one + 1,
+                field_offset_size_minus_one + 1
+            );
+        } else {
+            panic!("Expected object variant");
+        }
+    }
+
+    #[test]
+    fn test_variant_object_child_data_0_byte_offsets_minus_one() {
+        test_variant_object_with_large_data(10, 0);
+    }
+
+    #[test]
+    fn test_variant_object_256_bytes_child_data_3_byte_offsets() {
+        test_variant_object_with_large_data(256 + 1, 1); // 2^8 - 2^16 elements
+    }
+
+    #[test]
+    fn test_variant_object_16777216_bytes_child_data_4_byte_offsets() {
+        test_variant_object_with_large_data(65536 + 1, 2); // 2^16 - 2^24 elements
+    }
+
+    #[test]
+    fn test_variant_object_65535_bytes_child_data_2_byte_offsets() {
+        test_variant_object_with_large_data(16777216 + 1, 3); // 2^24
+    }
 }