From eec97cab55b14f3a8c7cebe197ea42e6da9a31a1 Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Fri, 13 Aug 2021 19:55:53 +0000
Subject: [PATCH] Added memory estimation Union.

---
 src/compute/aggregate/memory.rs | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/compute/aggregate/memory.rs b/src/compute/aggregate/memory.rs
index cbace0d8f7d..d49411c3b34 100644
--- a/src/compute/aggregate/memory.rs
+++ b/src/compute/aggregate/memory.rs
@@ -109,7 +109,22 @@ pub fn estimated_bytes_size(array: &dyn Array) -> usize {
                 .sum::<usize>()
                 + validity_size(array.validity())
         }
-        Union(_, _, _) => unreachable!(),
+        Union(_, _, _) => {
+            let array = array.as_any().downcast_ref::<UnionArray>().unwrap();
+            let types = array.types().len() * std::mem::size_of::<i8>();
+            let offsets = array
+                .offsets()
+                .as_ref()
+                .map(|x| x.len() * std::mem::size_of::<i32>())
+                .unwrap_or_default();
+            let fields = array
+                .fields()
+                .iter()
+                .map(|x| x.as_ref())
+                .map(estimated_bytes_size)
+                .sum::<usize>();
+            types + offsets + fields
+        }
         Dictionary(keys, _) => match keys.as_ref() {
             Int8 => dyn_dict!(array, i8),
             Int16 => dyn_dict!(array, i16),