diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock index b72ced1ebf3..abef723d857 100644 --- a/vortex-array/public-api.lock +++ b/vortex-array/public-api.lock @@ -9938,12 +9938,98 @@ 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::detail_lines(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::BufferExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + 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::MetadataExtractor + +impl vortex_array::display::TreeExtractor for vortex_array::display::MetadataExtractor + +pub fn vortex_array::display::MetadataExtractor::detail_lines(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::MetadataExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub struct vortex_array::display::NbytesExtractor + +impl vortex_array::display::TreeExtractor for vortex_array::display::NbytesExtractor + +pub fn vortex_array::display::NbytesExtractor::detail_lines(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::NbytesExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub struct vortex_array::display::StatsExtractor + +impl vortex_array::display::TreeExtractor for vortex_array::display::StatsExtractor + +pub fn vortex_array::display::StatsExtractor::detail_lines(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::StatsExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub struct vortex_array::display::TreeContext + +impl vortex_array::display::TreeContext + +pub fn vortex_array::display::TreeContext::parent_total_size(&self) -> core::option::Option + +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(self, extractor: E) -> Self + +pub fn vortex_array::display::TreeDisplay::with_boxed(self, extractor: alloc::boxed::Box) -> 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::detail_lines(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::TreeExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +impl vortex_array::display::TreeExtractor for vortex_array::display::BufferExtractor + +pub fn vortex_array::display::BufferExtractor::detail_lines(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::BufferExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +impl vortex_array::display::TreeExtractor for vortex_array::display::MetadataExtractor + +pub fn vortex_array::display::MetadataExtractor::detail_lines(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::MetadataExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +impl vortex_array::display::TreeExtractor for vortex_array::display::NbytesExtractor + +pub fn vortex_array::display::NbytesExtractor::detail_lines(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::NbytesExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +impl vortex_array::display::TreeExtractor for vortex_array::display::StatsExtractor + +pub fn vortex_array::display::StatsExtractor::detail_lines(&self, array: &dyn vortex_array::DynArray, ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + +pub fn vortex_array::display::StatsExtractor::header_annotations(&self, array: &dyn vortex_array::DynArray, _ctx: &vortex_array::display::TreeContext) -> alloc::vec::Vec + pub mod vortex_array::dtype pub use vortex_array::dtype::half diff --git a/vortex-array/src/display/extractor.rs b/vortex-array/src/display/extractor.rs new file mode 100644 index 00000000000..264c9197e8f --- /dev/null +++ b/vortex-array/src/display/extractor.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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>, +} + +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 { + self.ancestor_sizes.last().cloned().flatten() + } + + pub(crate) fn push(&mut self, size: Option) { + self.ancestor_sizes.push(size); + } + + pub(crate) fn pop(&mut self) { + self.ancestor_sizes.pop(); + } +} + +/// 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 { + /// Annotations appended to the header line (e.g., `nbytes=10 B (100.00%)`). + fn header_annotations(&self, array: &dyn DynArray, ctx: &TreeContext) -> Vec { + let _ = (array, ctx); + vec![] + } + + /// Additional detail lines shown below the header (e.g., `metadata: EmptyMetadata`). + fn detail_lines(&self, array: &dyn DynArray, ctx: &TreeContext) -> Vec { + let _ = (array, ctx); + vec![] + } +} diff --git a/vortex-array/src/display/extractors.rs b/vortex-array/src/display/extractors.rs new file mode 100644 index 00000000000..6f0447ee4d1 --- /dev/null +++ b/vortex-array/src/display/extractors.rs @@ -0,0 +1,219 @@ +// 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::DynArray; +use crate::display::extractor::TreeContext; +use crate::display::extractor::TreeExtractor; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProvider; + +/// Display wrapper for array statistics in compact format. +/// +/// Produces output like ` [nulls=3, min=5, max=100]` (with leading space). +pub(crate) struct StatsDisplay<'a>(pub(crate) &'a dyn DynArray); + +impl fmt::Display for StatsDisplay<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let stats = self.0.statistics(); + let mut first = true; + + 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() { + match self.0.all_valid() { + Ok(true) => { + sep(f)?; + f.write_str("all_valid")?; + } + Ok(false) => { + if self.0.all_invalid().unwrap_or(false) { + sep(f)?; + f.write_str("all_invalid")?; + } + } + Err(e) => { + tracing::warn!("Failed to check validity: {e}"); + sep(f)?; + f.write_str("validity_failed")?; + } + } + } + + // 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(()) + } +} + +/// Extractor that adds `nbytes=X (Y%)` to the header line. +pub struct NbytesExtractor; + +impl TreeExtractor for NbytesExtractor { + fn header_annotations(&self, array: &dyn DynArray, ctx: &TreeContext) -> Vec { + let nbytes = array.nbytes(); + let total_size = ctx.parent_total_size().unwrap_or(nbytes); + let percent = if total_size == 0 { + 0.0 + } else { + 100_f64 * nbytes as f64 / total_size as f64 + }; + vec![format!( + "nbytes={} ({:.2}%)", + format_size(nbytes, DECIMAL), + percent + )] + } +} + +/// Extractor that adds stats annotations (e.g. `[nulls=3, min=5]`) to the header line. +pub struct StatsExtractor; + +impl TreeExtractor for StatsExtractor { + fn header_annotations(&self, array: &dyn DynArray, _ctx: &TreeContext) -> Vec { + let s = StatsDisplay(array).to_string(); + let trimmed = s.trim_start(); + if trimmed.is_empty() { + vec![] + } else { + vec![trimmed.to_string()] + } + } +} + +/// Extractor that adds a `metadata: ...` detail line. +pub struct MetadataExtractor; + +impl TreeExtractor for MetadataExtractor { + fn detail_lines(&self, array: &dyn DynArray, _ctx: &TreeContext) -> Vec { + // Capture the metadata_fmt output + let mut buf = String::new(); + // metadata_fmt writes directly to a Formatter, so we use a helper wrapper + struct FmtCapture<'a>(&'a dyn DynArray); + impl fmt::Display for FmtCapture<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.metadata_fmt(f) + } + } + let _ = write!(&mut buf, "{}", FmtCapture(array)); + vec![format!("metadata: {buf}")] + } +} + +/// 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 detail_lines(&self, array: &dyn DynArray, _ctx: &TreeContext) -> Vec { + let nbytes = array.nbytes(); + let mut lines = Vec::new(); + 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 + }; + lines.push(format!( + "buffer: {} {loc} {} (align={}) ({:.2}%)", + name, + format_size(buffer.len(), DECIMAL), + align, + buffer_percent, + )); + } else { + lines.push(format!( + "buffer: {} {loc} {} (align={})", + name, + format_size(buffer.len(), DECIMAL), + align, + )); + } + } + lines + } +} diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 9c34d6642b1..b957ddc51fc 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -1,12 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod extractor; +mod extractors; +mod node; mod tree; +mod tree_display; use std::fmt::Display; +pub use extractor::TreeContext; +pub use extractor::TreeExtractor; +pub use extractors::BufferExtractor; +pub use extractors::MetadataExtractor; +pub use extractors::NbytesExtractor; +pub use extractors::StatsExtractor; use itertools::Itertools as _; use tree::TreeDisplayWrapper; +pub use tree_display::TreeDisplay; use crate::DynArray; @@ -431,6 +442,61 @@ impl dyn DynArray + '_ { ) } + /// Create a tree display with all built-in extractors (nbytes, stats, metadata, buffers). + /// + /// This is the default, fully-detailed tree display. Use + /// `tree_display_builder()` for a blank slate. + /// + /// # Examples + /// ``` + /// # use vortex_array::IntoArray; + /// # use vortex_buffer::buffer; + /// let array = buffer![0_i16, 1, 2, 3, 4].into_array(); + /// let expected = "root: vortex.primitive(i16, len=5) nbytes=10 B (100.00%) + /// metadata: EmptyMetadata + /// buffer: values host 10 B (align=2) (100.00%) + /// "; + /// assert_eq!(array.tree_display().to_string(), expected); + /// ``` + pub fn tree_display(&self) -> TreeDisplay { + TreeDisplay::default_display(self.to_array()) + } + + /// Create a composable tree display builder with no extractors. + /// + /// With no extractors, only encoding headers and the tree structure are shown. + /// Add extractors with [`.with()`][TreeDisplay::with] to include additional information. + /// + /// # Examples + /// ``` + /// # use vortex_array::IntoArray; + /// # use vortex_buffer::buffer; + /// use vortex_array::display::{NbytesExtractor, MetadataExtractor, BufferExtractor}; + /// + /// let array = buffer![0_i16, 1, 2, 3, 4].into_array(); + /// + /// // Encodings only (no extractors) + /// let encodings = array.tree_display_builder().to_string(); + /// assert_eq!(encodings, "root: vortex.primitive(i16, len=5)\n"); + /// + /// // With nbytes + /// let with_nbytes = array.tree_display_builder() + /// .with(NbytesExtractor) + /// .to_string(); + /// assert_eq!(with_nbytes, "root: vortex.primitive(i16, len=5) nbytes=10 B (100.00%)\n"); + /// + /// // With metadata and buffers + /// let detailed = array.tree_display_builder() + /// .with(MetadataExtractor) + /// .with(BufferExtractor { show_percent: false }) + /// .to_string(); + /// let expected = "root: vortex.primitive(i16, len=5)\n metadata: EmptyMetadata\n buffer: values host 10 B (align=2)\n"; + /// assert_eq!(detailed, expected); + /// ``` + pub fn tree_display_builder(&self) -> TreeDisplay { + TreeDisplay::new(self.to_array()) + } + /// Display the array as a formatted table. /// /// For struct arrays, displays a column for each field in the struct. diff --git a/vortex-array/src/display/node.rs b/vortex-array/src/display/node.rs new file mode 100644 index 00000000000..fcf70f5cc40 --- /dev/null +++ b/vortex-array/src/display/node.rs @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; + +/// Internal intermediate representation of a single node in the display tree. +/// +/// Built by collecting annotations from all extractors, then rendered to text. +pub(crate) struct DisplayNode { + /// The name label for this node (e.g. "root", "x", "values"). + pub(crate) name: String, + /// The encoding summary (e.g. "vortex.primitive(i16, len=5)"). + pub(crate) encoding_summary: String, + /// Annotations appended to the header line, collected from extractors. + pub(crate) header_annotations: Vec, + /// Detail lines shown below the header, collected from extractors. + pub(crate) detail_lines: Vec, + /// Recursive children. + pub(crate) children: Vec, +} + +impl DisplayNode { + /// Render this node and all descendants to the formatter with the given indent prefix. + pub(crate) fn render(&self, f: &mut fmt::Formatter<'_>, indent: &str) -> fmt::Result { + // Header line: "{indent}{name}: {encoding_summary} {annotations...}\n" + write!(f, "{indent}{}: {}", self.name, self.encoding_summary)?; + for ann in &self.header_annotations { + write!(f, " {ann}")?; + } + writeln!(f)?; + + // Detail lines + let child_indent = format!("{indent} "); + for line in &self.detail_lines { + writeln!(f, "{child_indent}{line}")?; + } + + // Children + for child in &self.children { + child.render(f, &child_indent)?; + } + + Ok(()) + } +} diff --git a/vortex-array/src/display/tree.rs b/vortex-array/src/display/tree.rs index 7ff8b2991d9..87af9578e46 100644 --- a/vortex-array/src/display/tree.rs +++ b/vortex-array/src/display/tree.rs @@ -1,121 +1,16 @@ // 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 vortex_error::VortexExpect as _; +use std::fmt; use crate::ArrayRef; -use crate::ArrayVisitor; -use crate::DynArray; -use crate::arrays::Chunked; -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 DynArray); - -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() { - match self.0.all_valid() { - Ok(true) => { - sep(f)?; - f.write_str("all_valid")?; - } - Ok(false) => { - if self.0.all_invalid().unwrap_or(false) { - sep(f)?; - f.write_str("all_invalid")?; - } - } - Err(e) => { - tracing::warn!("Failed to check validity: {e}"); - sep(f)?; - f.write_str("validity_failed")?; - } - } - } - - // 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(()) - } -} +use crate::display::extractors::BufferExtractor; +use crate::display::extractors::MetadataExtractor; +use crate::display::extractors::NbytesExtractor; +use crate::display::extractors::StatsExtractor; +use crate::display::tree_display::TreeDisplay; +/// Backward-compatible wrapper that maps the old boolean flags to the new extractor-based system. #[derive(Clone)] pub(crate) struct TreeDisplayWrapper { pub(crate) array: ArrayRef, @@ -125,161 +20,24 @@ pub(crate) struct TreeDisplayWrapper { } impl fmt::Display for TreeDisplayWrapper { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - let TreeDisplayWrapper { - array, - buffers, - metadata, - stats, - } = self.clone(); - let mut array_fmt = TreeFormatter { - fmt, - indent: "".to_string(), - ancestor_sizes: Vec::new(), - buffers, - metadata, - stats, - }; - array_fmt.format("root", array) - } -} - -pub struct TreeFormatter<'a, 'b: 'a> { - fmt: &'a mut fmt::Formatter<'b>, - indent: String, - ancestor_sizes: Vec>, - buffers: bool, - metadata: bool, - stats: bool, -} - -impl<'a, 'b: 'a> TreeFormatter<'a, 'b> { - fn format(&mut self, name: &str, array: ArrayRef) -> fmt::Result { - if self.stats { - let nbytes = array.nbytes(); - let total_size = self - .ancestor_sizes - .last() - .cloned() - .flatten() - .unwrap_or(nbytes); - - self.ancestor_sizes.push(if array.is::() { - // Treat each chunk as a new root - None - } else { - // Children will present themselves as a percentage of our size. - Some(nbytes) - }); - let percent = if total_size == 0 { - 0.0 - } else { - 100_f64 * nbytes as f64 / total_size as f64 - }; - - writeln!( - self, - "{}: {} nbytes={} ({:.2}%){}", - name, - array.display_as(DisplayOptions::MetadataOnly), - format_size(nbytes, DECIMAL), - percent, - StatsDisplay(array.as_ref()), - )?; - } else { - writeln!( - self, - "{}: {}", - name, - array.display_as(DisplayOptions::MetadataOnly) - )?; - } - - self.indent(|i| { - if i.metadata { - write!(i, "metadata: ")?; - array.metadata_fmt(i.fmt)?; - writeln!(i.fmt)?; - } - - if i.buffers { - 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 { - "".to_string() - }; - - if i.stats { - let buffer_percent = if nbytes == 0 { - 0.0 - } else { - 100_f64 * buffer.len() as f64 / nbytes as f64 - }; - writeln!( - i, - "buffer: {} {loc} {} (align={}) ({:.2}%)", - name, - format_size(buffer.len(), DECIMAL), - align, - buffer_percent - )?; - } else { - writeln!( - i, - "buffer: {} {loc} {} (align={})", - name, - format_size(buffer.len(), DECIMAL), - align, - )?; - } - } - } - - Ok(()) - })?; - - self.indent(|i| { - for (name, child) in array - .children_names() - .into_iter() - .zip(array.children().into_iter()) - { - i.format(&name, child)?; + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let extractors: [(bool, Box); 4] = [ + (self.stats, Box::new(NbytesExtractor)), + (self.stats, Box::new(StatsExtractor)), + (self.metadata, Box::new(MetadataExtractor)), + ( + self.buffers, + Box::new(BufferExtractor { + show_percent: self.stats, + }), + ), + ]; + let mut display = TreeDisplay::new(self.array.clone()); + for (enabled, extractor) in extractors { + if enabled { + display = display.with_boxed(extractor); } - Ok(()) - })?; - - if self.stats { - let _ = self - .ancestor_sizes - .pop() - .vortex_expect("pushes and pops are matched"); } - - Ok(()) - } - - fn indent(&mut self, indented: F) -> fmt::Result - where - F: FnOnce(&mut TreeFormatter) -> fmt::Result, - { - let original_ident = self.indent.clone(); - self.indent += " "; - let res = indented(self); - self.indent = original_ident; - res - } - - fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> fmt::Result { - write!(self.fmt, "{}{}", self.indent, fmt) + write!(f, "{display}") } } diff --git a/vortex-array/src/display/tree_display.rs b/vortex-array/src/display/tree_display.rs new file mode 100644 index 00000000000..8cd1e70ee47 --- /dev/null +++ b/vortex-array/src/display/tree_display.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; + +use crate::ArrayRef; +use crate::arrays::Chunked; +use crate::display::DisplayOptions; +use crate::display::extractor::TreeContext; +use crate::display::extractor::TreeExtractor; +use crate::display::extractors::BufferExtractor; +use crate::display::extractors::MetadataExtractor; +use crate::display::extractors::NbytesExtractor; +use crate::display::extractors::StatsExtractor; +use crate::display::node::DisplayNode; + +/// Composable tree display builder. +/// +/// Use `tree_display()` for the default display with all built-in extractors, +/// or `tree_display_builder()` to start with a blank slate and compose your own: +/// +/// ``` +/// # use vortex_array::IntoArray; +/// # use vortex_buffer::buffer; +/// use vortex_array::display::{NbytesExtractor, MetadataExtractor, BufferExtractor}; +/// +/// let array = buffer![0_i16, 1, 2, 3, 4].into_array(); +/// +/// // Default: all built-in extractors +/// let full = array.tree_display(); +/// +/// // Custom: pick only what you need +/// let custom = array.tree_display_builder() +/// .with(NbytesExtractor) +/// .with(MetadataExtractor); +/// ``` +pub struct TreeDisplay { + array: ArrayRef, + extractors: Vec>, +} + +impl TreeDisplay { + /// Create a new tree display for the given array with no extractors. + /// + /// With no extractors, only encoding headers and the tree structure are shown. + /// Use [`Self::default_display`] for the standard set of all built-in extractors. + pub fn new(array: ArrayRef) -> Self { + Self { + array, + extractors: Vec::new(), + } + } + + /// Create a tree display with all built-in extractors: nbytes, stats, metadata, and buffers. + pub fn default_display(array: ArrayRef) -> Self { + Self::new(array) + .with(NbytesExtractor) + .with(StatsExtractor) + .with(MetadataExtractor) + .with(BufferExtractor { show_percent: true }) + } + + /// Add an extractor to the display pipeline. + pub fn with(mut self, extractor: E) -> Self { + self.extractors.push(Box::new(extractor)); + self + } + + /// Add a pre-boxed extractor to the display pipeline. + pub fn with_boxed(mut self, extractor: Box) -> Self { + self.extractors.push(extractor); + self + } + + /// Recursively build the display node tree. + fn build_node(&self, name: &str, array: &ArrayRef, ctx: &mut TreeContext) -> DisplayNode { + // Collect header annotations from all extractors + let header_annotations: Vec = self + .extractors + .iter() + .flat_map(|e| e.header_annotations(array.as_ref(), ctx)) + .collect(); + + // Collect detail lines from all extractors + let detail_lines: Vec = self + .extractors + .iter() + .flat_map(|e| e.detail_lines(array.as_ref(), ctx)) + .collect(); + + // Push context for children: chunked arrays reset the percentage root + let child_size = if array.is::() { + None + } else { + Some(array.nbytes()) + }; + ctx.push(child_size); + + // Recurse into children + let children: Vec = array + .children_names() + .into_iter() + .zip(array.children()) + .map(|(child_name, child)| self.build_node(&child_name, &child, ctx)) + .collect(); + + ctx.pop(); + + DisplayNode { + name: name.to_string(), + encoding_summary: format!("{}", array.display_as(DisplayOptions::MetadataOnly)), + header_annotations, + detail_lines, + children, + } + } +} + +impl fmt::Display for TreeDisplay { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut ctx = TreeContext::new(); + let root = self.build_node("root", &self.array, &mut ctx); + root.render(f, "") + } +}