use itertools::Itertools;
use vortex_error::VortexResult;
use crate::IntoArray;
use crate::array::ArrayView;
use crate::array::ValidityVTable;
use crate::arrays::Chunked;
use crate::arrays::ChunkedArray;
use crate::arrays::chunked::ChunkedArrayExt;
use crate::dtype::DType;
use crate::dtype::Nullability;
use crate::validity::Validity;
impl ValidityVTable<Chunked> for Chunked {
fn validity(array: ArrayView<'_, Chunked>) -> VortexResult<Validity> {
let validities: Vec<Validity> =
array.chunks().iter().map(|c| c.validity()).try_collect()?;
match validities.first() {
None => return Ok(array.dtype().nullability().into()),
Some(first) if !matches!(first, Validity::Array(_)) => {
let target = std::mem::discriminant(first);
if validities
.iter()
.all(|v| std::mem::discriminant(v) == target)
{
return Ok(first.clone());
}
}
_ => {
}
}
Ok(Validity::Array(
unsafe {
ChunkedArray::new_unchecked(
validities
.into_iter()
.zip(array.iter_chunks())
.map(|(v, chunk)| v.to_array(chunk.len()))
.collect(),
DType::Bool(Nullability::NonNullable),
)
}
.into_array(),
))
}
}