Skip to main content

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