use vortex_error::{VortexResult, vortex_err};
use vortex_scalar::Scalar;
use crate::arrays::{BoolArray, BoolVTable, ConstantArray};
use crate::compute::{FillNullKernel, FillNullKernelAdapter};
use crate::validity::Validity;
use crate::vtable::ValidityHelper;
use crate::{ArrayRef, IntoArray, ToCanonical, register_kernel};
impl FillNullKernel for BoolVTable {
fn fill_null(&self, array: &BoolArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
let fill = fill_value
.as_bool()
.value()
.ok_or_else(|| vortex_err!("Fill value must be non null"))?;
Ok(match array.validity() {
Validity::NonNullable | Validity::AllValid => BoolArray::from_bool_buffer(
array.boolean_buffer().clone(),
fill_value.dtype().nullability().into(),
)
.into_array(),
Validity::AllInvalid => {
ConstantArray::new(fill_value.clone(), array.len()).into_array()
}
Validity::Array(v) => {
let bool_buffer = if fill {
array.boolean_buffer() | &!v.to_bool().boolean_buffer()
} else {
array.boolean_buffer() & v.to_bool().boolean_buffer()
};
BoolArray::from_bool_buffer(bool_buffer, fill_value.dtype().nullability().into())
.into_array()
}
})
}
}
register_kernel!(FillNullKernelAdapter(BoolVTable).lift());
#[cfg(test)]
mod tests {
use arrow_buffer::BooleanBuffer;
use rstest::rstest;
use vortex_dtype::{DType, Nullability};
use crate::arrays::BoolArray;
use crate::canonical::ToCanonical;
use crate::compute::fill_null;
use crate::validity::Validity;
#[rstest]
#[case(true, vec![true, true, false, true])]
#[case(false, vec![true, false, false, false])]
fn bool_fill_null(#[case] fill_value: bool, #[case] expected: Vec<bool>) {
let bool_array = BoolArray::from_bool_buffer(
BooleanBuffer::from_iter([true, true, false, false]),
Validity::from_iter([true, false, true, false]),
);
let non_null_array = fill_null(bool_array.as_ref(), &fill_value.into())
.unwrap()
.to_bool();
assert_eq!(
non_null_array.boolean_buffer().iter().collect::<Vec<_>>(),
expected
);
assert_eq!(
non_null_array.dtype(),
&DType::Bool(Nullability::NonNullable)
);
}
}