feat[array]: validity in display array (#5799)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/compute/zip.rs

@@ -329,7 +329,7 @@ mod tests {
         let result = zip(&const1, &const2, &mask).unwrap();
 
         insta::assert_snapshot!(result.display_tree(), @r"
-        root: vortex.varbinview(utf8?, len=100) nbytes=1.66 kB (100.00%)
+        root: vortex.varbinview(utf8?, len=100) nbytes=1.66 kB (100.00%) [all_valid]
           metadata: EmptyMetadata
           buffer (align=1): 29 B (1.75%)
           buffer (align=1): 28 B (1.69%)
@@ -348,7 +348,7 @@ mod tests {
         insta::assert_snapshot!(wrapped_result.display_tree(), @r"
         root: vortex.struct({nested=utf8?}, len=100) nbytes=1.66 kB (100.00%)
           metadata: EmptyMetadata
-          nested: vortex.varbinview(utf8?, len=100) nbytes=1.66 kB (100.00%)
+          nested: vortex.varbinview(utf8?, len=100) nbytes=1.66 kB (100.00%) [all_valid]
             metadata: EmptyMetadata
             buffer (align=1): 29 B (1.75%)
             buffer (align=1): 28 B (1.69%)

vortex-array/src/display/tree.rs

@@ -1,15 +1,109 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::fmt::Write;
 use std::fmt::{self};
 
 use humansize::DECIMAL;
 use humansize::format_size;
 
+use crate::Array;
 use crate::ArrayRef;
 use crate::ArrayVisitor;
 use crate::arrays::ChunkedVTable;
 use crate::display::DisplayOptions;
+use crate::expr::stats::Stat;
+use crate::expr::stats::StatsProvider;
+
+/// Display wrapper for array statistics in compact format.
+struct StatsDisplay<'a>(&'a dyn Array);
+
+impl fmt::Display for StatsDisplay<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let stats = self.0.statistics();
+        let mut first = true;
+
+        // Helper to write separator
+        let mut sep = |f: &mut fmt::Formatter<'_>| -> fmt::Result {
+            if first {
+                first = false;
+                f.write_str(" [")
+            } else {
+                f.write_str(", ")
+            }
+        };
+
+        // Null count or validity fallback
+        if let Some(nc) = stats.get(Stat::NullCount) {
+            if let Ok(n) = usize::try_from(&nc.clone().into_inner()) {
+                sep(f)?;
+                write!(f, "nulls={}", n)?;
+            } else {
+                sep(f)?;
+                write!(f, "nulls={}", nc)?;
+            }
+        } else if self.0.dtype().is_nullable() {
+            if self.0.all_valid() {
+                sep(f)?;
+                f.write_str("all_valid")?;
+            } else if self.0.all_invalid() {
+                sep(f)?;
+                f.write_str("all_invalid")?;
+            }
+        }
+
+        // NaN count (only if > 0)
+        if let Some(nan) = stats.get(Stat::NaNCount)
+            && let Ok(n) = usize::try_from(&nan.into_inner())
+            && n > 0
+        {
+            sep(f)?;
+            write!(f, "nan={}", n)?;
+        }
+
+        // Min/Max
+        if let Some(min) = stats.get(Stat::Min) {
+            sep(f)?;
+            write!(f, "min={}", min)?;
+        }
+        if let Some(max) = stats.get(Stat::Max) {
+            sep(f)?;
+            write!(f, "max={}", max)?;
+        }
+
+        // Sum
+        if let Some(sum) = stats.get(Stat::Sum) {
+            sep(f)?;
+            write!(f, "sum={}", sum)?;
+        }
+
+        // Boolean flags (compact)
+        if let Some(c) = stats.get(Stat::IsConstant)
+            && bool::try_from(&c.into_inner()).unwrap_or(false)
+        {
+            sep(f)?;
+            f.write_str("const")?;
+        }
+        if let Some(s) = stats.get(Stat::IsStrictSorted) {
+            if bool::try_from(&s.into_inner()).unwrap_or(false) {
+                sep(f)?;
+                f.write_str("strict")?;
+            }
+        } else if let Some(s) = stats.get(Stat::IsSorted)
+            && bool::try_from(&s.into_inner()).unwrap_or(false)
+        {
+            sep(f)?;
+            f.write_str("sorted")?;
+        }
+
+        // Close bracket if we wrote anything
+        if !first {
+            f.write_char(']')?;
+        }
+
+        Ok(())
+    }
+}
 
 pub(crate) struct TreeDisplayWrapper(pub(crate) ArrayRef);
 
@@ -39,13 +133,15 @@ impl<'a, 'b: 'a> TreeFormatter<'a, 'b> {
         } else {
             100_f64 * nbytes as f64 / total_size as f64
         };
+
         writeln!(
             self,
-            "{}: {} nbytes={} ({:.2}%)",
+            "{}: {} nbytes={} ({:.2}%){}",
             name,
             array.display_as(DisplayOptions::MetadataOnly),
             format_size(nbytes, DECIMAL),
-            percent
+            percent,
+            StatsDisplay(array.as_ref()),
         )?;
 
         self.indent(|i| {

vortex-python/src/arrays/mod.rs

@@ -650,11 +650,11 @@ impl PyArray {
     /// ```python
     /// >>> import vortex as vx
     /// >>> arr = vx.array([1, 2, None, 3])
-    /// >>> print(arr.display_tree())
+    /// >>> print(arr.display_tree()) # doctest: +ELLIPSIS
     /// root: vortex.primitive(i64?, len=4) nbytes=33 B (100.00%)
     ///   metadata: EmptyMetadata
     ///   buffer (align=8): 32 B (96.97%)
-    ///   validity: vortex.bool(bool, len=4) nbytes=1 B (3.03%)
+    ///   validity: vortex.bool(bool, len=4) nbytes=1 B (3.03%)...
     ///     metadata: BoolMetadata { offset: 0 }
     ///     buffer (align=1): 1 B (100.00%)
     /// <BLANKLINE>