Fix serialization issue with nested unions (#110)

* Expand enum test coverage with nested unions * Fix implementation of len() in generated MutableArray for Unions * Fix test formatting * Test for dense-inside-sparse * Handle pushing None in Dense unions * CI cleanup for updated rust stable (1.69)
DataEngineeringLabs · May 9, 2023 · 90de932 · 90de932
1 parent d517f07
commit 90de932
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 13 deletions.
diff --git a/arrow2_convert/tests/test_enum.rs b/arrow2_convert/tests/test_enum.rs
@@ -57,32 +57,68 @@ fn test_nested_unit_variant() {
 
     #[derive(Debug, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)]
     #[arrow_field(type = "dense")]
-    enum TestEnum {
+    enum TestDenseEnum {
         VAL1,
         VAL2(i32),
         VAL3(f64),
         VAL4(TestStruct),
-        VAL5(ChildEnum),
+        VAL5(DenseChildEnum),
+        VAL6(SparseChildEnum),
     }
 
     #[derive(Debug, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)]
     #[arrow_field(type = "sparse")]
-    enum ChildEnum {
+    enum TestSparseEnum {
+        VAL1,
+        VAL2(i32),
+        VAL3(f64),
+        VAL4(TestStruct),
+        VAL5(DenseChildEnum),
+        VAL6(SparseChildEnum),
+    }
+
+    #[derive(Debug, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)]
+    #[arrow_field(type = "dense")]
+    enum DenseChildEnum {
+        VAL1,
+        VAL2(i32),
+        VAL3(f64),
+        VAL4(TestStruct),
+    }
+
+    #[derive(Debug, PartialEq, ArrowField, ArrowSerialize, ArrowDeserialize)]
+    #[arrow_field(type = "sparse")]
+    enum SparseChildEnum {
         VAL1,
         VAL2(i32),
         VAL3(f64),
         VAL4(TestStruct),
     }
 
     let enums = vec![
-        TestEnum::VAL1,
-        TestEnum::VAL2(2),
-        TestEnum::VAL3(1.2),
-        TestEnum::VAL4(TestStruct { a1: 10 }),
+        TestDenseEnum::VAL1,
+        TestDenseEnum::VAL2(2),
+        TestDenseEnum::VAL3(1.2),
+        TestDenseEnum::VAL4(TestStruct { a1: 10 }),
+        TestDenseEnum::VAL5(DenseChildEnum::VAL4(TestStruct { a1: 17 })),
+        TestDenseEnum::VAL6(SparseChildEnum::VAL4(TestStruct { a1: 42 })),
     ];
 
     let b: Box<dyn Array> = enums.try_into_arrow().unwrap();
-    let round_trip: Vec<TestEnum> = b.try_into_collection().unwrap();
+    let round_trip: Vec<TestDenseEnum> = b.try_into_collection().unwrap();
+    assert_eq!(round_trip, enums);
+
+    let enums = vec![
+        TestSparseEnum::VAL1,
+        TestSparseEnum::VAL2(2),
+        TestSparseEnum::VAL3(1.2),
+        TestSparseEnum::VAL4(TestStruct { a1: 10 }),
+        TestSparseEnum::VAL5(DenseChildEnum::VAL4(TestStruct { a1: 17 })),
+        TestSparseEnum::VAL6(SparseChildEnum::VAL4(TestStruct { a1: 42 })),
+    ];
+
+    let b: Box<dyn Array> = enums.try_into_arrow().unwrap();
+    let round_trip: Vec<TestSparseEnum> = b.try_into_collection().unwrap();
     assert_eq!(round_trip, enums);
 }
 

diff --git a/arrow2_convert/tests/test_serialize.rs b/arrow2_convert/tests/test_serialize.rs
@@ -73,14 +73,14 @@ fn test_array() {
 #[test]
 fn test_buffer() {
     // Buffer<u8> and Vec<u8> should serialize into BinaryArray
-    let dat: Vec<Buffer<u8>> = vec![(0..10).into_iter().collect()];
+    let dat: Vec<Buffer<u8>> = vec![(0..10).collect()];
     let r: Box<dyn Array> = dat.try_into_arrow().unwrap();
     assert_eq!(r.len(), 1);
     assert_eq!(r.data_type(), &<Buffer<u8> as ArrowField>::data_type());
     assert_eq!(r.data_type(), &<Vec<u8> as ArrowField>::data_type());
 
     // Buffer<u16> and Vec<u16> should serialize into ListArray
-    let dat: Vec<Buffer<u16>> = vec![(0..10).into_iter().collect()];
+    let dat: Vec<Buffer<u16>> = vec![(0..10).collect()];
     let r: Box<dyn Array> = dat.try_into_arrow().unwrap();
     assert_eq!(r.len(), 1);
     assert_eq!(r.data_type(), &<Buffer<u16> as ArrowField>::data_type());

diff --git a/arrow2_convert/tests/ui/struct_incorrect_type.stderr b/arrow2_convert/tests/ui/struct_incorrect_type.stderr
@@ -11,7 +11,7 @@ error[E0308]: mismatched types
  --> tests/ui/struct_incorrect_type.rs:4:45
   |
 4 | #[derive(Debug, ArrowField, ArrowSerialize, ArrowDeserialize)]
-  |                                             ^^^^^^^^^^^^^^^^ expected struct `String`, found struct `Vec`
+  |                                             ^^^^^^^^^^^^^^^^ expected `String`, found `Vec<u8>`
   |
   = note: expected struct `String`
              found struct `Vec<u8>`

diff --git a/arrow2_convert_derive/src/derive_enum.rs b/arrow2_convert_derive/src/derive_enum.rs
@@ -119,7 +119,6 @@ pub fn expand_serialize(input: DeriveEnum) -> TokenStream {
         ..
     } = (&input).into();
 
-    let first_variant = &variant_names[0];
     let is_dense = input.is_dense;
 
     let mutable_array_name = &input.common.mutable_array_name();
@@ -227,6 +226,7 @@ pub fn expand_serialize(input: DeriveEnum) -> TokenStream {
         let first_name = &variant_names[0];
         quote! {
             self.types.push(0);
+            self.offsets.push((self.#first_name.len()) as i32);
             <#first_array_type as MutableArray>::push_null(&mut self.#first_name);
         }
     } else {
@@ -313,7 +313,7 @@ pub fn expand_serialize(input: DeriveEnum) -> TokenStream {
             }
 
             fn len(&self) -> usize {
-                self.#first_variant.len()
+                self.types.len()
             }
 
             fn validity(&self) -> Option<&arrow2::bitmap::MutableBitmap> {