use polars_error::PolarsResult;
use crate::array::{ArrayRef, FixedSizeListArray, NullArray, new_null_array};
use crate::bitmap::BitmapBuilder;
use crate::compute::concatenate::concatenate_unchecked;
use crate::datatypes::ArrowDataType;
use crate::legacy::array::{convert_inner_type, is_nested_null};
#[derive(Default)]
pub struct AnonymousBuilder {
arrays: Vec<ArrayRef>,
validity: Option<BitmapBuilder>,
length: usize,
pub width: usize,
}
impl AnonymousBuilder {
pub fn new(capacity: usize, width: usize) -> Self {
Self {
arrays: Vec::with_capacity(capacity),
validity: None,
width,
length: 0,
}
}
pub fn is_empty(&self) -> bool {
self.arrays.is_empty()
}
#[inline]
pub fn push(&mut self, arr: ArrayRef) {
self.arrays.push(arr);
if let Some(validity) = &mut self.validity {
validity.push(true)
}
self.length += 1;
}
pub fn push_null(&mut self) {
self.arrays
.push(NullArray::new(ArrowDataType::Null, self.width).boxed());
match &mut self.validity {
Some(validity) => validity.push(false),
None => self.init_validity(),
}
self.length += 1;
}
fn init_validity(&mut self) {
let mut validity = BitmapBuilder::with_capacity(self.arrays.capacity());
if !self.arrays.is_empty() {
validity.extend_constant(self.arrays.len() - 1, true);
validity.push(false);
}
self.validity = Some(validity)
}
pub fn finish(self, inner_dtype: Option<&ArrowDataType>) -> PolarsResult<FixedSizeListArray> {
let mut inner_dtype = inner_dtype.unwrap_or_else(|| self.arrays[0].dtype());
if is_nested_null(inner_dtype) {
for arr in &self.arrays {
if !is_nested_null(arr.dtype()) {
inner_dtype = arr.dtype();
break;
}
}
};
let arrays = self
.arrays
.iter()
.map(|arr| {
if matches!(arr.dtype(), ArrowDataType::Null) {
new_null_array(inner_dtype.clone(), arr.len())
} else if is_nested_null(arr.dtype()) {
convert_inner_type(&**arr, inner_dtype)
} else {
arr.to_boxed()
}
})
.collect::<Vec<_>>();
let values = concatenate_unchecked(&arrays)?;
let dtype = FixedSizeListArray::default_datatype(inner_dtype.clone(), self.width);
Ok(FixedSizeListArray::new(
dtype,
self.length,
values,
self.validity
.and_then(|validity| validity.into_opt_validity()),
))
}
}