Correctly update child-offsets in GrowableUnion (jorgecarleitao#1360)

src/array/growable/binary.rs

@@ -81,6 +81,11 @@ impl<'a, O: Offset> Growable<'a> for GrowableBinary<'a, O> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.offsets.len() - 1
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         self.to().arced()
     }

src/array/growable/boolean.rs

@@ -71,6 +71,11 @@ impl<'a> Growable<'a> for GrowableBoolean<'a> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.values.len()
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(self.to())
     }

src/array/growable/dictionary.rs

@@ -132,6 +132,11 @@ impl<'a, T: DictionaryKey> Growable<'a> for GrowableDictionary<'a, T> {
         );
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.key_values.len()
+    }
+
     #[inline]
     fn extend_validity(&mut self, additional: usize) {
         self.key_values

src/array/growable/fixed_binary.rs

@@ -78,6 +78,11 @@ impl<'a> Growable<'a> for GrowableFixedSizeBinary<'a> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.values.len() / self.size
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(self.to())
     }

src/array/growable/fixed_size_list.rs

@@ -85,6 +85,11 @@ impl<'a> Growable<'a> for GrowableFixedSizeList<'a> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.values.len() / self.size
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(self.to())
     }

src/array/growable/list.rs

@@ -97,6 +97,11 @@ impl<'a, O: Offset> Growable<'a> for GrowableList<'a, O> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.offsets.len() - 1
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(self.to())
     }

src/array/growable/mod.rs

@@ -43,6 +43,9 @@ pub trait Growable<'a> {
     /// Extends this [`Growable`] with null elements, disregarding the bound arrays
     fn extend_validity(&mut self, additional: usize);
 
+    /// The current length of the [`Growable`].
+    fn len(&self) -> usize;
+
     /// Converts this [`Growable`] to an [`Arc<dyn Array>`], thereby finishing the mutation.
     /// Self will be empty after such operation.
     fn as_arc(&mut self) -> Arc<dyn Array> {

src/array/growable/null.rs

@@ -38,6 +38,11 @@ impl<'a> Growable<'a> for GrowableNull {
         self.length += additional;
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.length
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(NullArray::new(self.data_type.clone(), self.length))
     }

src/array/growable/primitive.rs

@@ -82,6 +82,11 @@ impl<'a, T: NativeType> Growable<'a> for GrowablePrimitive<'a, T> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.values.len()
+    }
+
     #[inline]
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(self.to())

src/array/growable/structure.rs

@@ -104,6 +104,18 @@ impl<'a> Growable<'a> for GrowableStruct<'a> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        // All children should have the same indexing, so just use the first
+        // one. If we don't have children, we might still have a validity
+        // array, so use that.
+        if let Some(child) = self.values.get(0) {
+            child.len()
+        } else {
+            self.validity.len()
+        }
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(self.to())
     }

src/array/growable/union.rs

@@ -73,10 +73,15 @@ impl<'a> Growable<'a> for GrowableUnion<'a> {
         if let Some(x) = self.offsets.as_mut() {
             let offsets = &array.offsets().unwrap()[start..start + len];
 
-            x.extend(offsets);
             // in a dense union, each slot has its own offset. We extend the fields accordingly.
             for (&type_, &offset) in types.iter().zip(offsets.iter()) {
-                self.fields[type_ as usize].extend(index, offset as usize, 1);
+                let field = &mut self.fields[type_ as usize];
+                // The offset for the element that is about to be extended is the current length
+                // of the child field of the corresponding type. Note that this may be very
+                // different than the original offset from the array we are extending from as
+                // it is a function of the previous extensions to this child.
+                x.push(field.len() as i32);
+                field.extend(index, offset as usize, 1);
             }
         } else {
             // in a sparse union, every field has the same length => extend all fields equally
@@ -88,6 +93,11 @@ impl<'a> Growable<'a> for GrowableUnion<'a> {
 
     fn extend_validity(&mut self, _additional: usize) {}
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.types.len()
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         self.to().arced()
     }

src/array/growable/utf8.rs

@@ -88,6 +88,11 @@ impl<'a, O: Offset> Growable<'a> for GrowableUtf8<'a, O> {
         self.validity.extend_constant(additional, false);
     }
 
+    #[inline]
+    fn len(&self) -> usize {
+        self.offsets.len() - 1
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(self.to())
     }

tests/it/array/growable/binary.rs

@@ -10,6 +10,7 @@ fn no_offsets() {
     let mut a = GrowableBinary::new(vec![&array], false, 0);
 
     a.extend(0, 1, 2);
+    assert_eq!(a.len(), 2);
 
     let result: BinaryArray<i32> = a.into();
 
@@ -27,6 +28,7 @@ fn with_offsets() {
     let mut a = GrowableBinary::new(vec![&array], false, 0);
 
     a.extend(0, 0, 3);
+    assert_eq!(a.len(), 3);
 
     let result: BinaryArray<i32> = a.into();
 
@@ -42,6 +44,7 @@ fn test_string_offsets() {
     let mut a = GrowableBinary::new(vec![&array], false, 0);
 
     a.extend(0, 0, 3);
+    assert_eq!(a.len(), 3);
 
     let result: BinaryArray<i32> = a.into();
 
@@ -58,6 +61,7 @@ fn test_multiple_with_validity() {
 
     a.extend(0, 0, 2);
     a.extend(1, 0, 2);
+    assert_eq!(a.len(), 4);
 
     let result: BinaryArray<i32> = a.into();
 
@@ -74,6 +78,7 @@ fn test_string_null_offset_validity() {
 
     a.extend(0, 1, 2);
     a.extend_validity(1);
+    assert_eq!(a.len(), 3);
 
     let result: BinaryArray<i32> = a.into();
 

tests/it/array/growable/boolean.rs

@@ -8,6 +8,7 @@ fn test_bool() {
     let mut a = GrowableBoolean::new(vec![&array], false, 0);
 
     a.extend(0, 1, 2);
+    assert_eq!(a.len(), 2);
 
     let result: BooleanArray = a.into();
 

tests/it/array/growable/dictionary.rs

@@ -21,6 +21,7 @@ fn test_single() -> Result<()> {
     let mut growable = GrowableDictionary::new(&[&array], false, 0);
 
     growable.extend(0, 1, 2);
+    assert_eq!(growable.len(), 2);
 
     let result: DictionaryArray<i32> = growable.into();
 
@@ -56,6 +57,7 @@ fn test_multi() -> Result<()> {
 
     growable.extend(0, 1, 2);
     growable.extend(1, 1, 2);
+    assert_eq!(growable.len(), 4);
 
     let result: DictionaryArray<i32> = growable.into();
 

tests/it/array/growable/fixed_binary.rs

@@ -12,6 +12,7 @@ fn basic() {
     let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0);
 
     a.extend(0, 1, 2);
+    assert_eq!(a.len(), 2);
 
     let result: FixedSizeBinaryArray = a.into();
 
@@ -30,6 +31,7 @@ fn offsets() {
     let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0);
 
     a.extend(0, 0, 3);
+    assert_eq!(a.len(), 3);
 
     let result: FixedSizeBinaryArray = a.into();
 
@@ -46,6 +48,7 @@ fn multiple_with_validity() {
 
     a.extend(0, 0, 2);
     a.extend(1, 0, 2);
+    assert_eq!(a.len(), 4);
 
     let result: FixedSizeBinaryArray = a.into();
 
@@ -63,6 +66,7 @@ fn null_offset_validity() {
 
     a.extend(0, 1, 2);
     a.extend_validity(1);
+    assert_eq!(a.len(), 3);
 
     let result: FixedSizeBinaryArray = a.into();
 
@@ -81,6 +85,7 @@ fn sized_offsets() {
 
     a.extend(0, 1, 1);
     a.extend(0, 0, 1);
+    assert_eq!(a.len(), 2);
 
     let result: FixedSizeBinaryArray = a.into();
 

tests/it/array/growable/fixed_size_list.rs

@@ -21,6 +21,7 @@ fn basic() {
 
     let mut a = GrowableFixedSizeList::new(vec![&array], false, 0);
     a.extend(0, 0, 1);
+    assert_eq!(a.len(), 1);
 
     let result: FixedSizeListArray = a.into();
 
@@ -42,6 +43,7 @@ fn null_offset() {
 
     let mut a = GrowableFixedSizeList::new(vec![&array], false, 0);
     a.extend(0, 1, 1);
+    assert_eq!(a.len(), 1);
 
     let result: FixedSizeListArray = a.into();
 
@@ -70,6 +72,7 @@ fn test_from_two_lists() {
     let mut a = GrowableFixedSizeList::new(vec![&array_1, &array_2], false, 6);
     a.extend(0, 0, 2);
     a.extend(1, 1, 1);
+    assert_eq!(a.len(), 3);
 
     let result: FixedSizeListArray = a.into();
 

tests/it/array/growable/list.rs

@@ -33,6 +33,7 @@ fn extension() {
 
     let mut a = GrowableList::new(vec![&array_ext], false, 0);
     a.extend(0, 0, 1);
+    assert_eq!(a.len(), 1);
 
     let result: ListArray<i32> = a.into();
     assert_eq!(array_ext.data_type(), result.data_type());
@@ -51,6 +52,7 @@ fn basic() {
 
     let mut a = GrowableList::new(vec![&array], false, 0);
     a.extend(0, 0, 1);
+    assert_eq!(a.len(), 1);
 
     let result: ListArray<i32> = a.into();
 
@@ -72,6 +74,7 @@ fn null_offset() {
 
     let mut a = GrowableList::new(vec![&array], false, 0);
     a.extend(0, 1, 1);
+    assert_eq!(a.len(), 1);
 
     let result: ListArray<i32> = a.into();
 
@@ -93,6 +96,7 @@ fn null_offsets() {
 
     let mut a = GrowableList::new(vec![&array], false, 0);
     a.extend(0, 1, 1);
+    assert_eq!(a.len(), 1);
 
     let result: ListArray<i32> = a.into();
 
@@ -121,6 +125,7 @@ fn test_from_two_lists() {
     let mut a = GrowableList::new(vec![&array_1, &array_2], false, 6);
     a.extend(0, 0, 2);
     a.extend(1, 1, 1);
+    assert_eq!(a.len(), 3);
 
     let result: ListArray<i32> = a.into();
 

tests/it/array/growable/null.rs

@@ -12,6 +12,7 @@ fn null() {
 
     mutable.extend(0, 1, 2);
     mutable.extend(1, 0, 1);
+    assert_eq!(mutable.len(), 3);
 
     let result: NullArray = mutable.into();
 

tests/it/array/growable/primitive.rs

@@ -9,6 +9,7 @@ fn basics() {
     let b = PrimitiveArray::<u8>::from(vec![Some(1), Some(2), Some(3)]);
     let mut a = GrowablePrimitive::new(vec![&b], false, 3);
     a.extend(0, 0, 2);
+    assert_eq!(a.len(), 2);
     let result: PrimitiveArray<u8> = a.into();
     let expected = PrimitiveArray::<u8>::from(vec![Some(1), Some(2)]);
     assert_eq!(result, expected);
@@ -21,6 +22,7 @@ fn offset() {
     let b = b.slice(1, 2);
     let mut a = GrowablePrimitive::new(vec![&b], false, 2);
     a.extend(0, 0, 2);
+    assert_eq!(a.len(), 2);
     let result: PrimitiveArray<u8> = a.into();
     let expected = PrimitiveArray::<u8>::from(vec![Some(2), Some(3)]);
     assert_eq!(result, expected);
@@ -33,6 +35,7 @@ fn null_offset() {
     let b = b.slice(1, 2);
     let mut a = GrowablePrimitive::new(vec![&b], false, 2);
     a.extend(0, 0, 2);
+    assert_eq!(a.len(), 2);
     let result: PrimitiveArray<u8> = a.into();
     let expected = PrimitiveArray::<u8>::from(vec![None, Some(3)]);
     assert_eq!(result, expected);
@@ -46,6 +49,7 @@ fn null_offset_validity() {
     a.extend(0, 0, 2);
     a.extend_validity(3);
     a.extend(0, 1, 1);
+    assert_eq!(a.len(), 6);
     let result: PrimitiveArray<u8> = a.into();
     let expected = PrimitiveArray::<u8>::from(&[Some(2), Some(3), None, None, None, Some(3)]);
     assert_eq!(result, expected);
@@ -58,6 +62,7 @@ fn joining_arrays() {
     let mut a = GrowablePrimitive::new(vec![&b, &c], false, 4);
     a.extend(0, 0, 2);
     a.extend(1, 1, 2);
+    assert_eq!(a.len(), 4);
     let result: PrimitiveArray<u8> = a.into();
 
     let expected = PrimitiveArray::<u8>::from(&[Some(1), Some(2), Some(5), Some(6)]);