vortex_array/stats/
rewrite.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Session-registered rewrite rules for aggregate-backed stats expressions.
5
6use std::fmt::Debug;
7use std::sync::Arc;
8
9use vortex_error::VortexResult;
10use vortex_error::vortex_ensure;
11use vortex_session::VortexSession;
12use vortex_utils::iter::ReduceBalancedIterExt;
13
14use crate::dtype::DType;
15use crate::expr::BoundExpression;
16use crate::scalar_fn::ScalarFnId;
17use crate::scalar_fn::ScalarFnVTableExt;
18use crate::scalar_fn::fns::binary::Binary;
19use crate::scalar_fn::fns::operators::Operator;
20use crate::stats::session::StatsSessionExt;
21
22mod builtins;
23
24pub(crate) use builtins::register_builtins;
25
26/// Shared reference to a stats rewrite rule.
27pub type StatsRewriteRuleRef = Arc<dyn StatsRewriteRule>;
28
29/// A plugin-provided rule for predicates whose root scalar function matches this rule.
30///
31/// Rules do not produce expressions equivalent to `expr`. They produce optional sufficient
32/// conditions over stats for the current scope:
33///
34/// - a falsifier evaluating to `true` proves that `expr` is false for every row in the scope;
35/// - a satisfier evaluating to `true` proves that `expr` is true for every row in the scope.
36///
37/// Returning `None` means this rule cannot prove anything for the expression. A returned proof
38/// expression that evaluates to `false` or `null` is also inconclusive.
39///
40/// Multiple rules may be registered for the same scalar function. Their proofs are combined with
41/// `OR`, so every proof returned by an individual rule must be sound on its own.
42///
43/// `expr` is the full predicate expression whose root scalar function id is
44/// [`Self::scalar_fn_id`]. Use [`StatsRewriteCtx`] to resolve dtypes and recursively rewrite child
45/// predicates.
46pub trait StatsRewriteRule: Debug + Send + Sync + 'static {
47    /// Returns the scalar function id handled by this rule.
48    fn scalar_fn_id(&self) -> ScalarFnId;
49
50    /// Returns a stats-backed proof that `expr` is false for the current scope.
51    ///
52    /// If the returned expression evaluates to `true` against the scope's stats, then `expr` is
53    /// guaranteed to be false for every row in that scope. A returned proof expression that
54    /// evaluates to `false` or `null` is inconclusive.
55    ///
56    /// Returns `Ok(None)` when this rule cannot construct a sound falsity proof for `expr`.
57    fn falsify(
58        &self,
59        expr: &BoundExpression,
60        ctx: &StatsRewriteCtx<'_>,
61    ) -> VortexResult<Option<BoundExpression>> {
62        _ = expr;
63        _ = ctx;
64        Ok(None)
65    }
66
67    /// Returns a stats-backed proof that `expr` is true for the current scope.
68    ///
69    /// If the returned expression evaluates to `true` against the scope's stats, then `expr` is
70    /// guaranteed to be true for every row in that scope. A returned proof expression that
71    /// evaluates to `false` or `null` is inconclusive.
72    ///
73    /// This is not the complement of [`Self::falsify`]; both methods are one-way proofs and may be
74    /// implemented independently.
75    ///
76    /// Returns `Ok(None)` when this rule cannot construct a sound truth proof for `expr`.
77    fn satisfy(
78        &self,
79        expr: &BoundExpression,
80        ctx: &StatsRewriteCtx<'_>,
81    ) -> VortexResult<Option<BoundExpression>> {
82        _ = expr;
83        _ = ctx;
84        Ok(None)
85    }
86}
87
88/// Context passed to stats rewrite rules.
89pub struct StatsRewriteCtx<'a> {
90    session: &'a VortexSession,
91}
92
93impl<'a> StatsRewriteCtx<'a> {
94    /// Create a rewrite context for `session`.
95    pub fn new(session: &'a VortexSession) -> Self {
96        Self { session }
97    }
98
99    /// Returns the session that owns the rewrite registry.
100    pub fn session(&self) -> &'a VortexSession {
101        self.session
102    }
103
104    /// Return the dtype of `expr` within this rewrite scope.
105    pub fn return_dtype(&self, expr: &BoundExpression) -> VortexResult<DType> {
106        Ok(expr.dtype().clone())
107    }
108
109    /// Rewrite `expr` into a stats-backed falsifier.
110    pub fn falsify(&self, expr: &BoundExpression) -> VortexResult<Option<BoundExpression>> {
111        self.ensure_predicate(expr)?;
112        rewrite(expr, self, StatsRewriteRule::falsify)
113    }
114
115    /// Rewrite `expr` into a stats-backed satisfier.
116    pub fn satisfy(&self, expr: &BoundExpression) -> VortexResult<Option<BoundExpression>> {
117        self.ensure_predicate(expr)?;
118        rewrite(expr, self, StatsRewriteRule::satisfy)
119    }
120
121    fn ensure_predicate(&self, expr: &BoundExpression) -> VortexResult<()> {
122        let dtype = self.return_dtype(expr)?;
123        vortex_ensure!(
124            matches!(dtype, DType::Bool(_)),
125            "Stats rewrites require a boolean predicate, got {dtype}",
126        );
127        Ok(())
128    }
129}
130
131fn rewrite(
132    expr: &BoundExpression,
133    ctx: &StatsRewriteCtx<'_>,
134    apply: fn(
135        &dyn StatsRewriteRule,
136        &BoundExpression,
137        &StatsRewriteCtx<'_>,
138    ) -> VortexResult<Option<BoundExpression>>,
139) -> VortexResult<Option<BoundExpression>> {
140    // The scope alone proves nothing about the rows it contains.
141    let Some(scalar_fn) = expr.as_scalar() else {
142        return Ok(None);
143    };
144    let rules = ctx.session().stats().rewrite_rules_for(scalar_fn.id());
145    let Some(rules) = rules else {
146        return Ok(None);
147    };
148
149    let mut rewrites = Vec::new();
150    for rule in rules.iter() {
151        if let Some(rewrite) = apply(rule.as_ref(), expr, ctx)? {
152            rewrites.push(rewrite);
153        }
154    }
155
156    rewrites
157        .into_iter()
158        .try_reduce_balanced(|lhs, rhs| Binary.try_new_bound_expr(Operator::Or, [lhs, rhs]))
159}
160
161#[cfg(test)]
162mod tests {
163    use vortex_error::VortexResult;
164
165    use super::StatsRewriteCtx;
166    use super::StatsRewriteRule;
167    use crate::dtype::DType;
168    use crate::dtype::Nullability;
169    use crate::dtype::PType;
170    use crate::expr::BoundExpression;
171    use crate::expr::lit;
172    use crate::expr::or;
173    use crate::scalar_fn::ScalarFnId;
174    use crate::scalar_fn::ScalarFnVTable;
175    use crate::scalar_fn::fns::literal::Literal;
176    use crate::stats::session::StatsSessionExt;
177
178    #[derive(Debug)]
179    struct StaticLiteralRule {
180        falsifier: Option<BoundExpression>,
181        satisfier: Option<BoundExpression>,
182    }
183
184    impl StatsRewriteRule for StaticLiteralRule {
185        fn scalar_fn_id(&self) -> ScalarFnId {
186            Literal.id()
187        }
188
189        fn falsify(
190            &self,
191            _expr: &BoundExpression,
192            _ctx: &StatsRewriteCtx<'_>,
193        ) -> VortexResult<Option<BoundExpression>> {
194            Ok(self.falsifier.clone())
195        }
196
197        fn satisfy(
198            &self,
199            _expr: &BoundExpression,
200            _ctx: &StatsRewriteCtx<'_>,
201        ) -> VortexResult<Option<BoundExpression>> {
202            Ok(self.satisfier.clone())
203        }
204    }
205
206    #[test]
207    fn combines_multiple_falsifiers_with_or() -> VortexResult<()> {
208        let session = crate::array_session();
209        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
210        session.stats().register_rewrite(StaticLiteralRule {
211            falsifier: Some(lit(false).bind(&dtype)?),
212            satisfier: None,
213        });
214        session.stats().register_rewrite(StaticLiteralRule {
215            falsifier: Some(lit(true).bind(&dtype)?),
216            satisfier: None,
217        });
218
219        assert_eq!(
220            lit(true).bind(&dtype)?.falsify(&session)?,
221            Some(or(lit(false), lit(true)).bind(&dtype)?)
222        );
223        Ok(())
224    }
225
226    #[test]
227    fn combines_multiple_satisfiers_with_or() -> VortexResult<()> {
228        let session = crate::array_session();
229        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
230        session.stats().register_rewrite(StaticLiteralRule {
231            falsifier: None,
232            satisfier: Some(lit(false).bind(&dtype)?),
233        });
234        session.stats().register_rewrite(StaticLiteralRule {
235            falsifier: None,
236            satisfier: Some(lit(true).bind(&dtype)?),
237        });
238
239        assert_eq!(
240            lit(true).bind(&dtype)?.satisfy(&session)?,
241            Some(or(lit(false), lit(true)).bind(&dtype)?)
242        );
243        Ok(())
244    }
245
246    #[test]
247    fn unregistered_expression_has_no_rewrite() -> VortexResult<()> {
248        let session = crate::array_session();
249        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
250
251        let expr = lit(true).bind(&dtype)?;
252        assert_eq!(expr.falsify(&session)?, None);
253        assert_eq!(expr.satisfy(&session)?, None);
254        Ok(())
255    }
256
257    #[test]
258    fn non_predicate_expression_errors() -> VortexResult<()> {
259        let session = crate::array_session();
260        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
261
262        let expr = lit(7).bind(&dtype)?;
263        assert!(expr.falsify(&session).is_err());
264        assert!(expr.satisfy(&session).is_err());
265        Ok(())
266    }
267}
vortex_array/stats/rewrite.rs

vortex_array/stats/
rewrite.rs
