use arrow_array::{Array, ArrowPrimitiveType, PrimitiveArray};
use arrow_buffer::{BooleanBuffer, NullBuffer, ScalarBuffer};
use arrow_select::filter::FilterPredicate;
use crate::{
Result,
util::{bitmap, kernel},
value::{Value, is::IsNumber, to_value::ToValue, value_type::get::GetType},
};
pub fn get_value<A>(array: &PrimitiveArray<A>, index: usize) -> Value
where
A: ArrowPrimitiveType,
A::Native: IsNumber + GetType,
{
if index < array.len() {
array.value(index).to_value()
} else {
Value::none_of(A::Native::get_type())
}
}
pub fn as_string<A>(array: &PrimitiveArray<A>, index: usize) -> String
where
A: ArrowPrimitiveType,
A::Native: IsNumber,
{
if index < array.len() {
array.value(index).to_string()
} else {
"none".to_string()
}
}
pub fn slice<A>(array: &PrimitiveArray<A>, start: usize, end: usize) -> PrimitiveArray<A>
where
A: ArrowPrimitiveType,
{
let end = end.min(array.len());
let start = start.min(end);
array.slice(start, end - start)
}
pub fn take<A>(array: &PrimitiveArray<A>, num: usize) -> PrimitiveArray<A>
where
A: ArrowPrimitiveType,
{
slice(array, 0, num)
}
pub fn filter<A>(array: &PrimitiveArray<A>, mask: &BooleanBuffer) -> PrimitiveArray<A>
where
A: ArrowPrimitiveType,
{
filter_with(array, &kernel::predicate(mask, array.len()))
}
pub fn filter_with<A>(array: &PrimitiveArray<A>, predicate: &FilterPredicate) -> PrimitiveArray<A>
where
A: ArrowPrimitiveType,
{
let selected = kernel::filtered(array, predicate);
let nulls = kernel::kept_nulls(array, &selected);
let (data_type, values, _) = selected.into_parts();
PrimitiveArray::new(values, nulls).with_data_type(data_type)
}
pub fn reorder<A>(array: &PrimitiveArray<A>, indices: &[usize]) -> Result<PrimitiveArray<A>>
where
A: ArrowPrimitiveType,
{
kernel::rows_in_range(indices, array.len())?;
let values = array.values();
let reordered: Vec<A::Native> = indices.iter().map(|&idx| values[idx]).collect();
Ok(PrimitiveArray::new(
ScalarBuffer::from(reordered),
bitmap::reorder_nulls(array.logical_nulls().as_ref(), indices),
)
.with_data_type(array.data_type().clone()))
}
pub fn attach_nulls<A>(array: PrimitiveArray<A>, nulls: Option<NullBuffer>) -> PrimitiveArray<A>
where
A: ArrowPrimitiveType,
{
bitmap::assert_nulls_len(nulls.as_ref(), array.len());
let (data_type, values, _) = array.into_parts();
PrimitiveArray::new(values, nulls).with_data_type(data_type)
}
pub fn capacity<A>(array: &PrimitiveArray<A>) -> usize
where
A: ArrowPrimitiveType,
{
let buffer = array.values().inner();
if buffer.strong_count() == 1 {
buffer.capacity() / size_of::<A::Native>()
} else {
array.len()
}
}
pub fn heap_size<A>(array: &PrimitiveArray<A>) -> usize
where
A: ArrowPrimitiveType,
{
capacity(array) * size_of::<A::Native>()
}