From 45a9e162d53f44bb3a5428af1d2d090106c5be40 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 24 Jul 2026 17:13:57 -0400 Subject: [PATCH] feat(expr): fold cast of a literal to a literal in Cast simplify Implement Cast::simplify_untyped so cast(lit(x) as T) folds to lit(x cast to T) at plan time. A failing scalar cast (e.g. null to a non-nullable dtype) leaves the expression unchanged so the error still surfaces at execution time. Prune rules pattern-match bare Literal nodes and silently decline cast-wrapped ones, so filters pushed with a cast around the literal (e.g. DuckDB decimal literals bound as cast(dec as f64)) scan every chunk. Since ScanBuilder already runs optimize_recursive on every pushed filter, this one fold fixes all producers at once. Signed-off-by: Nemo Yu --- vortex-array/src/expr/optimize.rs | 33 ++++++++++ vortex-array/src/scalar_fn/fns/cast/mod.rs | 77 ++++++++++++++++++++++ 2 files changed, 110 insertions(+) diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index 27959a96070..ea1fc796d2e 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -315,16 +315,21 @@ impl ReduceCtx for ExpressionReduceCtx { #[cfg(test)] mod tests { use vortex_error::VortexResult; + use vortex_error::vortex_err; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::expr::cast; use crate::expr::eq; use crate::expr::get_item; use crate::expr::lit; + use crate::expr::lt_eq; use crate::expr::or; use crate::expr::root; + use crate::scalar::Scalar; + use crate::scalar_fn::fns::literal::Literal; #[test] fn optimize_or_chain_correctness() -> VortexResult<()> { @@ -347,4 +352,32 @@ mod tests { assert!(s.contains("2i32") || s.contains('2'), "expected 2 in {s}"); Ok(()) } + + #[test] + fn optimize_folds_cast_of_literal_in_comparison() -> VortexResult<()> { + let expr = lt_eq( + get_item("x", root()), + cast( + lit(3i32), + DType::Primitive(PType::F64, Nullability::NonNullable), + ), + ); + let scope = DType::Struct( + StructFields::new( + ["x"].into(), + vec![DType::Primitive(PType::F64, Nullability::NonNullable)], + ), + Nullability::NonNullable, + ); + let optimized = expr.optimize_recursive(&scope)?; + + // Prune rules pattern-match a bare Literal on the comparison RHS; a cast wrapper + // silently disables pruning. + let rhs = optimized + .child(1) + .as_opt::() + .ok_or_else(|| vortex_err!("expected a bare literal RHS, got {optimized}"))?; + assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable)); + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 394adcc9924..859e76b4932 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -42,6 +42,7 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ReduceNodeRef; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::literal::Literal; /// A cast expression that converts values to a target data type. #[derive(Clone)] @@ -149,6 +150,19 @@ impl ScalarFnVTable for Cast { Ok(None) } + fn simplify_untyped( + &self, + target_dtype: &DType, + expr: &Expression, + ) -> VortexResult> { + let Some(scalar) = expr.child(0).as_opt::() else { + return Ok(None); + }; + // A failing cast (e.g. null to a non-nullable dtype) is left in place so the error + // surfaces at execution time rather than during optimization. + Ok(scalar.cast(target_dtype).ok().map(lit)) + } + fn validity(&self, dtype: &DType, expression: &Expression) -> VortexResult> { Ok(Some(if dtype.is_nullable() { expression.child(0).validity()? @@ -207,17 +221,25 @@ fn cast_constant(array: ArrayView, dtype: &DType) -> VortexResult VortexResult<()> { + let expr = cast( + lit(3i32), + DType::Primitive(PType::F64, Nullability::NonNullable), + ); + let optimized = expr.optimize(&test_harness::struct_dtype())?; + + let scalar = optimized + .as_opt::() + .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?; + assert_eq!(scalar, &Scalar::primitive(3.0f64, Nullability::NonNullable)); + Ok(()) + } + + #[test] + fn simplify_folds_cast_of_decimal_literal() -> VortexResult<()> { + let decimal = Scalar::decimal( + DecimalValue::I128(319), + DecimalDType::new(3, 2), + Nullability::NonNullable, + ); + let expr = cast( + lit(decimal), + DType::Primitive(PType::F64, Nullability::NonNullable), + ); + let optimized = expr.optimize(&test_harness::struct_dtype())?; + + let scalar = optimized + .as_opt::() + .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?; + assert_eq!( + scalar, + &Scalar::primitive(3.19f64, Nullability::NonNullable) + ); + Ok(()) + } + + #[test] + fn simplify_leaves_failing_cast_unchanged() -> VortexResult<()> { + let target = DType::Primitive(PType::F64, Nullability::NonNullable); + let expr = cast( + lit(Scalar::null(DType::Primitive( + PType::I32, + Nullability::Nullable, + ))), + target.clone(), + ); + let optimized = expr.optimize(&test_harness::struct_dtype())?; + + assert!(optimized.as_opt::().is_none()); + assert_eq!(optimized.as_opt::(), Some(&target)); + Ok(()) + } + #[test] fn test_display() { let expr = cast(