Drop the fallible dtype lookup from stats rewrites

vortex-array/src/stats/rewrite.rs

@@ -41,8 +41,8 @@ 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`]. Read dtypes from [`BoundExpression::dtype`], and use
+/// [`StatsRewriteCtx`] to recursively rewrite child predicates.
 pub trait StatsRewriteRule: Debug + Send + Sync + 'static {
     /// Returns the scalar function id handled by this rule.
     fn scalar_fn_id(&self) -> ScalarFnId;
@@ -101,11 +101,6 @@ impl<'a> StatsRewriteCtx<'a> {
         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)?;
@@ -119,7 +114,7 @@ impl<'a> StatsRewriteCtx<'a> {
     }
 
     fn ensure_predicate(&self, expr: &BoundExpression) -> VortexResult<()> {
-        let dtype = self.return_dtype(expr)?;
+        let dtype = expr.dtype();
         vortex_ensure!(
             matches!(dtype, DType::Bool(_)),
             "Stats rewrites require a boolean predicate, got {dtype}",

vortex-array/src/stats/rewrite/builtins.rs

@@ -556,11 +556,9 @@ impl NonNanProof for NanCountProof {
     const EMIT_UNGUARDED_REWRITES: bool = true;
 
     fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult<NanCheck> {
-        non_nan_check(ctx, expr, |expr| {
-            match stat_expr(expr, Stat::NaNCount, ctx) {
-                Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))),
-                None => NanCheck::Unavailable,
-            }
+        non_nan_check(expr, |expr| match stat_expr(expr, Stat::NaNCount, ctx) {
+            Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))),
+            None => NanCheck::Unavailable,
         })
     }
 }
@@ -570,8 +568,8 @@ struct AllNonNanProof;
 impl NonNanProof for AllNonNanProof {
     const EMIT_UNGUARDED_REWRITES: bool = false;
 
-    fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult<NanCheck> {
-        non_nan_check(ctx, expr, |expr| {
+    fn check(_ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult<NanCheck> {
+        non_nan_check(expr, |expr| {
             NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions)))
         })
     }
@@ -581,7 +579,6 @@ impl NonNanProof for AllNonNanProof {
 // candidate value is known to be non-NaN. Cast result dtypes are not enough: a cast
 // from float to non-float still needs a proof about the float source values.
 fn non_nan_check(
-    ctx: &StatsRewriteCtx<'_>,
     expr: &BoundExpression,
     proof: impl FnOnce(&BoundExpression) -> NanCheck,
 ) -> VortexResult<NanCheck> {
@@ -597,14 +594,14 @@ fn non_nan_check(
     }
 
     if expr.is::<Cast>() {
-        if !has_nans(&ctx.return_dtype(expr.child(0))?) {
+        if !has_nans(expr.child(0).dtype()) {
             return Ok(NanCheck::NotNeeded);
         }
 
-        return non_nan_check(ctx, expr.child(0), proof);
+        return non_nan_check(expr.child(0), proof);
     }
 
-    if !has_nans(&ctx.return_dtype(expr)?) {
+    if !has_nans(expr.dtype()) {
         return Ok(NanCheck::NotNeeded);
     }
 
@@ -639,9 +636,8 @@ fn stat_expr(
     // The aggregate may not support the expression's dtype, e.g. min/max over structs,
     // even when the predicate itself is well-typed. Such stats cannot be lowered later,
     // so do not reference them in the rewrite.
-    let input_dtype = ctx.return_dtype(expr).ok()?;
     aggregate_fn
-        .return_dtype(&input_dtype)
+        .return_dtype(expr.dtype())
         .is_some()
         .then(|| stat_fn(expr.clone(), aggregate_fn))
 }

vortex-spatial/src/prune/distance.rs

@@ -77,7 +77,7 @@ impl StatsRewriteRule for SpatialDistancePrune {
             return Ok(None);
         }
 
-        let Some((geom, constant)) = geometry_and_constant(distance, ctx)? else {
+        let Some((geom, constant)) = geometry_and_constant(distance) else {
             return Ok(None);
         };
         let Some(query) = query_aabb(constant, ctx)? else {

vortex-spatial/src/prune/intersects.rs

@@ -37,7 +37,7 @@ impl StatsRewriteRule for SpatialIntersectsPrune {
         expr: &BoundExpression,
         ctx: &StatsRewriteCtx<'_>,
     ) -> VortexResult<Option<BoundExpression>> {
-        let Some((geom, constant)) = geometry_and_constant(expr, ctx)? else {
+        let Some((geom, constant)) = geometry_and_constant(expr) else {
             return Ok(None);
         };
         let Some(query) = query_aabb(constant, ctx)? else {

vortex-spatial/src/prune/mod.rs

@@ -51,10 +51,7 @@ use crate::extension::single_geometry;
 /// shape (in either operand order), or the column's dtype carries no [`GeometryAabb`] statistic.
 /// An asymmetric predicate (e.g. a future contains) must recover which operand is the column
 /// itself instead of calling this.
-fn geometry_and_constant<'a>(
-    expr: &'a BoundExpression,
-    ctx: &StatsRewriteCtx<'_>,
-) -> VortexResult<Option<(&'a BoundExpression, &'a Scalar)>> {
+fn geometry_and_constant(expr: &BoundExpression) -> Option<(&BoundExpression, &Scalar)> {
     // The predicate is symmetric, so the column (scope root) and the constant may be on either
     // side.
     let (lhs, rhs) = (expr.child(0), expr.child(1));
@@ -63,16 +60,16 @@ fn geometry_and_constant<'a>(
     } else if rhs.is_root() {
         (rhs, lhs)
     } else {
-        return Ok(None);
+        return None;
     };
 
     // A `GeometryAabb` stat reference only binds for dtypes it supports; anything else (e.g. a
     // WKB column) must fall through to the scan.
-    if !is_native_geometry(&ctx.return_dtype(geom)?) {
-        return Ok(None);
+    if !is_native_geometry(geom.dtype()) {
+        return None;
     }
 
-    Ok(constant.as_opt::<Literal>().map(|scalar| (geom, scalar)))
+    constant.as_opt::<Literal>().map(|scalar| (geom, scalar))
 }
 
 /// The 2D bounding box of a constant geometry of any type, or `None` for one without an extent