16355473372

Committed 17 Jul 2025 08:38PM UTC coverage: 80.907% (-0.02%) from 80.926%

Build # 16355473372

Build Type

Pull #3909

github

Committed by

web-flow

Commit Message

Merge ed2add1a5 into ac7014e8c

Pull Request Pull Request #3909: feat: impl optimize for more arrays

Coverage Stats

51 of 82 new or added lines in 6 files covered. (62.2%)

41951 of 51851 relevant lines covered (80.91%)

157671.18 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

76.11

/vortex-array/src/validity.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Array validity and nullability behavior, used by arrays and compute functions.

use std::fmt::Debug;
use std::ops::{BitAnd, Not};

use arrow_buffer::{BooleanBuffer, NullBuffer};
use vortex_dtype::{DType, Nullability};
use vortex_error::{VortexExpect as _, VortexResult, vortex_bail, vortex_err, vortex_panic};
use vortex_mask::{AllOr, Mask, MaskValues};
use vortex_scalar::Scalar;

use crate::arrays::{BoolArray, ConstantArray};
use crate::compute::{fill_null, filter, sum, take};
use crate::patches::Patches;
use crate::{Array, ArrayRef, IntoArray, ToCanonical};

/// Validity information for an array
#[derive(Clone, Debug)]
pub enum Validity {
    /// Items *can't* be null
    NonNullable,
    /// All items are valid
    AllValid,
    /// All items are null
    AllInvalid,
    /// Specified items are null
    Array(ArrayRef),
}

impl Validity {
    /// The [`DType`] of the underlying validity array (if it exists).
    pub const DTYPE: DType = DType::Bool(Nullability::NonNullable);

    pub fn null_count(&self, length: usize) -> VortexResult<usize> {
        match self {
            Self::NonNullable | Self::AllValid => Ok(0),
            Self::AllInvalid => Ok(length),
            Self::Array(a) => {
                let validity_len = a.len();
                if validity_len != length {
                    vortex_bail!(
                        "Validity array length {} doesn't match array length {}",
                        validity_len,
                        length
                    )
                }
                let true_count = sum(a)?
                    .as_primitive()
                    .as_::<usize>()?
                    .ok_or_else(|| vortex_err!("Failed to compute true count"))?;
                Ok(length - true_count)
            }
        }
    }

    /// If Validity is [`Validity::Array`], returns the array, otherwise returns `None`.
    pub fn into_array(self) -> Option<ArrayRef> {
        match self {
            Self::Array(a) => Some(a),
            _ => None,
        }
    }

    /// If Validity is [`Validity::Array`], returns a reference to the array array, otherwise returns `None`.
    pub fn as_array(&self) -> Option<&ArrayRef> {
        match self {
            Self::Array(a) => Some(a),
            _ => None,
        }
    }

    pub fn nullability(&self) -> Nullability {
        match self {
            Self::NonNullable => Nullability::NonNullable,
            _ => Nullability::Nullable,
        }
    }

    /// The union nullability and validity.
    pub fn union_nullability(self, nullability: Nullability) -> Self {
        match nullability {
            Nullability::NonNullable => self,
            Nullability::Nullable => self.into_nullable(),
        }
    }

    pub fn all_valid(&self) -> VortexResult<bool> {
        Ok(match self {
            Validity::NonNullable | Validity::AllValid => true,
            Validity::AllInvalid => false,
            Validity::Array(array) => {
                // TODO(ngates): replace with SUM compute function
                array.to_bool()?.boolean_buffer().count_set_bits() == array.len()
            }
        })
    }

    pub fn all_invalid(&self) -> VortexResult<bool> {
        Ok(match self {
            Validity::NonNullable | Validity::AllValid => false,
            Validity::AllInvalid => true,
            Validity::Array(array) => {
                // TODO(ngates): replace with SUM compute function
                array.to_bool()?.boolean_buffer().count_set_bits() == 0
            }
        })
    }

    /// Returns whether the `index` item is valid.
    #[inline]
    pub fn is_valid(&self, index: usize) -> VortexResult<bool> {
        Ok(match self {
            Self::NonNullable | Self::AllValid => true,
            Self::AllInvalid => false,
            Self::Array(a) => {
                let scalar = a.scalar_at(index)?;
                scalar
                    .as_bool()
                    .value()
                    .vortex_expect("Validity must be non-nullable")
            }
        })
    }

    #[inline]
    pub fn is_null(&self, index: usize) -> VortexResult<bool> {
        Ok(!self.is_valid(index)?)
    }

    pub fn slice(&self, start: usize, stop: usize) -> VortexResult<Self> {
        match self {
            Self::Array(a) => Ok(Self::Array(a.slice(start, stop)?)),
            _ => Ok(self.clone()),
        }
    }

    pub fn take(&self, indices: &dyn Array) -> VortexResult<Self> {
        match self {
            Self::NonNullable => match indices.validity_mask()?.boolean_buffer() {
                AllOr::All => {
                    if indices.dtype().is_nullable() {
                        Ok(Self::AllValid)
                    } else {
                        Ok(Self::NonNullable)
                    }
                }
                AllOr::None => Ok(Self::AllInvalid),
                AllOr::Some(buf) => Ok(Validity::from(buf.clone())),
            },
            Self::AllValid => match indices.validity_mask()?.boolean_buffer() {
                AllOr::All => Ok(Self::AllValid),
                AllOr::None => Ok(Self::AllInvalid),
                AllOr::Some(buf) => Ok(Validity::from(buf.clone())),
            },
            Self::AllInvalid => Ok(Self::AllInvalid),
            Self::Array(is_valid) => {
                let maybe_is_valid = take(is_valid, indices)?;
                // Null indices invalidate that position.
                let is_valid = fill_null(&maybe_is_valid, &Scalar::from(false))?;
                Ok(Self::Array(is_valid))
            }
        }
    }

    /// Keep only the entries for which the mask is true.
    ///
    /// The result has length equal to the number of true values in mask.
    pub fn filter(&self, mask: &Mask) -> VortexResult<Self> {
        // NOTE(ngates): we take the mask as a reference to avoid the caller cloning unnecessarily
        //  if we happen to be NonNullable, AllValid, or AllInvalid.
        match self {
            v @ (Validity::NonNullable | Validity::AllValid | Validity::AllInvalid) => {
                Ok(v.clone())
            }
            Validity::Array(arr) => Ok(Validity::Array(filter(arr, mask)?)),
        }
    }

    /// If this validity is backed by an array, [`optimize`][crate::vtable::operations::OperationsVTable]
    /// the underlying array. Otherwise, returns the original validity unchanged.
    pub fn optimize(&self) -> VortexResult<Self> {
        match self {
            Validity::Array(array) => array.optimize().map(Validity::Array),
            x => Ok(x.clone()),
        }
    }

    /// Set to false any entries for which the mask is true.
    ///
    /// The result is always nullable. The result has the same length as self.
    pub fn mask(&self, mask: &Mask) -> VortexResult<Self> {
        match mask.boolean_buffer() {
            AllOr::All => Ok(Validity::AllInvalid),
            AllOr::None => Ok(self.clone()),
            AllOr::Some(make_invalid) => Ok(match self {
                Validity::NonNullable | Validity::AllValid => {
                    Validity::Array(BoolArray::from(make_invalid.not()).into_array())
                }
                Validity::AllInvalid => Validity::AllInvalid,
                Validity::Array(is_valid) => {
                    let is_valid = is_valid.to_bool()?;
                    let keep_valid = make_invalid.not();
                    Validity::from(is_valid.boolean_buffer().bitand(&keep_valid))
                }
            }),
        }
    }

    pub fn to_mask(&self, length: usize) -> VortexResult<Mask> {
        Ok(match self {
            Self::NonNullable | Self::AllValid => Mask::AllTrue(length),
            Self::AllInvalid => Mask::AllFalse(length),
            Self::Array(is_valid) => {
                assert_eq!(
                    is_valid.len(),
                    length,
                    "Validity::Array length must equal to_logical's argument: {}, {}.",
                    is_valid.len(),
                    length,
                );
                Mask::try_from(&is_valid.to_bool()?)?
            }
        })
    }

    /// Logically & two Validity values of the same length
    pub fn and(self, rhs: Validity) -> VortexResult<Validity> {
        let validity = match (self, rhs) {
            // Should be pretty clear
            (Validity::NonNullable, Validity::NonNullable) => Validity::NonNullable,
            // Any `AllInvalid` makes the output all invalid values
            (Validity::AllInvalid, _) | (_, Validity::AllInvalid) => Validity::AllInvalid,
            // All truthy values on one side, which makes no effect on an `Array` variant
            (Validity::Array(a), Validity::AllValid)
            | (Validity::Array(a), Validity::NonNullable)
            | (Validity::NonNullable, Validity::Array(a))
            | (Validity::AllValid, Validity::Array(a)) => Validity::Array(a),
            // Both sides are all valid
            (Validity::NonNullable, Validity::AllValid)
            | (Validity::AllValid, Validity::NonNullable)
            | (Validity::AllValid, Validity::AllValid) => Validity::AllValid,
            // Here we actually have to do some work
            (Validity::Array(lhs), Validity::Array(rhs)) => {
                let lhs = lhs.to_bool()?;
                let rhs = rhs.to_bool()?;

                let lhs = lhs.boolean_buffer();
                let rhs = rhs.boolean_buffer();

                Validity::from(lhs.bitand(rhs))
            }
        };

        Ok(validity)
    }

    pub fn patch(
        self,
        len: usize,
        indices_offset: usize,
        indices: &dyn Array,
        patches: &Validity,
    ) -> VortexResult<Self> {
        match (&self, patches) {
            (Validity::NonNullable, Validity::NonNullable) => return Ok(Validity::NonNullable),
            (Validity::NonNullable, _) => {
                vortex_bail!("Can't patch a non-nullable validity with nullable validity")
            }
            (_, Validity::NonNullable) => {
                vortex_bail!("Can't patch a nullable validity with non-nullable validity")
            }
            (Validity::AllValid, Validity::AllValid) => return Ok(Validity::AllValid),
            (Validity::AllInvalid, Validity::AllInvalid) => return Ok(Validity::AllInvalid),
            _ => {}
        };

        let own_nullability = if self == Validity::NonNullable {
            Nullability::NonNullable
        } else {
            Nullability::Nullable
        };

        let source = match self {
            Validity::NonNullable => BoolArray::from(BooleanBuffer::new_set(len)),
            Validity::AllValid => BoolArray::from(BooleanBuffer::new_set(len)),
            Validity::AllInvalid => BoolArray::from(BooleanBuffer::new_unset(len)),
            Validity::Array(a) => a.to_bool()?,
        };

        let patch_values = match patches {
            Validity::NonNullable => BoolArray::from(BooleanBuffer::new_set(indices.len())),
            Validity::AllValid => BoolArray::from(BooleanBuffer::new_set(indices.len())),
            Validity::AllInvalid => BoolArray::from(BooleanBuffer::new_unset(indices.len())),
            Validity::Array(a) => a.to_bool()?,
        };

        let patches = Patches::new(
            len,
            indices_offset,
            indices.to_array(),
            patch_values.into_array(),
        );

        Ok(Self::from_array(
            source.patch(&patches)?.into_array(),
            own_nullability,
        ))
    }

    /// Convert into a nullable variant
    pub fn into_nullable(self) -> Validity {
        match self {
            Self::NonNullable => Self::AllValid,
            _ => self,
        }
    }

    /// Convert into a non-nullable variant
    pub fn into_non_nullable(self) -> Option<Validity> {
        match self {
            Self::NonNullable => Some(Self::NonNullable),
            Self::AllValid => Some(Self::NonNullable),
            Self::AllInvalid => None,
            Self::Array(is_valid) => {
                is_valid
                    .statistics()
                    .compute_min::<bool>()
                    .vortex_expect("validity array must support min")
                    .then(|| {
                        // min true => all true
                        Self::NonNullable
                    })
            }
        }
    }

    /// Convert into a variant compatible with the given nullability, if possible.
    pub fn cast_nullability(self, nullability: Nullability) -> VortexResult<Validity> {
        match nullability {
            Nullability::NonNullable => self.into_non_nullable().ok_or_else(|| {
                vortex_err!("Cannot cast array with invalid values to non-nullable type.")
            }),
            Nullability::Nullable => Ok(self.into_nullable()),
        }
    }

    /// Create Validity by copying the given array's validity.
    pub fn copy_from_array(array: &dyn Array) -> VortexResult<Self> {
        Ok(Validity::from_mask(
            array.validity_mask()?,
            array.dtype().nullability(),
        ))
    }

    /// Create Validity from boolean array with given nullability of the array.
    ///
    /// Note: You want to pass the nullability of parent array and not the nullability of the validity array itself
    ///     as that is always nonnullable
    fn from_array(value: ArrayRef, nullability: Nullability) -> Self {
        if !matches!(value.dtype(), DType::Bool(Nullability::NonNullable)) {
            vortex_panic!("Expected a non-nullable boolean array")
        }
        match nullability {
            Nullability::NonNullable => Self::NonNullable,
            Nullability::Nullable => Self::Array(value),
        }
    }

    /// Returns the length of the validity array, if it exists.
    pub fn maybe_len(&self) -> Option<usize> {
        match self {
            Self::NonNullable | Self::AllValid | Self::AllInvalid => None,
            Self::Array(a) => Some(a.len()),
        }
    }

    pub fn uncompressed_size(&self) -> usize {
        if let Validity::Array(a) = self {
            a.len().div_ceil(8)
        } else {
            0
        }
    }

    pub fn is_array(&self) -> bool {
        matches!(self, Validity::Array(_))
    }
}

impl PartialEq for Validity {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::NonNullable, Self::NonNullable) => true,
            (Self::AllValid, Self::AllValid) => true,
            (Self::AllInvalid, Self::AllInvalid) => true,
            (Self::Array(a), Self::Array(b)) => {
                let a = a
                    .to_bool()
                    .vortex_expect("Failed to get Validity Array as BoolArray");
                let b = b
                    .to_bool()
                    .vortex_expect("Failed to get Validity Array as BoolArray");
                a.boolean_buffer() == b.boolean_buffer()
            }
            _ => false,
        }
    }
}

impl From<BooleanBuffer> for Validity {
    fn from(value: BooleanBuffer) -> Self {
        if value.count_set_bits() == value.len() {
            Self::AllValid
        } else if value.count_set_bits() == 0 {
            Self::AllInvalid
        } else {
            Self::Array(BoolArray::from(value).into_array())
        }
    }
}

impl From<NullBuffer> for Validity {
    fn from(value: NullBuffer) -> Self {
        value.into_inner().into()
    }
}

impl FromIterator<Mask> for Validity {
    fn from_iter<T: IntoIterator<Item = Mask>>(iter: T) -> Self {
        Validity::from_mask(iter.into_iter().collect(), Nullability::Nullable)
    }
}

impl FromIterator<bool> for Validity {
    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {
        Validity::from(BooleanBuffer::from_iter(iter))
    }
}

impl From<Nullability> for Validity {
    fn from(value: Nullability) -> Self {
        match value {
            Nullability::NonNullable => Validity::NonNullable,
            Nullability::Nullable => Validity::AllValid,
        }
    }
}

impl Validity {
    pub fn from_mask(mask: Mask, nullability: Nullability) -> Self {
        assert!(
            nullability == Nullability::Nullable || matches!(mask, Mask::AllTrue(_)),
            "NonNullable validity must be AllValid",
        );
        match mask {
            Mask::AllTrue(_) => match nullability {
                Nullability::NonNullable => Validity::NonNullable,
                Nullability::Nullable => Validity::AllValid,
            },
            Mask::AllFalse(_) => Validity::AllInvalid,
            Mask::Values(values) => Validity::Array(values.into_array()),
        }
    }
}

impl IntoArray for Mask {
    fn into_array(self) -> ArrayRef {
        match self {
            Self::AllTrue(len) => ConstantArray::new(true, len).into_array(),
            Self::AllFalse(len) => ConstantArray::new(false, len).into_array(),
            Self::Values(a) => a.into_array(),
        }
    }
}

impl IntoArray for &MaskValues {
    fn into_array(self) -> ArrayRef {
        BoolArray::new(self.boolean_buffer().clone(), Validity::NonNullable).into_array()
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;
    use vortex_buffer::{Buffer, buffer};
    use vortex_dtype::Nullability;
    use vortex_mask::Mask;

    use crate::arrays::{BoolArray, PrimitiveArray};
    use crate::validity::Validity;
    use crate::{ArrayRef, IntoArray};

    #[rstest]
    #[case(Validity::AllValid, 5, &[2, 4], Validity::AllValid, Validity::AllValid)]
    #[case(Validity::AllValid, 5, &[2, 4], Validity::AllInvalid, Validity::Array(BoolArray::from_iter([true, true, false, true, false]).into_array())
    )]
    #[case(Validity::AllValid, 5, &[2, 4], Validity::Array(BoolArray::from_iter([true, false]).into_array()), Validity::Array(BoolArray::from_iter([true, true, true, true, false]).into_array())
    )]
    #[case(Validity::AllInvalid, 5, &[2, 4], Validity::AllValid, Validity::Array(BoolArray::from_iter([false, false, true, false, true]).into_array())
    )]
    #[case(Validity::AllInvalid, 5, &[2, 4], Validity::AllInvalid, Validity::AllInvalid)]
    #[case(Validity::AllInvalid, 5, &[2, 4], Validity::Array(BoolArray::from_iter([true, false]).into_array()), Validity::Array(BoolArray::from_iter([false, false, true, false, false]).into_array())
    )]
    #[case(Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array()), 5, &[2, 4], Validity::AllValid, Validity::Array(BoolArray::from_iter([false, true, true, true, true]).into_array())
    )]
    #[case(Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array()), 5, &[2, 4], Validity::AllInvalid, Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array())
    )]
    #[case(Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array()), 5, &[2, 4], Validity::Array(BoolArray::from_iter([true, false]).into_array()), Validity::Array(BoolArray::from_iter([false, true, true, true, false]).into_array())
    )]
    fn patch_validity(
        #[case] validity: Validity,
        #[case] len: usize,
        #[case] positions: &[u64],
        #[case] patches: Validity,
        #[case] expected: Validity,
    ) {
        let indices =
            PrimitiveArray::new(Buffer::copy_from(positions), Validity::NonNullable).into_array();
        assert_eq!(
            validity.patch(len, 0, &indices, &patches).unwrap(),
            expected
        );
    }

    #[test]
    #[should_panic]
    fn out_of_bounds_patch() {
        Validity::NonNullable
            .patch(2, 0, &buffer![4].into_array(), &Validity::AllInvalid)
            .unwrap();
    }

    #[test]
    #[should_panic]
    fn into_validity_nullable() {
        Validity::from_mask(Mask::AllFalse(10), Nullability::NonNullable);
    }

    #[test]
    #[should_panic]
    fn into_validity_nullable_array() {
        Validity::from_mask(Mask::from_iter(vec![true, false]), Nullability::NonNullable);
    }

    #[rstest]
    #[case(Validity::AllValid, PrimitiveArray::new(buffer![0, 1], Validity::from_iter(vec![true, false])).into_array(), Validity::from_iter(vec![true, false]))]
    #[case(Validity::AllValid, buffer![0, 1].into_array(), Validity::AllValid)]
    #[case(Validity::AllValid, PrimitiveArray::new(buffer![0, 1], Validity::AllInvalid).into_array(), Validity::AllInvalid)]
    #[case(Validity::NonNullable, PrimitiveArray::new(buffer![0, 1], Validity::from_iter(vec![true, false])).into_array(), Validity::from_iter(vec![true, false]))]
    #[case(Validity::NonNullable, buffer![0, 1].into_array(), Validity::NonNullable)]
    #[case(Validity::NonNullable, PrimitiveArray::new(buffer![0, 1], Validity::AllInvalid).into_array(), Validity::AllInvalid)]
    fn validity_take(
        #[case] validity: Validity,
        #[case] indices: ArrayRef,
        #[case] expected: Validity,
    ) {
        assert_eq!(validity.take(&indices).unwrap(), expected);
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	//! Array validity and nullability behavior, used by arrays and compute functions.
5
6	use std::fmt::Debug;
7	use std::ops::{BitAnd, Not};
8
9	use arrow_buffer::{BooleanBuffer, NullBuffer};
10	use vortex_dtype::{DType, Nullability};
11	use vortex_error::{VortexExpect as _, VortexResult, vortex_bail, vortex_err, vortex_panic};
12	use vortex_mask::{AllOr, Mask, MaskValues};
13	use vortex_scalar::Scalar;
14
15	use crate::arrays::{BoolArray, ConstantArray};
16	use crate::compute::{fill_null, filter, sum, take};
17	use crate::patches::Patches;
18	use crate::{Array, ArrayRef, IntoArray, ToCanonical};
19
20	/// Validity information for an array
21	#[derive(Clone, Debug)]
22	pub enum Validity {
23	/// Items can't be null
24	NonNullable,
25	/// All items are valid
26	AllValid,
27	/// All items are null
28	AllInvalid,
29	/// Specified items are null
30	Array(ArrayRef),
31	}
32
33	impl Validity {
34	/// The [`DType`] of the underlying validity array (if it exists).
35	pub const DTYPE: DType = DType::Bool(Nullability::NonNullable);
36
37	pub fn null_count(&self, length: usize) -> VortexResult<usize> {	×
38	match self {	×
39	Self::NonNullable \| Self::AllValid => Ok(0),	×
40	Self::AllInvalid => Ok(length),	×
41	Self::Array(a) => {	×
42	let validity_len = a.len();	×
43	if validity_len != length {	×
44	vortex_bail!(	×
45	"Validity array length {} doesn't match array length {}",	×
46	validity_len,
47	length
48	)
49	}	×
50	let true_count = sum(a)?	×
51	.as_primitive()	×
52	.as_::<usize>()?	×
53	.ok_or_else(\|\| vortex_err!("Failed to compute true count"))?;	×
54	Ok(length - true_count)	×
55	}
56	}
57	}	×
58
59	/// If Validity is [`Validity::Array`], returns the array, otherwise returns `None`.
60	pub fn into_array(self) -> Option<ArrayRef> {	×
61	match self {	×
62	Self::Array(a) => Some(a),	×
63	_ => None,	×
64	}
65	}	×
66
67	/// If Validity is [`Validity::Array`], returns a reference to the array array, otherwise returns `None`.
68	pub fn as_array(&self) -> Option<&ArrayRef> {	×
69	match self {	×
70	Self::Array(a) => Some(a),	×
71	_ => None,	×
72	}
73	}	×
74
75	pub fn nullability(&self) -> Nullability {	312,524✔
76	match self {	312,524✔
77	Self::NonNullable => Nullability::NonNullable,	250,461✔
78	_ => Nullability::Nullable,	62,063✔
79	}
80	}	312,524✔
81
82	/// The union nullability and validity.
83	pub fn union_nullability(self, nullability: Nullability) -> Self {	405✔
84	match nullability {	405✔
85	Nullability::NonNullable => self,	6✔
86	Nullability::Nullable => self.into_nullable(),	399✔
87	}
88	}	405✔
89
90	pub fn all_valid(&self) -> VortexResult<bool> {	3,150,493✔
91	Ok(match self {	3,150,493✔
92	Validity::NonNullable \| Validity::AllValid => true,	3,148,236✔
93	Validity::AllInvalid => false,	36✔
94	Validity::Array(array) => {	2,221✔
95	// TODO(ngates): replace with SUM compute function
96	array.to_bool()?.boolean_buffer().count_set_bits() == array.len()	2,221✔
97	}
98	})
99	}	3,150,493✔
100
101	pub fn all_invalid(&self) -> VortexResult<bool> {	152,043✔
102	Ok(match self {	152,043✔
103	Validity::NonNullable \| Validity::AllValid => false,	145,629✔
104	Validity::AllInvalid => true,	401✔
105	Validity::Array(array) => {	6,013✔
106	// TODO(ngates): replace with SUM compute function
107	array.to_bool()?.boolean_buffer().count_set_bits() == 0	6,013✔
108	}
109	})
110	}	152,043✔
111
112	/// Returns whether the `index` item is valid.
113	#[inline]
114	pub fn is_valid(&self, index: usize) -> VortexResult<bool> {	3,559,697✔
115	Ok(match self {	3,559,697✔
116	Self::NonNullable \| Self::AllValid => true,	3,546,049✔
117	Self::AllInvalid => false,	171✔
118	Self::Array(a) => {	13,477✔
119	let scalar = a.scalar_at(index)?;	13,477✔
120	scalar	13,477✔
121	.as_bool()	13,477✔
122	.value()	13,477✔
123	.vortex_expect("Validity must be non-nullable")	13,477✔
124	}
125	})
126	}	3,559,697✔
127
128	#[inline]
129	pub fn is_null(&self, index: usize) -> VortexResult<bool> {	216✔
130	Ok(!self.is_valid(index)?)	216✔
131	}	216✔
132
133	pub fn slice(&self, start: usize, stop: usize) -> VortexResult<Self> {	49,920✔
134	match self {	49,920✔
135	Self::Array(a) => Ok(Self::Array(a.slice(start, stop)?)),	1,298✔
136	_ => Ok(self.clone()),	48,622✔
137	}
138	}	49,920✔
139
140	pub fn take(&self, indices: &dyn Array) -> VortexResult<Self> {	3,704✔
141	match self {	3,704✔
142	Self::NonNullable => match indices.validity_mask()?.boolean_buffer() {	2,433✔
143	AllOr::All => {
144	if indices.dtype().is_nullable() {	2,135✔
145	Ok(Self::AllValid)	145✔
146	} else {
147	Ok(Self::NonNullable)	1,990✔
148	}
149	}
150	AllOr::None => Ok(Self::AllInvalid),	1✔
151	AllOr::Some(buf) => Ok(Validity::from(buf.clone())),	297✔
152	},
153	Self::AllValid => match indices.validity_mask()?.boolean_buffer() {	1,085✔
154	AllOr::All => Ok(Self::AllValid),	219✔
155	AllOr::None => Ok(Self::AllInvalid),	1✔
156	AllOr::Some(buf) => Ok(Validity::from(buf.clone())),	865✔
157	},
158	Self::AllInvalid => Ok(Self::AllInvalid),	×
159	Self::Array(is_valid) => {	186✔
160	let maybe_is_valid = take(is_valid, indices)?;	186✔
161	// Null indices invalidate that position.
162	let is_valid = fill_null(&maybe_is_valid, &Scalar::from(false))?;	186✔
163	Ok(Self::Array(is_valid))	186✔
164	}
165	}
166	}	3,704✔
167
168	/// Keep only the entries for which the mask is true.
169	///
170	/// The result has length equal to the number of true values in mask.
171	pub fn filter(&self, mask: &Mask) -> VortexResult<Self> {	5,249✔
172	// NOTE(ngates): we take the mask as a reference to avoid the caller cloning unnecessarily
173	// if we happen to be NonNullable, AllValid, or AllInvalid.
174	match self {	5,249✔
175	v @ (Validity::NonNullable \| Validity::AllValid \| Validity::AllInvalid) => {	4,815✔
176	Ok(v.clone())	4,815✔
177	}
178	Validity::Array(arr) => Ok(Validity::Array(filter(arr, mask)?)),	434✔
179	}
180	}	5,249✔
181
182	/// If this validity is backed by an array, [`optimize`][crate::vtable::operations::OperationsVTable]
183	/// the underlying array. Otherwise, returns the original validity unchanged.
NEW 184	pub fn optimize(&self) -> VortexResult<Self> {	×
NEW 185	match self {	×
NEW 186	Validity::Array(array) => array.optimize().map(Validity::Array),	×
NEW 187	x => Ok(x.clone()),	×
188	}
NEW 189	}	×
190
191	/// Set to false any entries for which the mask is true.
192	///
193	/// The result is always nullable. The result has the same length as self.
194	pub fn mask(&self, mask: &Mask) -> VortexResult<Self> {	592✔
195	match mask.boolean_buffer() {	592✔
196	AllOr::All => Ok(Validity::AllInvalid),	×
197	AllOr::None => Ok(self.clone()),	×
198	AllOr::Some(make_invalid) => Ok(match self {	592✔
199	Validity::NonNullable \| Validity::AllValid => {
200	Validity::Array(BoolArray::from(make_invalid.not()).into_array())	94✔
201	}
202	Validity::AllInvalid => Validity::AllInvalid,	4✔
203	Validity::Array(is_valid) => {	494✔
204	let is_valid = is_valid.to_bool()?;	494✔
205	let keep_valid = make_invalid.not();	494✔
206	Validity::from(is_valid.boolean_buffer().bitand(&keep_valid))	494✔
207	}
208	}),
209	}
210	}	592✔
211
212	pub fn to_mask(&self, length: usize) -> VortexResult<Mask> {	311,971✔
213	Ok(match self {	311,971✔
214	Self::NonNullable \| Self::AllValid => Mask::AllTrue(length),	290,729✔
215	Self::AllInvalid => Mask::AllFalse(length),	228✔
216	Self::Array(is_valid) => {	21,014✔
217	assert_eq!(	21,014✔
218	is_valid.len(),	21,014✔
219	length,
220	"Validity::Array length must equal to_logical's argument: {}, {}.",	×
221	is_valid.len(),	×
222	length,
223	);
224	Mask::try_from(&is_valid.to_bool()?)?	21,014✔
225	}
226	})
227	}	311,971✔
228
229	/// Logically & two Validity values of the same length
230	pub fn and(self, rhs: Validity) -> VortexResult<Validity> {	1✔
231	let validity = match (self, rhs) {	1✔
232	// Should be pretty clear
233	(Validity::NonNullable, Validity::NonNullable) => Validity::NonNullable,	1✔
234	// Any `AllInvalid` makes the output all invalid values
235	(Validity::AllInvalid, _) \| (_, Validity::AllInvalid) => Validity::AllInvalid,	×
236	// All truthy values on one side, which makes no effect on an `Array` variant
237	(Validity::Array(a), Validity::AllValid)	×
238	\| (Validity::Array(a), Validity::NonNullable)	×
239	\| (Validity::NonNullable, Validity::Array(a))	×
240	\| (Validity::AllValid, Validity::Array(a)) => Validity::Array(a),	×
241	// Both sides are all valid
242	(Validity::NonNullable, Validity::AllValid)
243	\| (Validity::AllValid, Validity::NonNullable)
244	\| (Validity::AllValid, Validity::AllValid) => Validity::AllValid,	×
245	// Here we actually have to do some work
246	(Validity::Array(lhs), Validity::Array(rhs)) => {	×
247	let lhs = lhs.to_bool()?;	×
248	let rhs = rhs.to_bool()?;	×
249
250	let lhs = lhs.boolean_buffer();	×
251	let rhs = rhs.boolean_buffer();	×
252
253	Validity::from(lhs.bitand(rhs))	×
254	}
255	};
256
257	Ok(validity)	1✔
258	}	1✔
259
260	pub fn patch(	2,081✔
261	self,	2,081✔
262	len: usize,	2,081✔
263	indices_offset: usize,	2,081✔
264	indices: &dyn Array,	2,081✔
265	patches: &Validity,	2,081✔
266	) -> VortexResult<Self> {	2,081✔
267	match (&self, patches) {	2,081✔
268	(Validity::NonNullable, Validity::NonNullable) => return Ok(Validity::NonNullable),	1,343✔
269	(Validity::NonNullable, _) => {
270	vortex_bail!("Can't patch a non-nullable validity with nullable validity")	1✔
271	}
272	(_, Validity::NonNullable) => {
273	vortex_bail!("Can't patch a nullable validity with non-nullable validity")	×
274	}
275	(Validity::AllValid, Validity::AllValid) => return Ok(Validity::AllValid),	1✔
276	(Validity::AllInvalid, Validity::AllInvalid) => return Ok(Validity::AllInvalid),	1✔
277	_ => {}	735✔
278	};
279
280	let own_nullability = if self == Validity::NonNullable {	735✔
281	Nullability::NonNullable	×
282	} else {
283	Nullability::Nullable	735✔
284	};
285
286	let source = match self {	735✔
287	Validity::NonNullable => BoolArray::from(BooleanBuffer::new_set(len)),	×
288	Validity::AllValid => BoolArray::from(BooleanBuffer::new_set(len)),	362✔
289	Validity::AllInvalid => BoolArray::from(BooleanBuffer::new_unset(len)),	334✔
290	Validity::Array(a) => a.to_bool()?,	39✔
291	};
292
293	let patch_values = match patches {	735✔
294	Validity::NonNullable => BoolArray::from(BooleanBuffer::new_set(indices.len())),	×
295	Validity::AllValid => BoolArray::from(BooleanBuffer::new_set(indices.len())),	154✔
296	Validity::AllInvalid => BoolArray::from(BooleanBuffer::new_unset(indices.len())),	2✔
297	Validity::Array(a) => a.to_bool()?,	579✔
298	};
299
300	let patches = Patches::new(	735✔
301	len,	735✔
302	indices_offset,	735✔
303	indices.to_array(),	735✔
304	patch_values.into_array(),	735✔
305	);
306
307	Ok(Self::from_array(	735✔
308	source.patch(&patches)?.into_array(),	735✔
309	own_nullability,	735✔
310	))
311	}	2,081✔
312
313	/// Convert into a nullable variant
314	pub fn into_nullable(self) -> Validity {	1,831✔
315	match self {	1,831✔
316	Self::NonNullable => Self::AllValid,	1,648✔
317	_ => self,	183✔
318	}
319	}	1,831✔
320
321	/// Convert into a non-nullable variant
322	pub fn into_non_nullable(self) -> Option<Validity> {	51✔
323	match self {	51✔
324	Self::NonNullable => Some(Self::NonNullable),	1✔
325	Self::AllValid => Some(Self::NonNullable),	44✔
326	Self::AllInvalid => None,	×
327	Self::Array(is_valid) => {	6✔
328	is_valid	6✔
329	.statistics()	6✔
330	.compute_min::<bool>()	6✔
331	.vortex_expect("validity array must support min")	6✔
332	.then(\|\| {	6✔
333	// min true => all true
334	Self::NonNullable	×
335	})	×
336	}
337	}
338	}	51✔
339
340	/// Convert into a variant compatible with the given nullability, if possible.
341	pub fn cast_nullability(self, nullability: Nullability) -> VortexResult<Validity> {	70✔
342	match nullability {	70✔
343	Nullability::NonNullable => self.into_non_nullable().ok_or_else(\|\| {	51✔
344	vortex_err!("Cannot cast array with invalid values to non-nullable type.")	6✔
345	}),	6✔
346	Nullability::Nullable => Ok(self.into_nullable()),	19✔
347	}
348	}	70✔
349
350	/// Create Validity by copying the given array's validity.
351	pub fn copy_from_array(array: &dyn Array) -> VortexResult<Self> {	2,112✔
352	Ok(Validity::from_mask(	2,112✔
353	array.validity_mask()?,	2,112✔
354	array.dtype().nullability(),	2,112✔
355	))
356	}	2,112✔
357
358	/// Create Validity from boolean array with given nullability of the array.
359	///
360	/// Note: You want to pass the nullability of parent array and not the nullability of the validity array itself
361	/// as that is always nonnullable
362	fn from_array(value: ArrayRef, nullability: Nullability) -> Self {	735✔
363	if !matches!(value.dtype(), DType::Bool(Nullability::NonNullable)) {	735✔
364	vortex_panic!("Expected a non-nullable boolean array")	×
365	}	735✔
366	match nullability {	735✔
367	Nullability::NonNullable => Self::NonNullable,	×
368	Nullability::Nullable => Self::Array(value),	735✔
369	}
370	}	735✔
371
372	/// Returns the length of the validity array, if it exists.
373	pub fn maybe_len(&self) -> Option<usize> {	548,349✔
374	match self {	548,349✔
375	Self::NonNullable \| Self::AllValid \| Self::AllInvalid => None,	523,766✔
376	Self::Array(a) => Some(a.len()),	24,583✔
377	}
378	}	548,349✔
379
380	pub fn uncompressed_size(&self) -> usize {	×
381	if let Validity::Array(a) = self {	×
382	a.len().div_ceil(8)	×
383	} else {
384	0	×
385	}
386	}	×
387
388	pub fn is_array(&self) -> bool {	×
389	matches!(self, Validity::Array(_))	×
390	}	×
391	}
392
393	impl PartialEq for Validity {
394	fn eq(&self, other: &Self) -> bool {	36,412✔
395	match (self, other) {	36,412✔
396	(Self::NonNullable, Self::NonNullable) => true,	20,473✔
397	(Self::AllValid, Self::AllValid) => true,	116✔
398	(Self::AllInvalid, Self::AllInvalid) => true,	114✔
399	(Self::Array(a), Self::Array(b)) => {	487✔
400	let a = a	487✔
401	.to_bool()	487✔
402	.vortex_expect("Failed to get Validity Array as BoolArray");	487✔
403	let b = b	487✔
404	.to_bool()	487✔
405	.vortex_expect("Failed to get Validity Array as BoolArray");	487✔
406	a.boolean_buffer() == b.boolean_buffer()	487✔
407	}
408	_ => false,	15,222✔
409	}
410	}	36,412✔
411	}
412
413	impl From<BooleanBuffer> for Validity {
414	fn from(value: BooleanBuffer) -> Self {	11,341✔
415	if value.count_set_bits() == value.len() {	11,341✔
416	Self::AllValid	256✔
417	} else if value.count_set_bits() == 0 {	11,085✔
418	Self::AllInvalid	478✔
419	} else {
420	Self::Array(BoolArray::from(value).into_array())	10,607✔
421	}
422	}	11,341✔
423	}
424
425	impl From<NullBuffer> for Validity {
426	fn from(value: NullBuffer) -> Self {	3,272✔
427	value.into_inner().into()	3,272✔
428	}	3,272✔
429	}
430
431	impl FromIterator<Mask> for Validity {
432	fn from_iter<T: IntoIterator<Item = Mask>>(iter: T) -> Self {	×
433	Validity::from_mask(iter.into_iter().collect(), Nullability::Nullable)	×
434	}	×
435	}
436
437	impl FromIterator<bool> for Validity {
438	fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {	188✔
439	Validity::from(BooleanBuffer::from_iter(iter))	188✔
440	}	188✔
441	}
442
443	impl From<Nullability> for Validity {
444	fn from(value: Nullability) -> Self {	19,003✔
445	match value {	19,003✔
446	Nullability::NonNullable => Validity::NonNullable,	14,943✔
447	Nullability::Nullable => Validity::AllValid,	4,060✔
448	}
449	}	19,003✔
450	}
451
452	impl Validity {
453	pub fn from_mask(mask: Mask, nullability: Nullability) -> Self {	2,479✔
454	assert!(	2,479✔
455	nullability == Nullability::Nullable \|\| matches!(mask, Mask::AllTrue(_)),	2,479✔
456	"NonNullable validity must be AllValid",	2✔
457	);
458	match mask {	2,477✔
459	Mask::AllTrue(_) => match nullability {	1,036✔
460	Nullability::NonNullable => Validity::NonNullable,	927✔
461	Nullability::Nullable => Validity::AllValid,	109✔
462	},
463	Mask::AllFalse(_) => Validity::AllInvalid,	72✔
464	Mask::Values(values) => Validity::Array(values.into_array()),	1,369✔
465	}
466	}	2,477✔
467	}
468
469	impl IntoArray for Mask {
470	fn into_array(self) -> ArrayRef {	36✔
471	match self {	36✔
472	Self::AllTrue(len) => ConstantArray::new(true, len).into_array(),	×
473	Self::AllFalse(len) => ConstantArray::new(false, len).into_array(),	×
474	Self::Values(a) => a.into_array(),	36✔
475	}
476	}	36✔
477	}
478
479	impl IntoArray for &MaskValues {
480	fn into_array(self) -> ArrayRef {	1,405✔
481	BoolArray::new(self.boolean_buffer().clone(), Validity::NonNullable).into_array()	1,405✔
482	}	1,405✔
483	}
484
485	#[cfg(test)]
486	mod tests {
487	use rstest::rstest;
488	use vortex_buffer::{Buffer, buffer};
489	use vortex_dtype::Nullability;
490	use vortex_mask::Mask;
491
492	use crate::arrays::{BoolArray, PrimitiveArray};
493	use crate::validity::Validity;
494	use crate::{ArrayRef, IntoArray};
495
496	#[rstest]
497	#[case(Validity::AllValid, 5, &[2, 4], Validity::AllValid, Validity::AllValid)]
498	#[case(Validity::AllValid, 5, &[2, 4], Validity::AllInvalid, Validity::Array(BoolArray::from_iter([true, true, false, true, false]).into_array())
499	)]
500	#[case(Validity::AllValid, 5, &[2, 4], Validity::Array(BoolArray::from_iter([true, false]).into_array()), Validity::Array(BoolArray::from_iter([true, true, true, true, false]).into_array())
501	)]
502	#[case(Validity::AllInvalid, 5, &[2, 4], Validity::AllValid, Validity::Array(BoolArray::from_iter([false, false, true, false, true]).into_array())
503	)]
504	#[case(Validity::AllInvalid, 5, &[2, 4], Validity::AllInvalid, Validity::AllInvalid)]
505	#[case(Validity::AllInvalid, 5, &[2, 4], Validity::Array(BoolArray::from_iter([true, false]).into_array()), Validity::Array(BoolArray::from_iter([false, false, true, false, false]).into_array())
506	)]
507	#[case(Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array()), 5, &[2, 4], Validity::AllValid, Validity::Array(BoolArray::from_iter([false, true, true, true, true]).into_array())
508	)]
509	#[case(Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array()), 5, &[2, 4], Validity::AllInvalid, Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array())
510	)]
511	#[case(Validity::Array(BoolArray::from_iter([false, true, false, true, false]).into_array()), 5, &[2, 4], Validity::Array(BoolArray::from_iter([true, false]).into_array()), Validity::Array(BoolArray::from_iter([false, true, true, true, false]).into_array())
512	)]
513	fn patch_validity(
514	#[case] validity: Validity,
515	#[case] len: usize,
516	#[case] positions: &[u64],
517	#[case] patches: Validity,
518	#[case] expected: Validity,
519	) {
520	let indices =
521	PrimitiveArray::new(Buffer::copy_from(positions), Validity::NonNullable).into_array();
522	assert_eq!(
523	validity.patch(len, 0, &indices, &patches).unwrap(),
524	expected
525	);
526	}
527
528	#[test]
529	#[should_panic]
530	fn out_of_bounds_patch() {	1✔
531	Validity::NonNullable	1✔
532	.patch(2, 0, &buffer![4].into_array(), &Validity::AllInvalid)	1✔
533	.unwrap();	1✔
534	}	1✔
535
536	#[test]
537	#[should_panic]
538	fn into_validity_nullable() {	1✔
539	Validity::from_mask(Mask::AllFalse(10), Nullability::NonNullable);	1✔
540	}	1✔
541
542	#[test]
543	#[should_panic]
544	fn into_validity_nullable_array() {	1✔
545	Validity::from_mask(Mask::from_iter(vec![true, false]), Nullability::NonNullable);	1✔
546	}	1✔
547
548	#[rstest]
549	#[case(Validity::AllValid, PrimitiveArray::new(buffer![0, 1], Validity::from_iter(vec![true, false])).into_array(), Validity::from_iter(vec![true, false]))]
550	#[case(Validity::AllValid, buffer![0, 1].into_array(), Validity::AllValid)]
551	#[case(Validity::AllValid, PrimitiveArray::new(buffer![0, 1], Validity::AllInvalid).into_array(), Validity::AllInvalid)]
552	#[case(Validity::NonNullable, PrimitiveArray::new(buffer![0, 1], Validity::from_iter(vec![true, false])).into_array(), Validity::from_iter(vec![true, false]))]
553	#[case(Validity::NonNullable, buffer![0, 1].into_array(), Validity::NonNullable)]
554	#[case(Validity::NonNullable, PrimitiveArray::new(buffer![0, 1], Validity::AllInvalid).into_array(), Validity::AllInvalid)]
555	fn validity_take(
556	#[case] validity: Validity,
557	#[case] indices: ArrayRef,
558	#[case] expected: Validity,
559	) {
560	assert_eq!(validity.take(&indices).unwrap(), expected);
561	}
562	}

vortex-data / vortex / 16355473372

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous