Skip to content

Commit

Permalink
Add into_parts-type functions for consuming PrimitiveArray and Boolea…
Browse files Browse the repository at this point in the history
…nArray
  • Loading branch information
srh committed Nov 11, 2024
1 parent 24097bd commit 2de3c3b
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 0 deletions.
7 changes: 7 additions & 0 deletions arrow/src/array/array_boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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 {
Expand Down
4 changes: 4 additions & 0 deletions arrow/src/array/array_primitive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down
9 changes: 9 additions & 0 deletions arrow/src/array/data.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 2de3c3b

Please sign in to comment.