diff --git a/vortex-array/src/stats/rewrite.rs b/vortex-array/src/stats/rewrite.rs index ddf74ee5dab..54b4182d5ce 100644 --- a/vortex-array/src/stats/rewrite.rs +++ b/vortex-array/src/stats/rewrite.rs @@ -41,8 +41,8 @@ pub type StatsRewriteRuleRef = Arc; /// `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 { - Ok(expr.dtype().clone()) - } - /// Rewrite `expr` into a stats-backed falsifier. pub fn falsify(&self, expr: &BoundExpression) -> VortexResult> { 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}", diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index 3cb5fdb06df..7ce88ad11c7 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/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 { - 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 { - non_nan_check(ctx, expr, |expr| { + fn check(_ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult { + 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 { @@ -597,14 +594,14 @@ fn non_nan_check( } if expr.is::() { - 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)) } diff --git a/vortex-spatial/src/prune/distance.rs b/vortex-spatial/src/prune/distance.rs index 9231b33875e..3a985030ecc 100644 --- a/vortex-spatial/src/prune/distance.rs +++ b/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 { diff --git a/vortex-spatial/src/prune/intersects.rs b/vortex-spatial/src/prune/intersects.rs index 74103003b2f..e98a9c7a670 100644 --- a/vortex-spatial/src/prune/intersects.rs +++ b/vortex-spatial/src/prune/intersects.rs @@ -37,7 +37,7 @@ impl StatsRewriteRule for SpatialIntersectsPrune { expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, ) -> VortexResult> { - 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 { diff --git a/vortex-spatial/src/prune/mod.rs b/vortex-spatial/src/prune/mod.rs index 32c2058b384..57d7f7b52ce 100644 --- a/vortex-spatial/src/prune/mod.rs +++ b/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> { +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::().map(|scalar| (geom, scalar))) + constant.as_opt::().map(|scalar| (geom, scalar)) } /// The 2D bounding box of a constant geometry of any type, or `None` for one without an extent