Thread scope dtype through stats rewrites

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-array/public-api.lock

@@ -12918,13 +12918,13 @@ pub fn vortex_array::expr::Expression::children(&self) -> &alloc::sync::Arc<allo
 
 pub fn vortex_array::expr::Expression::display_tree(&self) -> impl core::fmt::Display
 
-pub fn vortex_array::expr::Expression::falsify(&self, &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+pub fn vortex_array::expr::Expression::falsify(&self, &vortex_array::dtype::DType, &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
 
 pub fn vortex_array::expr::Expression::fmt_sql(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
 pub fn vortex_array::expr::Expression::return_dtype(&self, &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
 
-pub fn vortex_array::expr::Expression::satisfy(&self, &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
+pub fn vortex_array::expr::Expression::satisfy(&self, &vortex_array::dtype::DType, &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::Expression>>
 
 pub fn vortex_array::expr::Expression::scalar_fn(&self) -> &vortex_array::scalar_fn::ScalarFnRef
 

vortex-array/src/expr/expression.rs

@@ -138,18 +138,30 @@ impl Expression {
 
     /// Returns an expression that proves this predicate is definitely false from stats.
     ///
+    /// `scope` is the dtype of the row this expression evaluates over.
+    ///
     /// If the returned expression evaluates to `true` for a stats scope, this expression is
     /// guaranteed to be false for every row in that scope. `false` and `null` are unknown.
-    pub fn falsify(&self, session: &VortexSession) -> VortexResult<Option<Expression>> {
-        crate::stats::rewrite::StatsRewriteCtx::new(session).falsify(self)
+    pub fn falsify(
+        &self,
+        scope: &DType,
+        session: &VortexSession,
+    ) -> VortexResult<Option<Expression>> {
+        crate::stats::rewrite::StatsRewriteCtx::new(session, scope).falsify(self)
     }
 
     /// Returns an expression that proves this predicate is definitely true from stats.
     ///
+    /// `scope` is the dtype of the row this expression evaluates over.
+    ///
     /// If the returned expression evaluates to `true` for a stats scope, this expression is
     /// guaranteed to be true for every row in that scope. `false` and `null` are unknown.
-    pub fn satisfy(&self, session: &VortexSession) -> VortexResult<Option<Expression>> {
-        crate::stats::rewrite::StatsRewriteCtx::new(session).satisfy(self)
+    pub fn satisfy(
+        &self,
+        scope: &DType,
+        session: &VortexSession,
+    ) -> VortexResult<Option<Expression>> {
+        crate::stats::rewrite::StatsRewriteCtx::new(session, scope).satisfy(self)
     }
 
     /// Returns an expression representing the zoned statistic for the given stat, if available.

vortex-array/src/stats/rewrite.rs

@@ -7,8 +7,10 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 use vortex_session::VortexSession;
 
+use crate::dtype::DType;
 use crate::expr::Expression;
 use crate::expr::or_collect;
 use crate::scalar_fn::ScalarFnId;
@@ -54,28 +56,45 @@ pub(crate) trait StatsRewriteRule: Debug + Send + Sync + 'static {
 /// Context passed to stats rewrite rules.
 pub(crate) struct StatsRewriteCtx<'a> {
     session: &'a VortexSession,
+    scope: &'a DType,
 }
 
 impl<'a> StatsRewriteCtx<'a> {
     /// Create a rewrite context for `session`.
-    pub(crate) fn new(session: &'a VortexSession) -> Self {
-        Self { session }
+    pub(crate) fn new(session: &'a VortexSession, scope: &'a DType) -> Self {
+        Self { session, scope }
     }
 
     /// Returns the session that owns the rewrite registry.
     pub(crate) fn session(&self) -> &'a VortexSession {
         self.session
     }
 
+    /// Return the dtype of `expr` within this rewrite scope.
+    pub(crate) fn return_dtype(&self, expr: &Expression) -> VortexResult<DType> {
+        expr.return_dtype(self.scope)
+    }
+
     /// Rewrite `expr` into a stats-backed falsifier.
     pub(crate) fn falsify(&self, expr: &Expression) -> VortexResult<Option<Expression>> {
+        self.ensure_predicate(expr)?;
         rewrite(expr, self, StatsRewriteRule::falsify)
     }
 
     /// Rewrite `expr` into a stats-backed satisfier.
     pub(crate) fn satisfy(&self, expr: &Expression) -> VortexResult<Option<Expression>> {
+        self.ensure_predicate(expr)?;
         rewrite(expr, self, StatsRewriteRule::satisfy)
     }
+
+    fn ensure_predicate(&self, expr: &Expression) -> VortexResult<()> {
+        let dtype = self.return_dtype(expr)?;
+        vortex_ensure!(
+            matches!(dtype, DType::Bool(_)),
+            "Stats rewrites require a boolean predicate, got {dtype}",
+        );
+        Ok(())
+    }
 }
 
 fn rewrite(
@@ -112,6 +131,9 @@ mod tests {
 
     use super::StatsRewriteCtx;
     use super::StatsRewriteRule;
+    use crate::dtype::DType;
+    use crate::dtype::Nullability;
+    use crate::dtype::PType;
     use crate::expr::Expression;
     use crate::expr::lit;
     use crate::expr::or;
@@ -152,6 +174,7 @@ mod tests {
     #[test]
     fn combines_multiple_falsifiers_with_or() -> VortexResult<()> {
         let session = VortexSession::empty().with::<StatsRewriteSession>();
+        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
         session.stats_rewrites().register(StaticLiteralRule {
             falsifier: Some(lit(false)),
             satisfier: None,
@@ -161,13 +184,17 @@ mod tests {
             satisfier: None,
         });
 
-        assert_eq!(lit(7).falsify(&session)?, Some(or(lit(false), lit(true))));
+        assert_eq!(
+            lit(true).falsify(&dtype, &session)?,
+            Some(or(lit(false), lit(true)))
+        );
         Ok(())
     }
 
     #[test]
     fn combines_multiple_satisfiers_with_or() -> VortexResult<()> {
         let session = VortexSession::empty().with::<StatsRewriteSession>();
+        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
         session.stats_rewrites().register(StaticLiteralRule {
             falsifier: None,
             satisfier: Some(lit(false)),
@@ -177,16 +204,29 @@ mod tests {
             satisfier: Some(lit(true)),
         });
 
-        assert_eq!(lit(7).satisfy(&session)?, Some(or(lit(false), lit(true))));
+        assert_eq!(
+            lit(true).satisfy(&dtype, &session)?,
+            Some(or(lit(false), lit(true)))
+        );
         Ok(())
     }
 
     #[test]
     fn unregistered_expression_has_no_rewrite() -> VortexResult<()> {
         let session = VortexSession::empty().with::<StatsRewriteSession>();
+        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
 
-        assert_eq!(lit(7).falsify(&session)?, None);
-        assert_eq!(lit(7).satisfy(&session)?, None);
+        assert_eq!(lit(true).falsify(&dtype, &session)?, None);
+        assert_eq!(lit(true).satisfy(&dtype, &session)?, None);
         Ok(())
     }
+
+    #[test]
+    fn non_predicate_expression_errors() {
+        let session = VortexSession::empty().with::<StatsRewriteSession>();
+        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
+
+        assert!(lit(7).falsify(&dtype, &session).is_err());
+        assert!(lit(7).satisfy(&dtype, &session).is_err());
+    }
 }