Arrays must return Validity (#5785)

gatesn · web-flow · commit 31026ce58862 · 2025-12-18T13:01:30.000Z
Signed-off-by: Nicholas Gates &lt;nick@nickgates.com&gt;
diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs
@@ -432,7 +432,7 @@ impl ALPArray {
 }
 
 impl ValidityChild<ALPVTable> for ALPVTable {
-    fn validity_child(array: &ALPArray) -> &dyn Array {
+    fn validity_child(array: &ALPArray) -> &ArrayRef {
         array.encoded()
     }
 }
diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs
@@ -368,7 +368,7 @@ impl ALPRDArray {
 }
 
 impl ValidityChild<ALPRDVTable> for ALPRDVTable {
-    fn validity_child(array: &ALPRDArray) -> &dyn Array {
+    fn validity_child(array: &ALPRDArray) -> &ArrayRef {
         array.left_parts()
     }
 }
diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs
@@ -288,7 +288,7 @@ impl BaseArrayVTable<DateTimePartsVTable> for DateTimePartsVTable {
 }
 
 impl ValidityChild<DateTimePartsVTable> for DateTimePartsVTable {
-    fn validity_child(array: &DateTimePartsArray) -> &dyn Array {
+    fn validity_child(array: &DateTimePartsArray) -> &ArrayRef {
         array.days()
     }
 }
diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs
@@ -269,9 +269,9 @@ impl OperationsVTable<DecimalBytePartsVTable> for DecimalBytePartsVTable {
 }
 
 impl ValidityChild<DecimalBytePartsVTable> for DecimalBytePartsVTable {
-    fn validity_child(array: &DecimalBytePartsArray) -> &dyn Array {
+    fn validity_child(array: &DecimalBytePartsArray) -> &ArrayRef {
         // validity stored in 0th child
-        array.msp.as_ref()
+        &array.msp
     }
 }
 
diff --git a/encodings/fastlanes/src/for/vtable/validity.rs b/encodings/fastlanes/src/for/vtable/validity.rs
@@ -1,14 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_array::Array;
+use vortex_array::ArrayRef;
 use vortex_array::vtable::ValidityChild;
 
 use super::FoRVTable;
 use crate::FoRArray;
 
 impl ValidityChild<FoRVTable> for FoRVTable {
-    fn validity_child(array: &FoRArray) -> &dyn Array {
-        array.encoded().as_ref()
+    fn validity_child(array: &FoRArray) -> &ArrayRef {
+        array.encoded()
     }
 }
diff --git a/encodings/fastlanes/src/rle/vtable/validity.rs b/encodings/fastlanes/src/rle/vtable/validity.rs
@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_array::Array;
 use vortex_array::ArrayRef;
 use vortex_array::vtable::ValidityChild;
 use vortex_array::vtable::ValidityChildSliceHelper;
@@ -10,8 +9,8 @@ use super::RLEVTable;
 use crate::RLEArray;
 
 impl ValidityChild<RLEVTable> for RLEVTable {
-    fn validity_child(array: &RLEArray) -> &dyn Array {
-        array.indices().as_ref()
+    fn validity_child(array: &RLEArray) -> &ArrayRef {
+        array.indices()
     }
 }
 
diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs
@@ -198,6 +198,8 @@ pub struct FSSTArray {
     symbols: Buffer<Symbol>,
     symbol_lengths: Buffer<u8>,
     codes: VarBinArray,
+    /// NOTE(ngates): this === codes, but is stored as an ArrayRef so we can return &ArrayRef!
+    codes_array: ArrayRef,
     /// Lengths of the original values before compression, can be compressed.
     uncompressed_lengths: ArrayRef,
     stats_set: ArrayStats,
@@ -283,12 +285,14 @@ impl FSSTArray {
             Compressor::rebuild_from(symbols2.as_slice(), symbol_lengths2.as_slice())
         })
             as Box<dyn Fn() -> Compressor + Send>));
+        let codes_array = codes.to_array();
 
         Self {
             dtype,
             symbols,
             symbol_lengths,
             codes,
+            codes_array,
             uncompressed_lengths,
             stats_set: Default::default(),
             compressor,
@@ -377,8 +381,8 @@ impl BaseArrayVTable<FSSTVTable> for FSSTVTable {
 }
 
 impl ValidityChild<FSSTVTable> for FSSTVTable {
-    fn validity_child(array: &FSSTArray) -> &dyn Array {
-        array.codes().as_ref()
+    fn validity_child(array: &FSSTArray) -> &ArrayRef {
+        &array.codes_array
     }
 }
 
diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs
@@ -21,6 +21,7 @@ use vortex_array::search_sorted::SearchSortedSide;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::ArrayStats;
 use vortex_array::stats::StatsSetRef;
+use vortex_array::validity::Validity;
 use vortex_array::vtable;
 use vortex_array::vtable::ArrayId;
 use vortex_array::vtable::ArrayVTable;
@@ -427,6 +428,22 @@ impl ValidityVTable<RunEndVTable> for RunEndVTable {
         array.values().all_invalid()
     }
 
+    fn validity(array: &RunEndArray) -> VortexResult<Validity> {
+        Ok(match array.values().validity()? {
+            Validity::NonNullable | Validity::AllValid => Validity::AllValid,
+            Validity::AllInvalid => Validity::AllInvalid,
+            Validity::Array(values_validity) => Validity::Array(unsafe {
+                RunEndArray::new_unchecked(
+                    array.ends().clone(),
+                    values_validity,
+                    array.offset(),
+                    array.len(),
+                )
+                .into_array()
+            }),
+        })
+    }
+
     fn validity_mask(array: &RunEndArray) -> Mask {
         match array.values().validity_mask() {
             Mask::AllTrue(_) => Mask::AllTrue(array.len()),
diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs
@@ -19,6 +19,7 @@ use vortex_array::arrays::PrimitiveArray;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::ArrayStats;
 use vortex_array::stats::StatsSetRef;
+use vortex_array::validity::Validity;
 use vortex_array::vtable;
 use vortex_array::vtable::ArrayId;
 use vortex_array::vtable::ArrayVTable;
@@ -413,6 +414,10 @@ impl ValidityVTable<SequenceVTable> for SequenceVTable {
         false
     }
 
+    fn validity(_array: &SequenceArray) -> VortexResult<Validity> {
+        Ok(Validity::AllValid)
+    }
+
     fn validity_mask(array: &SequenceArray) -> Mask {
         Mask::AllTrue(array.len())
     }
diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs
@@ -29,6 +29,7 @@ use vortex_array::patches::PatchesMetadata;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::ArrayStats;
 use vortex_array::stats::StatsSetRef;
+use vortex_array::validity::Validity;
 use vortex_array::vtable;
 use vortex_array::vtable::ArrayId;
 use vortex_array::vtable::ArrayVTable;
@@ -410,6 +411,28 @@ impl ValidityVTable<SparseVTable> for SparseVTable {
         array.patches().values().all_invalid()
     }
 
+    fn validity(array: &SparseArray) -> VortexResult<Validity> {
+        let patches = unsafe {
+            Patches::new_unchecked(
+                array.patches.array_len(),
+                array.patches.offset(),
+                array.patches.indices().clone(),
+                array
+                    .patches
+                    .values()
+                    .validity()?
+                    .to_array(array.patches.values().len()),
+                array.patches.chunk_offsets().clone(),
+                array.patches.offset_within_chunk(),
+            )
+        };
+
+        Ok(Validity::Array(
+            unsafe { SparseArray::new_unchecked(patches, array.fill_value.is_valid().into()) }
+                .into_array(),
+        ))
+    }
+
     fn validity_mask(array: &SparseArray) -> Mask {
         let fill_is_valid = array.fill_scalar().is_valid();
         let values_validity = array.patches().values().validity_mask();
diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs
@@ -206,7 +206,7 @@ impl OperationsVTable<ZigZagVTable> for ZigZagVTable {
 }
 
 impl ValidityChild<ZigZagVTable> for ZigZagVTable {
-    fn validity_child(array: &ZigZagArray) -> &dyn Array {
+    fn validity_child(array: &ZigZagArray) -> &ArrayRef {
         array.encoded()
     }
 }
diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs
@@ -55,6 +55,7 @@ use crate::expr::stats::StatsProviderExt;
 use crate::hash;
 use crate::optimizer::ArrayOptimizer;
 use crate::stats::StatsSetRef;
+use crate::validity::Validity;
 use crate::vtable::ArrayId;
 use crate::vtable::ArrayVTable;
 use crate::vtable::BaseArrayVTable;
@@ -160,6 +161,9 @@ pub trait Array:
     /// Returns the number of invalid elements in the array.
     fn invalid_count(&self) -> usize;
 
+    /// Returns the [`Validity`] of the array.
+    fn validity(&self) -> VortexResult<Validity>;
+
     /// Returns the canonical validity mask for the array.
     fn validity_mask(&self) -> Mask;
 
@@ -281,6 +285,11 @@ impl Array for Arc<dyn Array> {
         self.as_ref().invalid_count()
     }
 
+    #[inline]
+    fn validity(&self) -> VortexResult<Validity> {
+        self.as_ref().validity()
+    }
+
     #[inline]
     fn validity_mask(&self) -> Mask {
         self.as_ref().validity_mask()
@@ -601,6 +610,14 @@ impl<V: VTable> Array for ArrayAdapter<V> {
         count
     }
 
+    fn validity(&self) -> VortexResult<Validity> {
+        if self.dtype().is_nullable() {
+            <V::ValidityVTable as ValidityVTable<V>>::validity(&self.0)
+        } else {
+            Ok(Validity::NonNullable)
+        }
+    }
+
     fn validity_mask(&self) -> Mask {
         let mask = <V::ValidityVTable as ValidityVTable<V>>::validity_mask(&self.0);
         assert_eq!(mask.len(), self.len(), "Validity mask length mismatch");
diff --git a/vortex-array/src/arrays/chunked/vtable/validity.rs b/vortex-array/src/arrays/chunked/vtable/validity.rs
@@ -1,11 +1,17 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use itertools::Itertools;
+use vortex_dtype::DType;
+use vortex_dtype::Nullability;
+use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::Array;
+use crate::IntoArray;
 use crate::arrays::ChunkedArray;
 use crate::arrays::ChunkedVTable;
+use crate::validity::Validity;
 use crate::vtable::ValidityVTable;
 
 impl ValidityVTable<ChunkedVTable> for ChunkedVTable {
@@ -41,6 +47,22 @@ impl ValidityVTable<ChunkedVTable> for ChunkedVTable {
         true
     }
 
+    fn validity(array: &ChunkedArray) -> VortexResult<Validity> {
+        Ok(Validity::Array(
+            unsafe {
+                ChunkedArray::new_unchecked(
+                    array
+                        .chunks()
+                        .iter()
+                        .map(|chunk| chunk.validity().map(|v| v.to_array(chunk.len())))
+                        .try_collect()?,
+                    DType::Bool(Nullability::NonNullable),
+                )
+            }
+            .into_array(),
+        ))
+    }
+
     fn validity_mask(array: &ChunkedArray) -> Mask {
         array.chunks().iter().map(|a| a.validity_mask()).collect()
     }
diff --git a/vortex-array/src/arrays/constant/vtable/validity.rs b/vortex-array/src/arrays/constant/vtable/validity.rs
@@ -1,10 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::arrays::ConstantArray;
 use crate::arrays::ConstantVTable;
+use crate::validity::Validity;
 use crate::vtable::ValidityVTable;
 
 impl ValidityVTable<ConstantVTable> for ConstantVTable {
@@ -20,6 +22,14 @@ impl ValidityVTable<ConstantVTable> for ConstantVTable {
         array.scalar().is_null()
     }
 
+    fn validity(array: &ConstantArray) -> VortexResult<Validity> {
+        Ok(if array.scalar().is_null() {
+            Validity::AllInvalid
+        } else {
+            Validity::AllValid
+        })
+    }
+
     fn validity_mask(array: &ConstantArray) -> Mask {
         match array.scalar().is_null() {
             true => Mask::AllFalse(array.len),
diff --git a/vortex-array/src/arrays/dict/vtable/validity.rs b/vortex-array/src/arrays/dict/vtable/validity.rs
@@ -4,13 +4,16 @@
 use vortex_buffer::BitBuffer;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
 use vortex_mask::AllOr;
 use vortex_mask::Mask;
 
 use super::DictVTable;
 use crate::Array;
+use crate::IntoArray;
 use crate::ToCanonical;
 use crate::arrays::dict::DictArray;
+use crate::validity::Validity;
 use crate::vtable::ValidityVTable;
 
 impl ValidityVTable<DictVTable> for DictVTable {
@@ -35,6 +38,36 @@ impl ValidityVTable<DictVTable> for DictVTable {
         array.codes().all_invalid() || array.values().all_invalid()
     }
 
+    fn validity(array: &DictArray) -> VortexResult<Validity> {
+        Ok(
+            match (array.codes().validity()?, array.values().validity()?) {
+                (
+                    Validity::NonNullable | Validity::AllValid,
+                    Validity::NonNullable | Validity::AllValid,
+                ) => {
+                    // Recall that we know the dictionary is nullable if we're in this function.
+                    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)) => {
+                    Validity::Array(
+                        unsafe { DictArray::new_unchecked(array.codes().clone(), values_validity) }
+                            .into_array(),
+                    )
+                }
+                (Validity::Array(_), Validity::Array(values_validity)) => {
+                    // We essentially create is_not_null(Dict(codes, is_not_null(values)))
+                    unsafe { DictArray::new_unchecked(array.codes().clone(), values_validity) }
+                        .into_array()
+                        .validity()?
+                }
+            },
+        )
+    }
+
     fn validity_mask(array: &DictArray) -> Mask {
         let codes_validity = array.codes().validity_mask();
         match codes_validity.bit_buffer() {
diff --git a/vortex-array/src/arrays/extension/vtable/validity.rs b/vortex-array/src/arrays/extension/vtable/validity.rs
@@ -1,13 +1,13 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use crate::Array;
-use crate::arrays::extension::ExtensionArray;
+use crate::ArrayRef;
+use crate::arrays::ExtensionArray;
 use crate::arrays::extension::ExtensionVTable;
 use crate::vtable::ValidityChild;
 
 impl ValidityChild<ExtensionVTable> for ExtensionVTable {
-    fn validity_child(array: &ExtensionArray) -> &dyn Array {
-        array.storage.as_ref()
+    fn validity_child(array: &ExtensionArray) -> &ArrayRef {
+        &array.storage
     }
 }
diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs
diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs
diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs
diff --git a/vortex-array/src/arrow/array.rs b/vortex-array/src/arrow/array.rs
diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs
diff --git a/vortex-array/src/vtable/validity.rs b/vortex-array/src/vtable/validity.rs
diff --git a/vortex-python/src/arrays/py/vtable.rs b/vortex-python/src/arrays/py/vtable.rs

Original file line number	Diff line number	Diff line change
`@@ -432,7 +432,7 @@ impl ALPArray {`
`432`	`432`	`}`
`433`	`433`
`434`	`434`	`impl ValidityChild<ALPVTable> for ALPVTable {`
`435`		`- fn validity_child(array: &ALPArray) -> &dyn Array {`
	`435`	`+ fn validity_child(array: &ALPArray) -> &ArrayRef {`
`436`	`436`	`array.encoded()`
`437`	`437`	`}`
`438`	`438`	`}`
Original file line number	Diff line number	Diff line change
`@@ -368,7 +368,7 @@ impl ALPRDArray {`
`368`	`368`	`}`
`369`	`369`
`370`	`370`	`impl ValidityChild<ALPRDVTable> for ALPRDVTable {`
`371`		`- fn validity_child(array: &ALPRDArray) -> &dyn Array {`
	`371`	`+ fn validity_child(array: &ALPRDArray) -> &ArrayRef {`
`372`	`372`	`array.left_parts()`
`373`	`373`	`}`
`374`	`374`	`}`
Original file line number	Diff line number	Diff line change
`@@ -288,7 +288,7 @@ impl BaseArrayVTable<DateTimePartsVTable> for DateTimePartsVTable {`
`288`	`288`	`}`
`289`	`289`
`290`	`290`	`impl ValidityChild<DateTimePartsVTable> for DateTimePartsVTable {`
`291`		`- fn validity_child(array: &DateTimePartsArray) -> &dyn Array {`
	`291`	`+ fn validity_child(array: &DateTimePartsArray) -> &ArrayRef {`
`292`	`292`	`array.days()`
`293`	`293`	`}`
`294`	`294`	`}`
Original file line number	Diff line number	Diff line change
`@@ -269,9 +269,9 @@ impl OperationsVTable<DecimalBytePartsVTable> for DecimalBytePartsVTable {`
`269`	`269`	`}`
`270`	`270`
`271`	`271`	`impl ValidityChild<DecimalBytePartsVTable> for DecimalBytePartsVTable {`
`272`		`- fn validity_child(array: &DecimalBytePartsArray) -> &dyn Array {`
	`272`	`+ fn validity_child(array: &DecimalBytePartsArray) -> &ArrayRef {`
`273`	`273`	`// validity stored in 0th child`
`274`		`- array.msp.as_ref()`
	`274`	`+ &array.msp`
`275`	`275`	`}`
`276`	`276`	`}`
`277`	`277`