feat: add dyn arrow_array::Array conversions for union arrays (#232)

mbrobbel · Sep 3, 2024 · 69b83cd · 69b83cd
1 parent a2eb1d9
commit 69b83cd
Showing 1 changed file with 92 additions and 2 deletions.
diff --git a/src/arrow/array/union.rs b/src/arrow/array/union.rs
@@ -100,6 +100,25 @@ where
     }
 }
 
+impl<
+        T: UnionArrayType<VARIANTS>,
+        const VARIANTS: usize,
+        Buffer: BufferType,
+        OffsetItem: OffsetElement,
+    > From<UnionArray<T, VARIANTS, SparseLayout, Buffer, OffsetItem>>
+    for Arc<dyn arrow_array::Array>
+where
+    for<'a> i8: From<&'a T>,
+    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>:
+        UnionArrayTypeFields<VARIANTS> + Into<Vec<Arc<dyn arrow_array::Array>>>,
+    arrow_buffer::ScalarBuffer<i8>: From<FixedSizePrimitiveArray<i8, false, Buffer>>,
+    UnionArray<T, VARIANTS, SparseLayout, Buffer, OffsetItem>: crate::arrow::Array,
+{
+    fn from(value: UnionArray<T, VARIANTS, SparseLayout, Buffer, OffsetItem>) -> Self {
+        Arc::new(arrow_array::UnionArray::from(value))
+    }
+}
+
 impl<
         T: UnionArrayType<VARIANTS>,
         const VARIANTS: usize,
@@ -132,6 +151,25 @@ where
     }
 }
 
+impl<
+        T: UnionArrayType<VARIANTS>,
+        const VARIANTS: usize,
+        Buffer: BufferType,
+        OffsetItem: OffsetElement,
+    > From<UnionArray<T, VARIANTS, DenseLayout, Buffer, OffsetItem>> for Arc<dyn arrow_array::Array>
+where
+    for<'a> i8: From<&'a T>,
+    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>:
+        UnionArrayTypeFields<VARIANTS> + Into<Vec<Arc<dyn arrow_array::Array>>>,
+    arrow_buffer::ScalarBuffer<i8>: From<FixedSizePrimitiveArray<i8, false, Buffer>>,
+    arrow_buffer::ScalarBuffer<i32>: From<FixedSizePrimitiveArray<i32, false, Buffer>>,
+    UnionArray<T, VARIANTS, SparseLayout, Buffer, OffsetItem>: crate::arrow::Array,
+{
+    fn from(value: UnionArray<T, VARIANTS, DenseLayout, Buffer, OffsetItem>) -> Self {
+        Arc::new(arrow_array::UnionArray::from(value))
+    }
+}
+
 impl<
         T: UnionArrayType<VARIANTS>,
         const VARIANTS: usize,
@@ -156,6 +194,25 @@ where
     }
 }
 
+impl<
+        T: UnionArrayType<VARIANTS>,
+        const VARIANTS: usize,
+        Buffer: BufferType,
+        OffsetItem: OffsetElement,
+    > From<Arc<dyn arrow_array::Array>>
+    for UnionArray<T, VARIANTS, SparseLayout, Buffer, OffsetItem>
+where
+    for<'a> i8: From<&'a T>,
+    FixedSizePrimitiveArray<i8, false, Buffer>: From<arrow_buffer::ScalarBuffer<i8>>,
+    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, SparseLayout>:
+        FromIterator<Arc<dyn arrow_array::Array>>,
+{
+    fn from(value: Arc<dyn arrow_array::Array>) -> Self {
+        let array = arrow_array::UnionArray::from(value.to_data());
+        Self::from(array)
+    }
+}
+
 impl<
         T: UnionArrayType<VARIANTS>,
         const VARIANTS: usize,
@@ -182,20 +239,53 @@ where
     }
 }
 
+impl<
+        T: UnionArrayType<VARIANTS>,
+        const VARIANTS: usize,
+        Buffer: BufferType,
+        OffsetItem: OffsetElement,
+    > From<Arc<dyn arrow_array::Array>> for UnionArray<T, VARIANTS, DenseLayout, Buffer, OffsetItem>
+where
+    for<'a> i8: From<&'a T>,
+    FixedSizePrimitiveArray<i8, false, Buffer>: From<arrow_buffer::ScalarBuffer<i8>>,
+    FixedSizePrimitiveArray<i32, false, Buffer>: From<arrow_buffer::ScalarBuffer<i32>>,
+    <T as UnionArrayType<VARIANTS>>::Array<Buffer, OffsetItem, DenseLayout>:
+        FromIterator<Arc<dyn arrow_array::Array>>,
+{
+    fn from(value: Arc<dyn arrow_array::Array>) -> Self {
+        let array = arrow_array::UnionArray::from(value.to_data());
+        Self::from(array)
+    }
+}
+
 #[cfg(test)]
 #[cfg(feature = "derive")]
 mod tests {
-    use crate::Length;
+    use arrow_array::RecordBatch;
+
+    use crate::{array::StructArray, Length};
 
     use super::*;
 
-    #[derive(crate::ArrayType, Clone)]
+    #[derive(crate::ArrayType, Clone, Debug, PartialEq)]
     enum FooBar {
         Foo,
         Bar(u8),
         Baz { a: bool },
     }
 
+    #[derive(crate::ArrayType, Clone, Debug, PartialEq)]
+    struct Wrap(FooBar);
+
+    #[test]
+    fn via_dyn_array() {
+        let input = [Wrap(FooBar::Foo), Wrap(FooBar::Bar(123))];
+        let struct_array = input.clone().into_iter().collect::<StructArray<Wrap>>();
+        let record_batch = RecordBatch::from(struct_array);
+        let read = StructArray::<Wrap>::from(record_batch);
+        assert_eq!(read.into_iter().collect::<Vec<_>>(), input);
+    }
+
     #[test]
     fn from() {
         let input = [