vortex_array/stats/
bind.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Bind abstract `vortex.stat` expressions to a concrete stats representation.
5//!
6//! Stats rewrite rules describe pruning in terms of `vortex.stat(input, aggregate_fn)` placeholders
7//! so the rewrite is independent of where statistics are stored. These stat placeholders are
8//! abstract because they name the statistic needed for a proof, but not how that statistic is
9//! represented by a specific layout or reader.
10//!
11//! Binding is the later pass that replaces each abstract placeholder with the representation used
12//! by a caller: zone-map field references, file-level stat literals, or typed nulls for missing
13//! stats. This lets all callers share the same falsification rules while keeping layout-specific
14//! stat storage behind [`StatBinder`].
15
16use vortex_error::VortexResult;
17
18use crate::aggregate_fn::AggregateFnRef;
19use crate::dtype::DType;
20use crate::expr::BoundExpression;
21use crate::expr::bound::lit;
22use crate::expr::traversal::NodeExt;
23use crate::expr::traversal::Transformed;
24use crate::scalar::Scalar;
25use crate::scalar_fn::fns::stat::StatFn;
26
27/// A target that can bind abstract statistics to concrete expressions.
28///
29/// Implementations define how a pruning proof should read stats from a specific backing
30/// representation. For example, a zone-map binder can translate a `max(col)` placeholder into a
31/// field reference in the per-zone stats table, while a file-stats binder can translate the same
32/// placeholder into a literal value from the file footer.
33pub trait StatBinder {
34    /// Bind `aggregate_fn(input)` to a concrete expression.
35    ///
36    /// Implementations should return `Ok(None)` when the requested aggregate
37    /// statistic is unavailable in their backing representation.
38    fn bind_aggregate(
39        &self,
40        input: &BoundExpression,
41        aggregate_fn: &AggregateFnRef,
42        stat_dtype: &DType,
43    ) -> VortexResult<Option<BoundExpression>>;
44
45    /// Expression to use when a stat is unavailable.
46    ///
47    /// The default is a nullable null literal, which preserves three-valued
48    /// pruning semantics for stats-table execution.
49    fn missing_stat(&self, dtype: DType) -> VortexResult<BoundExpression> {
50        null_expr(dtype)
51    }
52}
53
54/// Bind all `vortex.stat` expressions in `predicate`.
55///
56/// The predicate is usually the output of a stats rewrite rule. Rewrite rules
57/// are responsible for expressing stat semantics; binding maps aggregate-backed
58/// stat requests to the concrete stats representation supported by the binder.
59pub fn bind_stats<B: StatBinder + ?Sized>(
60    predicate: BoundExpression,
61    binder: &B,
62) -> VortexResult<BoundExpression> {
63    Ok(predicate
64        .transform_down(|expr| {
65            if !expr.is::<StatFn>() {
66                return Ok(Transformed::no(expr));
67            }
68
69            match bind_stat_fn(&expr, binder)? {
70                Some(bound) => Ok(Transformed::yes(bound)),
71                None => Ok(Transformed::yes(binder.missing_stat(expr.dtype().clone())?)),
72            }
73        })?
74        .into_inner())
75}
76
77fn bind_stat_fn(
78    expr: &BoundExpression,
79    binder: &(impl StatBinder + ?Sized),
80) -> VortexResult<Option<BoundExpression>> {
81    let options = expr.as_::<StatFn>();
82    let aggregate_fn = options.aggregate_fn();
83    // `StatFn` has exactly one child: the expression the aggregate statistic is computed over.
84    let input = expr.child(0);
85
86    binder.bind_aggregate(input, aggregate_fn, expr.dtype())
87}
88
89fn null_expr(dtype: DType) -> VortexResult<BoundExpression> {
90    Ok(lit(Scalar::null(dtype.as_nullable())))
91}
92
93#[cfg(test)]
94mod tests {
95    use vortex_error::VortexResult;
96
97    use super::*;
98    use crate::dtype::Nullability;
99    use crate::dtype::PType;
100    use crate::dtype::StructFields;
101    use crate::expr::and;
102    use crate::expr::col;
103    use crate::expr::get_item;
104    use crate::expr::is_null;
105    use crate::expr::lit;
106    use crate::expr::or;
107    use crate::expr::root;
108    use crate::expr::stats::Stat;
109    use crate::stats::all_non_nan;
110    use crate::stats::nan_count;
111
112    struct TestBinder {
113        input_scope: DType,
114        stats_scope: DType,
115        bind_nan_count: bool,
116    }
117
118    impl TestBinder {
119        fn new(bind_nan_count: bool) -> Self {
120            Self {
121                input_scope: DType::Struct(
122                    StructFields::from_iter([(
123                        "f",
124                        DType::Primitive(PType::F32, Nullability::NonNullable),
125                    )]),
126                    Nullability::NonNullable,
127                ),
128                stats_scope: DType::Struct(
129                    StructFields::from_iter([(
130                        "f_nan_count",
131                        DType::Primitive(PType::U64, Nullability::NonNullable),
132                    )]),
133                    Nullability::NonNullable,
134                ),
135                bind_nan_count,
136            }
137        }
138    }
139
140    impl StatBinder for TestBinder {
141        fn bind_aggregate(
142            &self,
143            _input: &BoundExpression,
144            aggregate_fn: &AggregateFnRef,
145            _stat_dtype: &DType,
146        ) -> VortexResult<Option<BoundExpression>> {
147            let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else {
148                return Ok(None);
149            };
150
151            if stat == Stat::NaNCount && self.bind_nan_count {
152                Ok(Some(
153                    get_item("f_nan_count", root()).bind(&self.stats_scope)?,
154                ))
155            } else {
156                Ok(None)
157            }
158        }
159    }
160
161    #[test]
162    fn nan_count_binds_to_direct_stat_slot() -> VortexResult<()> {
163        let binder = TestBinder::new(true);
164
165        let bound = bind_stats(nan_count(col("f")).bind(&binder.input_scope)?, &binder)?;
166
167        assert_eq!(bound, col("f_nan_count").bind(&binder.stats_scope)?);
168        Ok(())
169    }
170
171    #[test]
172    fn all_non_nan_does_not_derive_from_nan_count() -> VortexResult<()> {
173        let binder = TestBinder::new(true);
174
175        let bound = bind_stats(all_non_nan(col("f")).bind(&binder.input_scope)?, &binder)?;
176
177        assert_eq!(
178            bound,
179            lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)?
180        );
181        Ok(())
182    }
183
184    #[test]
185    fn missing_stats_bind_to_null_without_reducing() -> VortexResult<()> {
186        let binder = TestBinder::new(false);
187        let null_bool = lit(Scalar::null(DType::Bool(Nullability::Nullable)));
188
189        let bound = bind_stats(
190            and(lit(false), all_non_nan(col("f"))).bind(&binder.input_scope)?,
191            &binder,
192        )?;
193
194        assert_eq!(
195            bound,
196            and(lit(false), null_bool.clone()).bind(&binder.stats_scope)?
197        );
198
199        let bound = bind_stats(
200            or(lit(true), all_non_nan(col("f"))).bind(&binder.input_scope)?,
201            &binder,
202        )?;
203
204        assert_eq!(bound, or(lit(true), null_bool).bind(&binder.stats_scope)?);
205        Ok(())
206    }
207
208    #[test]
209    fn unrelated_expressions_do_not_request_nan_count() -> VortexResult<()> {
210        let binder = TestBinder::new(false);
211
212        let bound = bind_stats(is_null(col("f")).bind(&binder.input_scope)?, &binder)?;
213
214        assert_eq!(bound, is_null(col("f")).bind(&binder.input_scope)?);
215        Ok(())
216    }
217}
vortex_array/stats/bind.rs

vortex_array/stats/
bind.rs
