Add into_parts-type functions for consuming PrimitiveArray and Boolea…

…nArray
cube-js · Nov 11, 2024 · b8d1ee9 · b8d1ee9
1 parent 3e270a6
commit b8d1ee9
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 0 deletions.
diff --git a/arrow/src/array/array_boolean.rs b/arrow/src/array/array_boolean.rs
@@ -22,6 +22,7 @@ use std::{any::Any, fmt};
 
 use super::*;
 use super::{array::print_long_array, raw_pointer::RawPtrBox};
+use crate::bitmap::Bitmap;
 use crate::buffer::{Buffer, MutableBuffer};
 use crate::util::bit_util;
 
@@ -109,6 +110,12 @@ impl BooleanArray {
         debug_assert!(i < self.len());
         unsafe { self.value_unchecked(i) }
     }
+
+    /// Returns (_, _, offset, length)
+    pub fn into_parts(self) -> (Buffer, Option<Bitmap>, usize, usize) {
+        let data = self.data;
+        data.into_1_dimensional_parts()
+    }
 }
 
 impl Array for BooleanArray {

diff --git a/arrow/src/array/array_primitive.rs b/arrow/src/array/array_primitive.rs
@@ -146,6 +146,10 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
         );
         PrimitiveArray::from(data)
     }
+
+    pub fn into_data(self) -> ArrayData {
+        self.data
+    }
 }
 
 impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {

diff --git a/arrow/src/array/data.rs b/arrow/src/array/data.rs
@@ -478,6 +478,15 @@ impl ArrayData {
 
         Self::new(data_type.clone(), 0, Some(0), None, 0, buffers, child_data)
     }
+
+    pub fn into_1_dimensional_parts(self) -> (Buffer, Option<Bitmap>, usize, usize) {
+        let offset: usize = self.offset;
+        let length: usize = self.len;
+        let buffers: Vec<Buffer> = self.buffers;
+        let bitmap: Option<Bitmap> = self.null_bitmap;
+        let buffer0: Buffer = buffers.into_iter().next().unwrap();
+        (buffer0, bitmap, offset, length)
+    }
 }
 
 impl PartialEq for ArrayData {