vortex_array/arrays/chunked/vtable/
validity.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexResult;
6
7use crate::IntoArray;
8use crate::array::ArrayView;
9use crate::array::ValidityVTable;
10use crate::arrays::Chunked;
11use crate::arrays::ChunkedArray;
12use crate::arrays::chunked::ChunkedArrayExt;
13use crate::dtype::DType;
14use crate::dtype::Nullability;
15use crate::validity::Validity;
16
17impl ValidityVTable<Chunked> for Chunked {
18    fn validity(array: ArrayView<'_, Chunked>) -> VortexResult<Validity> {
19        let validities: Vec<Validity> =
20            array.chunks().iter().map(|c| c.validity()).try_collect()?;
21
22        match validities.first() {
23            // If there are no chunks, return the array's dtype nullability
24            None => return Ok(array.dtype().nullability().into()),
25            // If all chunks have the same non-array validity, return that validity directly
26            // We skip Validity::Array since equality is very expensive.
27            Some(first) if !matches!(first, Validity::Array(_)) => {
28                let target = std::mem::discriminant(first);
29                if validities
30                    .iter()
31                    .all(|v| std::mem::discriminant(v) == target)
32                {
33                    return Ok(first.clone());
34                }
35            }
36            _ => {
37                // Array validity or mixed validities, proceed to build the validity array
38            }
39        }
40
41        Ok(Validity::Array(
42            unsafe {
43                ChunkedArray::new_unchecked(
44                    validities
45                        .into_iter()
46                        .zip(array.iter_chunks())
47                        .map(|(v, chunk)| v.to_array(chunk.len()))
48                        .collect(),
49                    DType::Bool(Nullability::NonNullable),
50                )
51            }
52            .into_array(),
53        ))
54    }
55}
vortex_array/arrays/chunked/vtable/validity.rs

vortex_array/arrays/chunked/vtable/
validity.rs
