Skip to main content

vortex_array/display/extractors/
stats.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Write;
5use std::fmt::{self};
6
7use crate::ArrayRef;
8use crate::display::extractor::TreeContext;
9use crate::display::extractor::TreeExtractor;
10use crate::expr::stats::Stat;
11use crate::expr::stats::StatsProvider;
12use crate::validity::Validity;
13
14/// Display wrapper for array statistics in compact format.
15///
16/// Produces output like ` [nulls=3, min=5, max=100]` (with leading space).
17pub(crate) struct StatsDisplay<'a>(pub(crate) &'a ArrayRef);
18
19impl fmt::Display for StatsDisplay<'_> {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        let stats = self.0.statistics();
22        let mut first = true;
23
24        let mut sep = |f: &mut fmt::Formatter<'_>| -> fmt::Result {
25            if first {
26                first = false;
27                f.write_str(" [")
28            } else {
29                f.write_str(", ")
30            }
31        };
32
33        // Null count or validity fallback
34        if let Some(nc) = stats.get(Stat::NullCount) {
35            if let Ok(n) = usize::try_from(&nc.clone().into_inner()) {
36                sep(f)?;
37                write!(f, "nulls={}", n)?;
38            } else {
39                sep(f)?;
40                write!(f, "nulls={}", nc)?;
41            }
42        } else if self.0.dtype().is_nullable() {
43            match self.0.validity() {
44                Ok(Validity::NonNullable | Validity::AllValid) => {
45                    sep(f)?;
46                    f.write_str("all_valid")?;
47                }
48                Ok(Validity::AllInvalid) => {
49                    sep(f)?;
50                    f.write_str("all_invalid")?;
51                }
52                Ok(Validity::Array(_)) => {
53                    // Avoid computing validity-array stats as a side effect of display.
54                }
55                Err(e) => {
56                    tracing::warn!("Failed to check validity: {e}");
57                    sep(f)?;
58                    f.write_str("validity_failed")?;
59                }
60            }
61        }
62
63        // NaN count (only if > 0)
64        if let Some(nan) = stats.get(Stat::NaNCount)
65            && let Ok(n) = usize::try_from(&nan.into_inner())
66            && n > 0
67        {
68            sep(f)?;
69            write!(f, "nan={}", n)?;
70        }
71
72        // Min/Max
73        if let Some(min) = stats.get(Stat::Min) {
74            sep(f)?;
75            write!(f, "min={}", min)?;
76        }
77        if let Some(max) = stats.get(Stat::Max) {
78            sep(f)?;
79            write!(f, "max={}", max)?;
80        }
81
82        // Sum
83        if let Some(sum) = stats.get(Stat::Sum) {
84            sep(f)?;
85            write!(f, "sum={}", sum)?;
86        }
87
88        // Boolean flags (compact)
89        if let Some(c) = stats.get(Stat::IsConstant)
90            && bool::try_from(&c.into_inner()).unwrap_or(false)
91        {
92            sep(f)?;
93            f.write_str("const")?;
94        }
95        if let Some(s) = stats.get(Stat::IsStrictSorted) {
96            if bool::try_from(&s.into_inner()).unwrap_or(false) {
97                sep(f)?;
98                f.write_str("strict")?;
99            }
100        } else if let Some(s) = stats.get(Stat::IsSorted)
101            && bool::try_from(&s.into_inner()).unwrap_or(false)
102        {
103            sep(f)?;
104            f.write_str("sorted")?;
105        }
106
107        // Close bracket if we wrote anything
108        if !first {
109            f.write_char(']')?;
110        }
111
112        Ok(())
113    }
114}
115
116/// Extractor that adds stats annotations (e.g. `[nulls=3, min=5]`) to the header line.
117pub struct StatsExtractor;
118
119impl TreeExtractor for StatsExtractor {
120    fn write_header(
121        &self,
122        array: &ArrayRef,
123        _ctx: &TreeContext,
124        f: &mut fmt::Formatter<'_>,
125    ) -> fmt::Result {
126        write!(f, "{}", StatsDisplay(array))
127    }
128}