Fix incorrect buffers for UnionArray

apache · Apr 14, 2022 · d24bae3 · d24bae3
1 parent eb6b7c6
commit d24bae3
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/arrow/src/array/data.rs b/arrow/src/array/data.rs
@@ -219,8 +219,14 @@ pub(crate) fn into_buffers(
         DataType::Utf8
         | DataType::Binary
         | DataType::LargeUtf8
-        | DataType::LargeBinary
-        | DataType::Union(_, _) => vec![buffer1.into(), buffer2.into()],
+        | DataType::LargeBinary => vec![buffer1.into(), buffer2.into()],
+        DataType::Union(_, mode) => {
+            match mode {
+                // Based on Union's DataTypeLayout
+                UnionMode::Sparse => vec![buffer1.into()],
+                UnionMode::Dense => vec![buffer1.into(), buffer2.into()],
+            }
+        }
         _ => vec![buffer1.into()],
     }
 }