Simplify StatsRewriteCtx rules by leveraging BoundExpressions

vortex-array/src/expr/bound_expression.rs

@@ -22,7 +22,8 @@ use crate::expr::traversal::TraversalOrder;
 use crate::expr::traversal::pre_order_visit_down;
 use crate::scalar_fn::ScalarFnRef;
 use crate::scalar_fn::ScalarFnVTable;
-use crate::stats::rewrite::StatsRewriteCtx;
+use crate::stats::rewrite::falsify;
+use crate::stats::rewrite::satisfy;
 
 /// An [`Expression`] that has been type-checked against a [`Scope`].
 ///
@@ -242,12 +243,12 @@ impl BoundExpression {
 
     /// Return an expression that proves this predicate is definitely false from statistics.
     pub fn falsify(&self, session: &VortexSession) -> VortexResult<Option<BoundExpression>> {
-        StatsRewriteCtx::new(session).falsify(self)
+        falsify(self, session)
     }
 
     /// Return an expression that proves this predicate is definitely true from statistics.
     pub fn satisfy(&self, session: &VortexSession) -> VortexResult<Option<BoundExpression>> {
-        StatsRewriteCtx::new(session).satisfy(self)
+        satisfy(self, session)
     }
 
     /// Display the bound expression as a formatted tree structure.

vortex-array/src/stats/rewrite.rs

@@ -41,8 +41,7 @@ pub type StatsRewriteRuleRef = Arc<dyn StatsRewriteRule>;
 /// `OR`, so every proof returned by an individual rule must be sound on its own.
 ///
 /// `expr` is the full predicate expression whose root scalar function id is
-/// [`Self::scalar_fn_id`]. Use [`StatsRewriteCtx`] to resolve dtypes and recursively rewrite child
-/// predicates.
+/// [`Self::scalar_fn_id`].
 pub trait StatsRewriteRule: Debug + Send + Sync + 'static {
     /// Returns the scalar function id handled by this rule.
     fn scalar_fn_id(&self) -> ScalarFnId;
@@ -57,10 +56,10 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static {
     fn falsify(
         &self,
         expr: &BoundExpression,
-        ctx: &StatsRewriteCtx<'_>,
+        session: &VortexSession,
     ) -> VortexResult<Option<BoundExpression>> {
         _ = expr;
-        _ = ctx;
+        _ = session;
         Ok(None)
     }
 
@@ -77,78 +76,62 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static {
     fn satisfy(
         &self,
         expr: &BoundExpression,
-        ctx: &StatsRewriteCtx<'_>,
+        session: &VortexSession,
     ) -> VortexResult<Option<BoundExpression>> {
         _ = expr;
-        _ = ctx;
+        _ = session;
         Ok(None)
     }
 }
 
-/// Context passed to stats rewrite rules.
-pub struct StatsRewriteCtx<'a> {
-    session: &'a VortexSession,
+fn ensure_predicate(expr: &BoundExpression) -> VortexResult<()> {
+    let dtype = expr.dtype();
+    vortex_ensure!(
+        matches!(dtype, DType::Bool(_)),
+        "Stats rewrites require a boolean predicate, got {dtype}",
+    );
+    Ok(())
 }
 
-impl<'a> StatsRewriteCtx<'a> {
-    /// Create a rewrite context for `session`.
-    pub fn new(session: &'a VortexSession) -> Self {
-        Self { session }
-    }
-
-    /// Returns the session that owns the rewrite registry.
-    pub fn session(&self) -> &'a VortexSession {
-        self.session
-    }
-
-    /// Return the dtype of `expr` within this rewrite scope.
-    pub fn return_dtype(&self, expr: &BoundExpression) -> VortexResult<DType> {
-        Ok(expr.dtype().clone())
-    }
-
-    /// Rewrite `expr` into a stats-backed falsifier.
-    pub fn falsify(&self, expr: &BoundExpression) -> VortexResult<Option<BoundExpression>> {
-        self.ensure_predicate(expr)?;
-        rewrite(expr, self, StatsRewriteRule::falsify)
-    }
-
-    /// Rewrite `expr` into a stats-backed satisfier.
-    pub fn satisfy(&self, expr: &BoundExpression) -> VortexResult<Option<BoundExpression>> {
-        self.ensure_predicate(expr)?;
-        rewrite(expr, self, StatsRewriteRule::satisfy)
-    }
+/// Rewrite `expr` into a stats-backed falsifier.
+pub fn falsify(
+    expr: &BoundExpression,
+    session: &VortexSession,
+) -> VortexResult<Option<BoundExpression>> {
+    ensure_predicate(expr)?;
+    rewrite(expr, session, StatsRewriteRule::falsify)
+}
 
-    fn ensure_predicate(&self, expr: &BoundExpression) -> VortexResult<()> {
-        let dtype = self.return_dtype(expr)?;
-        vortex_ensure!(
-            matches!(dtype, DType::Bool(_)),
-            "Stats rewrites require a boolean predicate, got {dtype}",
-        );
-        Ok(())
-    }
+/// Rewrite `expr` into a stats-backed satisfier.
+pub fn satisfy(
+    expr: &BoundExpression,
+    session: &VortexSession,
+) -> VortexResult<Option<BoundExpression>> {
+    ensure_predicate(expr)?;
+    rewrite(expr, session, StatsRewriteRule::satisfy)
 }
 
 fn rewrite(
     expr: &BoundExpression,
-    ctx: &StatsRewriteCtx<'_>,
+    session: &VortexSession,
     apply: fn(
         &dyn StatsRewriteRule,
         &BoundExpression,
-        &StatsRewriteCtx<'_>,
+        &VortexSession,
     ) -> VortexResult<Option<BoundExpression>>,
 ) -> VortexResult<Option<BoundExpression>> {
     // The scope alone proves nothing about the rows it contains.
     let Some(scalar_fn) = expr.as_scalar() else {
         return Ok(None);
     };
-    let rules = ctx.session().stats().rewrite_rules_for(scalar_fn.id());
+    let rules = session.stats().rewrite_rules_for(scalar_fn.id());
     let Some(rules) = rules else {
         return Ok(None);
     };
 
     let mut rewrites = Vec::new();
     for rule in rules.iter() {
-        if let Some(rewrite) = apply(rule.as_ref(), expr, ctx)? {
+        if let Some(rewrite) = apply(rule.as_ref(), expr, session)? {
             rewrites.push(rewrite);
         }
     }
@@ -161,8 +144,8 @@ fn rewrite(
 #[cfg(test)]
 mod tests {
     use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
 
-    use super::StatsRewriteCtx;
     use super::StatsRewriteRule;
     use crate::dtype::DType;
     use crate::dtype::Nullability;
@@ -189,15 +172,15 @@ mod tests {
         fn falsify(
             &self,
             _expr: &BoundExpression,
-            _ctx: &StatsRewriteCtx<'_>,
+            _session: &VortexSession,
         ) -> VortexResult<Option<BoundExpression>> {
             Ok(self.falsifier.clone())
         }
 
         fn satisfy(
             &self,
             _expr: &BoundExpression,
-            _ctx: &StatsRewriteCtx<'_>,
+            _session: &VortexSession,
         ) -> VortexResult<Option<BoundExpression>> {
             Ok(self.satisfier.clone())
         }