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    let Some(scalar_fn) = expr.as_scalar() else {
141        return Ok(None);
142    };
143    let rules = ctx.session().stats().rewrite_rules_for(scalar_fn.id());
144    let Some(rules) = rules else {
145        return Ok(None);
146    };
147
148    let mut rewrites = Vec::new();
149    for rule in rules.iter() {
150        if let Some(rewrite) = apply(rule.as_ref(), expr, ctx)? {
151            rewrites.push(rewrite);
152        }
153    }
154
155    rewrites
156        .into_iter()
157        .try_reduce_balanced(|lhs, rhs| Binary.try_new_bound_expr(Operator::Or, [lhs, rhs]))
158}
159
160#[cfg(test)]
161mod tests {
162    use vortex_error::VortexResult;
163
164    use super::StatsRewriteCtx;
165    use super::StatsRewriteRule;
166    use crate::dtype::DType;
167    use crate::dtype::Nullability;
168    use crate::dtype::PType;
169    use crate::expr::BoundExpression;
170    use crate::expr::lit;
171    use crate::expr::or;
172    use crate::scalar_fn::ScalarFnId;
173    use crate::scalar_fn::ScalarFnVTable;
174    use crate::scalar_fn::fns::literal::Literal;
175    use crate::stats::session::StatsSessionExt;
176
177    #[derive(Debug)]
178    struct StaticLiteralRule {
179        falsifier: Option<BoundExpression>,
180        satisfier: Option<BoundExpression>,
181    }
182
183    impl StatsRewriteRule for StaticLiteralRule {
184        fn scalar_fn_id(&self) -> ScalarFnId {
185            Literal.id()
186        }
187
188        fn falsify(
189            &self,
190            _expr: &BoundExpression,
191            _ctx: &StatsRewriteCtx<'_>,
192        ) -> VortexResult<Option<BoundExpression>> {
193            Ok(self.falsifier.clone())
194        }
195
196        fn satisfy(
197            &self,
198            _expr: &BoundExpression,
199            _ctx: &StatsRewriteCtx<'_>,
200        ) -> VortexResult<Option<BoundExpression>> {
201            Ok(self.satisfier.clone())
202        }
203    }
204
205    #[test]
206    fn combines_multiple_falsifiers_with_or() -> VortexResult<()> {
207        let session = crate::array_session();
208        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
209        session.stats().register_rewrite(StaticLiteralRule {
210            falsifier: Some(lit(false).bind(&dtype)?),
211            satisfier: None,
212        });
213        session.stats().register_rewrite(StaticLiteralRule {
214            falsifier: Some(lit(true).bind(&dtype)?),
215            satisfier: None,
216        });
217
218        assert_eq!(
219            lit(true).bind(&dtype)?.falsify(&session)?,
220            Some(or(lit(false), lit(true)).bind(&dtype)?)
221        );
222        Ok(())
223    }
224
225    #[test]
226    fn combines_multiple_satisfiers_with_or() -> VortexResult<()> {
227        let session = crate::array_session();
228        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
229        session.stats().register_rewrite(StaticLiteralRule {
230            falsifier: None,
231            satisfier: Some(lit(false).bind(&dtype)?),
232        });
233        session.stats().register_rewrite(StaticLiteralRule {
234            falsifier: None,
235            satisfier: Some(lit(true).bind(&dtype)?),
236        });
237
238        assert_eq!(
239            lit(true).bind(&dtype)?.satisfy(&session)?,
240            Some(or(lit(false), lit(true)).bind(&dtype)?)
241        );
242        Ok(())
243    }
244
245    #[test]
246    fn unregistered_expression_has_no_rewrite() -> VortexResult<()> {
247        let session = crate::array_session();
248        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
249
250        let expr = lit(true).bind(&dtype)?;
251        assert_eq!(expr.falsify(&session)?, None);
252        assert_eq!(expr.satisfy(&session)?, None);
253        Ok(())
254    }
255
256    #[test]
257    fn non_predicate_expression_errors() -> VortexResult<()> {
258        let session = crate::array_session();
259        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
260
261        let expr = lit(7).bind(&dtype)?;
262        assert!(expr.falsify(&session).is_err());
263        assert!(expr.satisfy(&session).is_err());
264        Ok(())
265    }
266}
vortex_array/stats/rewrite.rs

vortex_array/stats/
rewrite.rs
