use vortex_error::VortexResult;
use super::Dict;
use crate::IntoArray;
use crate::array::ArrayView;
use crate::array::ValidityVTable;
use crate::arrays::DictArray;
use crate::arrays::dict::DictArraySlotsExt;
use crate::builtins::ArrayBuiltins;
use crate::dtype::Nullability;
use crate::scalar::Scalar;
use crate::validity::Validity;
impl ValidityVTable<Dict> for Dict {
fn validity(array: ArrayView<'_, Dict>) -> VortexResult<Validity> {
Ok(
match (array.codes().validity()?, array.values().validity()?) {
(
Validity::NonNullable | Validity::AllValid,
Validity::NonNullable | Validity::AllValid,
) => {
Validity::AllValid
}
(Validity::AllInvalid, _) | (_, Validity::AllInvalid) => Validity::AllInvalid,
(Validity::Array(codes_validity), Validity::NonNullable | Validity::AllValid) => {
Validity::Array(codes_validity)
}
(Validity::AllValid | Validity::NonNullable, Validity::Array(values_validity)) => {
let codes = array.codes().cast(array.codes().dtype().as_nonnullable())?;
Validity::Array(
unsafe { DictArray::new_unchecked(codes, values_validity) }.into_array(),
)
}
(Validity::Array(_codes_validity), Validity::Array(values_validity)) => {
let values_valid_mask =
unsafe { DictArray::new_unchecked(array.codes().clone(), values_validity) }
.into_array();
let values_valid_mask = values_valid_mask
.fill_null(Scalar::bool(false, Nullability::NonNullable))?;
Validity::Array(values_valid_mask)
}
},
)
}
}