From d806ae83174c55f25f98f8af04d08eb5cc1d1b56 Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Sat, 14 Aug 2021 07:52:31 +0000
Subject: [PATCH] Added UnionArray::new_null

---
 src/array/mod.rs       |  2 +-
 src/array/union/mod.rs | 24 +++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/array/mod.rs b/src/array/mod.rs
index 6cf45545bd2..37aa8f7d318 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -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(_, _, _) => todo!(),
+        DataType::Union(_, _, _) => Box::new(UnionArray::new_null(data_type, length)),
         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)),
diff --git a/src/array/union/mod.rs b/src/array/union/mod.rs
index 175827cdd8b..bccf74812e8 100644
--- a/src/array/union/mod.rs
+++ b/src/array/union/mod.rs
@@ -1,7 +1,7 @@
 use std::{collections::HashMap, sync::Arc};
 
 use crate::{
-    array::{display::get_value_display, display_fmt, new_empty_array, Array},
+    array::{display::get_value_display, display_fmt, new_empty_array, new_null_array, Array},
     bitmap::Bitmap,
     buffer::Buffer,
     datatypes::{DataType, Field},
@@ -35,6 +35,28 @@ pub struct UnionArray {
 }
 
 impl UnionArray {
+    pub fn new_null(data_type: DataType, length: usize) -> Self {
+        if let DataType::Union(f, _, is_sparse) = &data_type {
+            let fields = f
+                .iter()
+                .map(|x| new_null_array(x.data_type().clone(), length).into())
+                .collect();
+
+            let offsets = if *is_sparse {
+                None
+            } else {
+                Some((0..length as i32).collect::<Buffer<i32>>())
+            };
+
+            // all from the same field
+            let types = Buffer::new_zeroed(length);
+
+            Self::from_data(data_type, types, fields, offsets)
+        } else {
+            panic!("Union struct must be created with the corresponding Union DataType")
+        }
+    }
+
     pub fn new_empty(data_type: DataType) -> Self {
         if let DataType::Union(f, _, is_sparse) = &data_type {
             let fields = f