Added Union::new_empty

jorgecarleitao · Aug 13, 2021 · a0956c5 · a0956c5
1 parent d53f903
commit a0956c5
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 6 deletions.
diff --git a/src/array/ffi.rs b/src/array/ffi.rs
@@ -85,7 +85,7 @@ pub fn buffers_children_dictionary(array: &dyn Array) -> BuffersChildren {
         DataType::LargeList(_) => ffi_dyn!(array, ListArray::<i64>),
         DataType::FixedSizeList(_, _) => ffi_dyn!(array, FixedSizeListArray),
         DataType::Struct(_) => ffi_dyn!(array, StructArray),
-        DataType::Union(_, _, _) => unimplemented!(),
+        DataType::Union(_, _, _) => todo!(),
         DataType::Dictionary(key_type, _) => match key_type.as_ref() {
             DataType::Int8 => ffi_dict_dyn!(array, DictionaryArray::<i8>),
             DataType::Int16 => ffi_dict_dyn!(array, DictionaryArray::<i16>),

diff --git a/src/array/mod.rs b/src/array/mod.rs
@@ -242,7 +242,7 @@ pub fn new_empty_array(data_type: DataType) -> Box<dyn Array> {
         DataType::LargeList(_) => Box::new(ListArray::<i64>::new_empty(data_type)),
         DataType::FixedSizeList(_, _) => Box::new(FixedSizeListArray::new_empty(data_type)),
         DataType::Struct(fields) => Box::new(StructArray::new_empty(&fields)),
-        DataType::Union(_, _, _) => unimplemented!(),
+        DataType::Union(_, _, _) => Box::new(UnionArray::new_empty(data_type)),
         DataType::Dictionary(key_type, value_type) => match key_type.as_ref() {
             DataType::Int8 => Box::new(DictionaryArray::<i8>::new_empty(*value_type)),
             DataType::Int16 => Box::new(DictionaryArray::<i16>::new_empty(*value_type)),
@@ -296,7 +296,7 @@ pub fn new_null_array(data_type: DataType, length: usize) -> Box<dyn Array> {
         DataType::LargeList(_) => Box::new(ListArray::<i64>::new_null(data_type, length)),
         DataType::FixedSizeList(_, _) => Box::new(FixedSizeListArray::new_null(data_type, length)),
         DataType::Struct(fields) => Box::new(StructArray::new_null(&fields, length)),
-        DataType::Union(_, _, _) => unimplemented!(),
+        DataType::Union(_, _, _) => todo!(),
         DataType::Dictionary(key_type, value_type) => match key_type.as_ref() {
             DataType::Int8 => Box::new(DictionaryArray::<i8>::new_null(*value_type, length)),
             DataType::Int16 => Box::new(DictionaryArray::<i16>::new_null(*value_type, length)),
@@ -357,7 +357,7 @@ pub fn clone(array: &dyn Array) -> Box<dyn Array> {
         DataType::LargeList(_) => clone_dyn!(array, ListArray::<i64>),
         DataType::FixedSizeList(_, _) => clone_dyn!(array, FixedSizeListArray),
         DataType::Struct(_) => clone_dyn!(array, StructArray),
-        DataType::Union(_, _, _) => unimplemented!(),
+        DataType::Union(_, _, _) => clone_dyn!(array, UnionArray),
         DataType::Dictionary(key_type, _) => match key_type.as_ref() {
             DataType::Int8 => clone_dyn!(array, DictionaryArray::<i8>),
             DataType::Int16 => clone_dyn!(array, DictionaryArray::<i16>),
@@ -503,6 +503,8 @@ mod tests {
             DataType::Utf8,
             DataType::Binary,
             DataType::List(Box::new(Field::new("a", DataType::Binary, true))),
+            DataType::Union(vec![Field::new("a", DataType::Binary, true)], None, true),
+            DataType::Union(vec![Field::new("a", DataType::Binary, true)], None, false),
         ];
         let a = datatypes.into_iter().all(|x| new_empty_array(x).len() == 0);
         assert!(a);

diff --git a/src/array/union/mod.rs b/src/array/union/mod.rs
@@ -7,7 +7,7 @@ use crate::{
     scalar::{new_scalar, Scalar},
 };
 
-use super::Array;
+use super::{new_empty_array, Array};
 
 mod iterator;
 
@@ -31,6 +31,32 @@ pub struct UnionArray {
 }
 
 impl UnionArray {
+    pub fn new_empty(data_type: DataType) -> Self {
+        if let DataType::Union(f, _, is_sparse) = &data_type {
+            let fields = f
+                .iter()
+                .map(|x| new_empty_array(x.data_type().clone()).into())
+                .collect();
+
+            let offsets = if *is_sparse {
+                None
+            } else {
+                Some(Buffer::new())
+            };
+
+            Self {
+                data_type,
+                fields_hash: HashMap::new(),
+                fields,
+                offsets,
+                types: Buffer::new(),
+                offset: 0,
+            }
+        } else {
+            panic!("Union struct must be created with the corresponding Union DataType")
+        }
+    }
+
     pub fn from_data(
         data_type: DataType,
         types: Buffer<i8>,

diff --git a/src/io/json/read/deserialize.rs b/src/io/json/read/deserialize.rs
@@ -255,7 +255,6 @@ pub fn read(rows: &[&Value], data_type: DataType) -> Arc<dyn Array> {
         /*
         DataType::FixedSizeBinary(_) => Box::new(FixedSizeBinaryArray::new_empty(data_type)),
         DataType::FixedSizeList(_, _) => Box::new(FixedSizeListArray::new_empty(data_type)),
-        DataType::Union(_, _, _) => unimplemented!(),
         DataType::Decimal(_, _) => Box::new(PrimitiveArray::<i128>::new_empty(data_type)),
         */
     }