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`].
45pub trait StatsRewriteRule: Debug + Send + Sync + 'static {
46    /// Returns the scalar function id handled by this rule.
47    fn scalar_fn_id(&self) -> ScalarFnId;
48
49    /// Returns a stats-backed proof that `expr` is false for the current scope.
50    ///
51    /// If the returned expression evaluates to `true` against the scope's stats, then `expr` is
52    /// guaranteed to be false for every row in that scope. A returned proof expression that
53    /// evaluates to `false` or `null` is inconclusive.
54    ///
55    /// Returns `Ok(None)` when this rule cannot construct a sound falsity proof for `expr`.
56    fn falsify(
57        &self,
58        expr: &BoundExpression,
59        session: &VortexSession,
60    ) -> VortexResult<Option<BoundExpression>> {
61        _ = expr;
62        _ = session;
63        Ok(None)
64    }
65
66    /// Returns a stats-backed proof that `expr` is true for the current scope.
67    ///
68    /// If the returned expression evaluates to `true` against the scope's stats, then `expr` is
69    /// guaranteed to be true for every row in that scope. A returned proof expression that
70    /// evaluates to `false` or `null` is inconclusive.
71    ///
72    /// This is not the complement of [`Self::falsify`]; both methods are one-way proofs and may be
73    /// implemented independently.
74    ///
75    /// Returns `Ok(None)` when this rule cannot construct a sound truth proof for `expr`.
76    fn satisfy(
77        &self,
78        expr: &BoundExpression,
79        session: &VortexSession,
80    ) -> VortexResult<Option<BoundExpression>> {
81        _ = expr;
82        _ = session;
83        Ok(None)
84    }
85}
86
87fn ensure_predicate(expr: &BoundExpression) -> VortexResult<()> {
88    let dtype = expr.dtype();
89    vortex_ensure!(
90        matches!(dtype, DType::Bool(_)),
91        "Stats rewrites require a boolean predicate, got {dtype}",
92    );
93    Ok(())
94}
95
96/// Rewrite `expr` into a stats-backed falsifier.
97pub fn falsify(
98    expr: &BoundExpression,
99    session: &VortexSession,
100) -> VortexResult<Option<BoundExpression>> {
101    ensure_predicate(expr)?;
102    rewrite(expr, session, StatsRewriteRule::falsify)
103}
104
105/// Rewrite `expr` into a stats-backed satisfier.
106pub fn satisfy(
107    expr: &BoundExpression,
108    session: &VortexSession,
109) -> VortexResult<Option<BoundExpression>> {
110    ensure_predicate(expr)?;
111    rewrite(expr, session, StatsRewriteRule::satisfy)
112}
113
114fn rewrite(
115    expr: &BoundExpression,
116    session: &VortexSession,
117    apply: fn(
118        &dyn StatsRewriteRule,
119        &BoundExpression,
120        &VortexSession,
121    ) -> VortexResult<Option<BoundExpression>>,
122) -> VortexResult<Option<BoundExpression>> {
123    // The scope alone proves nothing about the rows it contains.
124    let Some(scalar_fn) = expr.as_scalar() else {
125        return Ok(None);
126    };
127    let rules = session.stats().rewrite_rules_for(scalar_fn.id());
128    let Some(rules) = rules else {
129        return Ok(None);
130    };
131
132    let mut rewrites = Vec::new();
133    for rule in rules.iter() {
134        if let Some(rewrite) = apply(rule.as_ref(), expr, session)? {
135            rewrites.push(rewrite);
136        }
137    }
138
139    rewrites
140        .into_iter()
141        .try_reduce_balanced(|lhs, rhs| Binary.try_new_bound_expr(Operator::Or, [lhs, rhs]))
142}
143
144#[cfg(test)]
145mod tests {
146    use vortex_error::VortexResult;
147    use vortex_session::VortexSession;
148
149    use super::StatsRewriteRule;
150    use crate::dtype::DType;
151    use crate::dtype::Nullability;
152    use crate::dtype::PType;
153    use crate::expr::BoundExpression;
154    use crate::expr::lit;
155    use crate::expr::not;
156    use crate::expr::or;
157    use crate::scalar_fn::ScalarFnId;
158    use crate::scalar_fn::ScalarFnVTable;
159    use crate::scalar_fn::fns::not::Not;
160    use crate::stats::session::StatsSessionExt;
161
162    #[derive(Debug)]
163    struct StaticNotRule {
164        falsifier: Option<BoundExpression>,
165        satisfier: Option<BoundExpression>,
166    }
167
168    impl StatsRewriteRule for StaticNotRule {
169        fn scalar_fn_id(&self) -> ScalarFnId {
170            Not.id()
171        }
172
173        fn falsify(
174            &self,
175            _expr: &BoundExpression,
176            _session: &VortexSession,
177        ) -> VortexResult<Option<BoundExpression>> {
178            Ok(self.falsifier.clone())
179        }
180
181        fn satisfy(
182            &self,
183            _expr: &BoundExpression,
184            _session: &VortexSession,
185        ) -> VortexResult<Option<BoundExpression>> {
186            Ok(self.satisfier.clone())
187        }
188    }
189
190    #[test]
191    fn combines_multiple_falsifiers_with_or() -> VortexResult<()> {
192        let session = crate::array_session();
193        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
194        session.stats().register_rewrite(StaticNotRule {
195            falsifier: Some(lit(false).bind(&dtype)?),
196            satisfier: None,
197        });
198        session.stats().register_rewrite(StaticNotRule {
199            falsifier: Some(lit(true).bind(&dtype)?),
200            satisfier: None,
201        });
202
203        assert_eq!(
204            not(lit(true)).bind(&dtype)?.falsify(&session)?,
205            Some(or(lit(false), lit(true)).bind(&dtype)?)
206        );
207        Ok(())
208    }
209
210    #[test]
211    fn combines_multiple_satisfiers_with_or() -> VortexResult<()> {
212        let session = crate::array_session();
213        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
214        session.stats().register_rewrite(StaticNotRule {
215            falsifier: None,
216            satisfier: Some(lit(false).bind(&dtype)?),
217        });
218        session.stats().register_rewrite(StaticNotRule {
219            falsifier: None,
220            satisfier: Some(lit(true).bind(&dtype)?),
221        });
222
223        assert_eq!(
224            not(lit(true)).bind(&dtype)?.satisfy(&session)?,
225            Some(or(lit(false), lit(true)).bind(&dtype)?)
226        );
227        Ok(())
228    }
229
230    #[test]
231    fn unregistered_expression_has_no_rewrite() -> VortexResult<()> {
232        let session = crate::array_session();
233        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
234
235        let expr = not(lit(true)).bind(&dtype)?;
236        assert_eq!(expr.falsify(&session)?, None);
237        assert_eq!(expr.satisfy(&session)?, None);
238        Ok(())
239    }
240
241    #[test]
242    fn non_predicate_expression_errors() -> VortexResult<()> {
243        let session = crate::array_session();
244        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
245
246        let expr = lit(7).bind(&dtype)?;
247        assert!(expr.falsify(&session).is_err());
248        assert!(expr.satisfy(&session).is_err());
249        Ok(())
250    }
251}
vortex_array/stats/rewrite.rs

vortex_array/stats/
rewrite.rs
