feat[array]: compositional array display logic (#7082)

Breaks up the display logic for array tree into multiple composable
extractors

---------

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-array/public-api.lock

@@ -10168,12 +10168,126 @@ impl core::default::Default for vortex_array::display::DisplayOptions
 
 pub fn vortex_array::display::DisplayOptions::default() -> Self
 
+pub struct vortex_array::display::BufferExtractor
+
+pub vortex_array::display::BufferExtractor::show_percent: bool
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::BufferExtractor
+
+pub fn vortex_array::display::BufferExtractor::write_details(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::BufferExtractor::write_header(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
 pub struct vortex_array::display::DisplayArrayAs<'a>(pub &'a dyn vortex_array::DynArray, pub vortex_array::display::DisplayOptions)
 
 impl core::fmt::Display for vortex_array::display::DisplayArrayAs<'_>
 
 pub fn vortex_array::display::DisplayArrayAs<'_>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
+pub struct vortex_array::display::EncodingSummaryExtractor
+
+impl vortex_array::display::EncodingSummaryExtractor
+
+pub fn vortex_array::display::EncodingSummaryExtractor::write(array: &dyn vortex_array::DynArray, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::EncodingSummaryExtractor
+
+pub fn vortex_array::display::EncodingSummaryExtractor::write_details(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::EncodingSummaryExtractor::write_header(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_array::display::IndentedFormatter<'a, 'b>
+
+impl<'a, 'b> vortex_array::display::IndentedFormatter<'a, 'b>
+
+pub fn vortex_array::display::IndentedFormatter<'a, 'b>::formatter(&mut self) -> &mut core::fmt::Formatter<'b>
+
+pub fn vortex_array::display::IndentedFormatter<'a, 'b>::indent(&self) -> &str
+
+pub fn vortex_array::display::IndentedFormatter<'a, 'b>::parts(&mut self) -> (&str, &mut core::fmt::Formatter<'b>)
+
+pub struct vortex_array::display::MetadataExtractor
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::MetadataExtractor
+
+pub fn vortex_array::display::MetadataExtractor::write_details(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::MetadataExtractor::write_header(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_array::display::NbytesExtractor
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::NbytesExtractor
+
+pub fn vortex_array::display::NbytesExtractor::write_details(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::NbytesExtractor::write_header(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_array::display::StatsExtractor
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::StatsExtractor
+
+pub fn vortex_array::display::StatsExtractor::write_details(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::StatsExtractor::write_header(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_array::display::TreeContext
+
+impl vortex_array::display::TreeContext
+
+pub fn vortex_array::display::TreeContext::parent_total_size(&self) -> core::option::Option<u64>
+
+pub struct vortex_array::display::TreeDisplay
+
+impl vortex_array::display::TreeDisplay
+
+pub fn vortex_array::display::TreeDisplay::default_display(array: vortex_array::ArrayRef) -> Self
+
+pub fn vortex_array::display::TreeDisplay::new(array: vortex_array::ArrayRef) -> Self
+
+pub fn vortex_array::display::TreeDisplay::with<E: vortex_array::display::TreeExtractor + 'static>(self, extractor: E) -> Self
+
+pub fn vortex_array::display::TreeDisplay::with_boxed(self, extractor: alloc::boxed::Box<dyn vortex_array::display::TreeExtractor>) -> Self
+
+impl core::fmt::Display for vortex_array::display::TreeDisplay
+
+pub fn vortex_array::display::TreeDisplay::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub trait vortex_array::display::TreeExtractor: core::marker::Send + core::marker::Sync
+
+pub fn vortex_array::display::TreeExtractor::write_details(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::TreeExtractor::write_header(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::BufferExtractor
+
+pub fn vortex_array::display::BufferExtractor::write_details(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::BufferExtractor::write_header(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::EncodingSummaryExtractor
+
+pub fn vortex_array::display::EncodingSummaryExtractor::write_details(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::EncodingSummaryExtractor::write_header(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::MetadataExtractor
+
+pub fn vortex_array::display::MetadataExtractor::write_details(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::MetadataExtractor::write_header(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::NbytesExtractor
+
+pub fn vortex_array::display::NbytesExtractor::write_details(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::NbytesExtractor::write_header(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::display::TreeExtractor for vortex_array::display::StatsExtractor
+
+pub fn vortex_array::display::StatsExtractor::write_details(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext, f: &mut vortex_array::display::IndentedFormatter<'_, '_>) -> core::fmt::Result
+
+pub fn vortex_array::display::StatsExtractor::write_header(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
 pub mod vortex_array::dtype
 
 pub use vortex_array::dtype::half

vortex-array/src/display/extractor.rs

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt;
+
+use crate::DynArray;
+
+/// Context threaded through tree traversal for percentage calculations etc.
+pub struct TreeContext {
+    /// Stack of ancestor nbytes values. `None` entries reset the percentage root
+    /// (e.g. for chunked arrays where each chunk is its own root).
+    pub(crate) ancestor_sizes: Vec<Option<u64>>,
+}
+
+impl TreeContext {
+    pub(crate) fn new() -> Self {
+        Self {
+            ancestor_sizes: Vec::new(),
+        }
+    }
+
+    /// The total size used as the denominator for percentage calculations.
+    /// Returns `None` if there is no ancestor (i.e., this node is the root or
+    /// a chunk boundary reset the percentage root).
+    pub fn parent_total_size(&self) -> Option<u64> {
+        self.ancestor_sizes.last().cloned().flatten()
+    }
+
+    pub(crate) fn push(&mut self, size: Option<u64>) {
+        self.ancestor_sizes.push(size);
+    }
+
+    pub(crate) fn pop(&mut self) {
+        self.ancestor_sizes.pop();
+    }
+}
+
+/// Wrapper providing access to a [`fmt::Formatter`] and the current indentation string.
+pub struct IndentedFormatter<'a, 'b> {
+    inner: &'a mut fmt::Formatter<'b>,
+    indent: &'a str,
+}
+
+impl<'a, 'b> IndentedFormatter<'a, 'b> {
+    pub(crate) fn new(f: &'a mut fmt::Formatter<'b>, indent: &'a str) -> Self {
+        Self { inner: f, indent }
+    }
+
+    /// Access the indent string and underlying [`fmt::Formatter`] together.
+    pub fn parts(&mut self) -> (&str, &mut fmt::Formatter<'b>) {
+        (self.indent, self.inner)
+    }
+
+    /// The current indentation string.
+    pub fn indent(&self) -> &str {
+        self.indent
+    }
+
+    /// Access the underlying [`fmt::Formatter`].
+    pub fn formatter(&mut self) -> &mut fmt::Formatter<'b> {
+        self.inner
+    }
+}
+
+/// Trait for contributing display information to tree nodes.
+///
+/// Each extractor represents one "dimension" of display (e.g., nbytes, stats, metadata, buffers).
+/// Extractors are composable: you can combine any number of them via [`TreeDisplay::with`].
+///
+/// [`TreeDisplay::with`]: super::TreeDisplay::with
+pub trait TreeExtractor: Send + Sync {
+    /// Write header annotations (space-prefixed) to the formatter.
+    fn write_header(
+        &self,
+        array: &dyn DynArray,
+        ctx: &TreeContext,
+        f: &mut fmt::Formatter<'_>,
+    ) -> fmt::Result {
+        let _ = (array, ctx, f);
+        Ok(())
+    }
+
+    /// Write detail lines below the header.
+    ///
+    /// Content written through `f` is automatically indented. Use
+    /// [`f.formatter()`](IndentedFormatter::formatter) to access the underlying
+    /// [`fmt::Formatter`] for formatting flags.
+    fn write_details(
+        &self,
+        array: &dyn DynArray,
+        ctx: &TreeContext,
+        f: &mut IndentedFormatter<'_, '_>,
+    ) -> fmt::Result {
+        let _ = (array, ctx, f);
+        Ok(())
+    }
+}

vortex-array/src/display/extractors/buffer.rs

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt;
+
+use humansize::DECIMAL;
+use humansize::format_size;
+
+use crate::DynArray;
+use crate::display::extractor::IndentedFormatter;
+use crate::display::extractor::TreeContext;
+use crate::display::extractor::TreeExtractor;
+
+/// Extractor that adds buffer detail lines.
+pub struct BufferExtractor {
+    /// Whether to show buffer-level percentage of parent nbytes.
+    pub show_percent: bool,
+}
+
+impl TreeExtractor for BufferExtractor {
+    fn write_details(
+        &self,
+        array: &dyn DynArray,
+        _ctx: &TreeContext,
+        f: &mut IndentedFormatter<'_, '_>,
+    ) -> fmt::Result {
+        let (indent, f) = f.parts();
+        let nbytes = array.nbytes();
+        for (name, buffer) in array.named_buffers() {
+            let loc = if buffer.is_on_device() {
+                "device"
+            } else if buffer.is_on_host() {
+                "host"
+            } else {
+                "location-unknown"
+            };
+            let align = if buffer.is_on_host() {
+                buffer.as_host().alignment().to_string()
+            } else {
+                String::new()
+            };
+
+            if self.show_percent {
+                let buffer_percent = if nbytes == 0 {
+                    0.0
+                } else {
+                    100_f64 * buffer.len() as f64 / nbytes as f64
+                };
+                writeln!(
+                    f,
+                    "{indent}buffer: {} {loc} {} (align={}) ({:.2}%)",
+                    name,
+                    format_size(buffer.len(), DECIMAL),
+                    align,
+                    buffer_percent,
+                )?;
+            } else {
+                writeln!(
+                    f,
+                    "{indent}buffer: {} {loc} {} (align={})",
+                    name,
+                    format_size(buffer.len(), DECIMAL),
+                    align,
+                )?;
+            }
+        }
+        Ok(())
+    }
+}

vortex-array/src/display/extractors/encoding_summary.rs

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt;
+
+use crate::DynArray;
+use crate::display::extractor::TreeContext;
+use crate::display::extractor::TreeExtractor;
+
+/// Extractor that adds the encoding summary (e.g. `vortex.primitive(i16, len=5)`) to the header.
+pub struct EncodingSummaryExtractor;
+
+impl EncodingSummaryExtractor {
+    /// Write the encoding summary for an array directly to a formatter.
+    pub fn write(array: &dyn DynArray, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}({}, len={})",
+            array.encoding_id(),
+            array.dtype(),
+            array.len()
+        )
+    }
+}
+
+impl TreeExtractor for EncodingSummaryExtractor {
+    fn write_header(
+        &self,
+        array: &dyn DynArray,
+        _ctx: &TreeContext,
+        f: &mut fmt::Formatter<'_>,
+    ) -> fmt::Result {
+        write!(f, " ")?;
+        Self::write(array, f)
+    }
+}

vortex-array/src/display/extractors/metadata.rs

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt;
+
+use crate::DynArray;
+use crate::display::extractor::IndentedFormatter;
+use crate::display::extractor::TreeContext;
+use crate::display::extractor::TreeExtractor;
+
+/// Extractor that adds a `metadata: ...` detail line.
+pub struct MetadataExtractor;
+
+impl TreeExtractor for MetadataExtractor {
+    fn write_details(
+        &self,
+        array: &dyn DynArray,
+        _ctx: &TreeContext,
+        f: &mut IndentedFormatter<'_, '_>,
+    ) -> fmt::Result {
+        let (indent, f) = f.parts();
+        write!(f, "{indent}metadata: ")?;
+        array.metadata_fmt(f)?;
+        writeln!(f)
+    }
+}

vortex-array/src/display/extractors/mod.rs

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod buffer;
+mod encoding_summary;
+mod metadata;
+mod nbytes;
+mod stats;
+
+pub use buffer::BufferExtractor;
+pub use encoding_summary::EncodingSummaryExtractor;
+pub use metadata::MetadataExtractor;
+pub use nbytes::NbytesExtractor;
+pub use stats::StatsExtractor;