vortex_array/scalar_fn/fns/
stat.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Scalar function implementation for aggregate-backed stat expressions.
5
6use std::fmt::Display;
7use std::fmt::Formatter;
8
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::IntoArray;
16use crate::aggregate_fn::AggregateFnRef;
17use crate::aggregate_fn::fns::all_nan::AllNan;
18use crate::aggregate_fn::fns::all_non_nan::AllNonNan;
19use crate::aggregate_fn::fns::all_non_null::AllNonNull;
20use crate::aggregate_fn::fns::all_null::AllNull;
21use crate::arrays::ConstantArray;
22use crate::dtype::DType;
23use crate::expr::display::ExprDisplay;
24use crate::expr::stats::Precision;
25use crate::expr::stats::Stat;
26use crate::expr::stats::StatsProvider;
27use crate::expr::stats::StatsProviderExt;
28use crate::scalar::Scalar;
29use crate::scalar::ScalarValue;
30use crate::scalar_fn::Arity;
31use crate::scalar_fn::ChildName;
32use crate::scalar_fn::ExecutionArgs;
33use crate::scalar_fn::ScalarFnId;
34use crate::scalar_fn::ScalarFnVTable;
35
36/// Options for the `stat` scalar function.
37#[derive(Clone, Debug, PartialEq, Eq, Hash)]
38pub struct StatOptions {
39    aggregate_fn: AggregateFnRef,
40}
41
42impl StatOptions {
43    /// Creates options for the provided aggregate statistic.
44    pub fn new(aggregate_fn: AggregateFnRef) -> Self {
45        Self { aggregate_fn }
46    }
47
48    /// Returns the aggregate function backing this statistic lookup.
49    pub fn aggregate_fn(&self) -> &AggregateFnRef {
50        &self.aggregate_fn
51    }
52}
53
54impl Display for StatOptions {
55    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
56        Display::fmt(&self.aggregate_fn, f)
57    }
58}
59
60/// Scalar function that broadcasts a stored aggregate partial over the input rows.
61///
62/// The only current consumer is **row-wise pruning**: substituting `stat(col, agg)` into a
63/// predicate produces a cheap, row-aligned approximation whose constant runs let downstream
64/// filters drop entire stretches at once. For example, `value < 10` is prunable as
65/// `stat(value, max) < 10` (rows where the bound is false are guaranteed false) or
66/// `stat(value, min) >= 10` (rows where it is true are guaranteed true) — the zone-map /
67/// min-max-index pattern, expressed as an ordinary expression so the existing scalar
68/// machinery can rewrite, fold, and execute it.
69///
70/// The result is row-aligned with the input, at whatever granularity the input carries the
71/// stat at: e.g. a flat array yields a single broadcast `ConstantArray`; a chunked array
72/// yields a constant per chunk; a zone-mapped array would yield a run-end-encoded array,
73/// one run per zone. If the requested stat is not available, the result is a null constant.
74///
75/// Pruning only makes sense for aggregates that can prove something about every row in the scope
76/// — `min`, `max`, `all_null`, `all_non_null`, bloom filters, etc. Non-idempotent aggregates like
77/// `sum`, `count`, `mean`, `null_count`, and `nan_count` still produce a meaningful per-chunk
78/// value but do **not** bound any single row.
79#[derive(Clone)]
80pub struct StatFn;
81
82impl ScalarFnVTable for StatFn {
83    type Options = StatOptions;
84
85    fn id(&self) -> ScalarFnId {
86        static ID: CachedId = CachedId::new("vortex.stat");
87        *ID
88    }
89
90    fn arity(&self, _options: &Self::Options) -> Arity {
91        Arity::Exact(1)
92    }
93
94    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
95        match child_idx {
96            0 => ChildName::from("input"),
97            _ => unreachable!("Invalid child index {} for Stat expression", child_idx),
98        }
99    }
100
101    fn fmt_sql(
102        &self,
103        options: &Self::Options,
104        expr: &dyn ExprDisplay,
105        f: &mut Formatter<'_>,
106    ) -> std::fmt::Result {
107        write!(f, "stat(")?;
108        Display::fmt(expr.display_child(0), f)?;
109        write!(f, ", {})", options.aggregate_fn())
110    }
111
112    fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
113        stat_dtype(options.aggregate_fn(), &arg_dtypes[0])
114    }
115
116    fn execute(
117        &self,
118        options: &Self::Options,
119        args: &dyn ExecutionArgs,
120        _ctx: &mut ExecutionCtx,
121    ) -> VortexResult<ArrayRef> {
122        let input = args.get(0)?;
123        let dtype = stat_dtype(options.aggregate_fn(), input.dtype())?;
124        stat_array(&input, options.aggregate_fn(), dtype, args.row_count())
125    }
126
127    fn is_strict(&self, _options: &Self::Options) -> bool {
128        false
129    }
130}
131
132fn stat_dtype(aggregate_fn: &AggregateFnRef, input_dtype: &DType) -> VortexResult<DType> {
133    let Some(dtype) = aggregate_fn.state_dtype(input_dtype) else {
134        vortex_bail!(
135            "Aggregate function {} does not support input dtype {}",
136            aggregate_fn,
137            input_dtype
138        );
139    };
140    Ok(dtype.as_nullable())
141}
142
143fn stat_array(
144    array: &ArrayRef,
145    aggregate_fn: &AggregateFnRef,
146    dtype: DType,
147    len: usize,
148) -> VortexResult<ArrayRef> {
149    let value = if aggregate_fn.is::<AllNull>() {
150        let len = u64::try_from(len)?;
151        match array.statistics().get_as::<u64>(Stat::NullCount) {
152            Precision::Exact(count) => Some(count == len),
153            Precision::Inexact(count) => (count < len).then_some(false),
154            Precision::Absent => None,
155        }
156        .map(ScalarValue::Bool)
157    } else if aggregate_fn.is::<AllNonNull>() {
158        match array.statistics().get_as::<u64>(Stat::NullCount) {
159            Precision::Exact(count) => Some(count == 0),
160            Precision::Inexact(0) => Some(true),
161            Precision::Inexact(_) | Precision::Absent => None,
162        }
163        .map(ScalarValue::Bool)
164    } else if aggregate_fn.is::<AllNan>() {
165        let len = u64::try_from(len)?;
166        match array.statistics().get_as::<u64>(Stat::NaNCount) {
167            Precision::Exact(count) => Some(count == len),
168            Precision::Inexact(count) => (count < len).then_some(false),
169            Precision::Absent => None,
170        }
171        .map(ScalarValue::Bool)
172    } else if aggregate_fn.is::<AllNonNan>() {
173        match array.statistics().get_as::<u64>(Stat::NaNCount) {
174            Precision::Exact(count) => Some(count == 0),
175            Precision::Inexact(0) => Some(true),
176            Precision::Inexact(_) | Precision::Absent => None,
177        }
178        .map(ScalarValue::Bool)
179    } else if let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) {
180        array
181            .statistics()
182            .with_typed_stats_set(|stats| stats.get(stat))
183            // We don't mind whether the stat is approxed or not, since these are row-wise bounds.
184            .into_inner()
185            .and_then(Scalar::into_value)
186    } else {
187        tracing::trace!(
188            "No legacy Stat slot for aggregate {}; stat expression will resolve to null",
189            aggregate_fn
190        );
191        None
192    };
193
194    let scalar = Scalar::try_new(dtype, value)?;
195    Ok(ConstantArray::new(scalar, len).into_array())
196}
vortex_array/scalar_fn/fns/stat.rs

vortex_array/scalar_fn/fns/
stat.rs
