use std::sync::Arc;
use crate::array::BooleanArray;
use crate::datatypes::DataType;
use crate::error::Error;
use crate::{
array::{FromFfi, PrimitiveArray},
types::NativeType,
};
use super::{ArrowArray, InternalArrowArray};
#[allow(dead_code)]
struct PrivateData<T> {
data: T,
buffers_ptr: Box<[*const std::os::raw::c_void]>,
children_ptr: Box<[*mut ArrowArray]>,
dictionary_ptr: Option<*mut ArrowArray>,
}
pub(crate) unsafe fn create_array<
T: AsRef<[u8]>,
I: Iterator<Item = Option<*const u8>>,
II: Iterator<Item = ArrowArray>,
>(
data: Arc<T>,
num_rows: usize,
null_count: usize,
buffers: I,
children: II,
dictionary: Option<ArrowArray>,
offset: Option<usize>,
) -> ArrowArray {
let buffers_ptr = buffers
.map(|maybe_buffer| match maybe_buffer {
Some(b) => b as *const std::os::raw::c_void,
None => std::ptr::null(),
})
.collect::<Box<[_]>>();
let n_buffers = buffers_ptr.len() as i64;
let children_ptr = children
.map(|child| Box::into_raw(Box::new(child)))
.collect::<Box<_>>();
let n_children = children_ptr.len() as i64;
let dictionary_ptr = dictionary.map(|array| Box::into_raw(Box::new(array)));
let mut private_data = Box::new(PrivateData::<Arc<T>> {
data,
buffers_ptr,
children_ptr,
dictionary_ptr,
});
ArrowArray {
length: num_rows as i64,
null_count: null_count as i64,
offset: offset.unwrap_or(0) as i64, n_buffers,
n_children,
buffers: private_data.buffers_ptr.as_mut_ptr(),
children: private_data.children_ptr.as_mut_ptr(),
dictionary: private_data.dictionary_ptr.unwrap_or(std::ptr::null_mut()),
release: Some(release::<Arc<T>>),
private_data: Box::into_raw(private_data) as *mut ::std::os::raw::c_void,
}
}
unsafe extern "C" fn release<T>(array: *mut ArrowArray) {
if array.is_null() {
return;
}
let array = &mut *array;
let private = Box::from_raw(array.private_data as *mut PrivateData<T>);
for child in private.children_ptr.iter() {
let _ = Box::from_raw(*child);
}
if let Some(ptr) = private.dictionary_ptr {
let _ = Box::from_raw(ptr);
}
array.release = None;
}
pub unsafe fn slice<T: NativeType>(slice: &[T]) -> PrimitiveArray<T> {
let num_rows = slice.len();
let null_count = 0;
let validity = None;
let data: &[u8] = bytemuck::cast_slice(slice);
let ptr = data.as_ptr();
let data = Arc::new(data);
let array = create_array(
data,
num_rows,
null_count,
[validity, Some(ptr)].into_iter(),
[].into_iter(),
None,
None,
);
let array = InternalArrowArray::new(array, T::PRIMITIVE.into());
unsafe { PrimitiveArray::<T>::try_from_ffi(array) }.unwrap()
}
pub unsafe fn bitmap(data: &[u8], offset: usize, length: usize) -> Result<BooleanArray, Error> {
if offset >= 8 {
return Err(Error::InvalidArgumentError("offset should be < 8".into()));
};
if length > data.len() * 8 - offset {
return Err(Error::InvalidArgumentError("given length is oob".into()));
}
let null_count = 0;
let validity = None;
let ptr = data.as_ptr();
let data = Arc::new(data);
let array = create_array(
data,
length,
null_count,
[validity, Some(ptr)].into_iter(),
[].into_iter(),
None,
Some(offset),
);
let array = InternalArrowArray::new(array, DataType::Boolean);
Ok(unsafe { BooleanArray::try_from_ffi(array) }.unwrap())
}