From ae0252e020ea45994a684a0aef4e51190e4d814a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 23 Jul 2026 15:50:50 -0400 Subject: [PATCH 1/3] first pass Signed-off-by: Matt Katz --- vortex-datafusion/src/convert/decimal.rs | 163 +++++++++++ vortex-datafusion/src/convert/exprs.rs | 203 +++++++++++++- vortex-datafusion/src/convert/mod.rs | 1 + vortex-datafusion/src/persistent/opener.rs | 2 +- .../src/tests/decimal_arithmetic.rs | 264 ++++++++++++++++++ vortex-datafusion/src/tests/mod.rs | 1 + vortex-datafusion/src/v2/source.rs | 2 +- 7 files changed, 619 insertions(+), 17 deletions(-) create mode 100644 vortex-datafusion/src/convert/decimal.rs create mode 100644 vortex-datafusion/src/tests/decimal_arithmetic.rs diff --git a/vortex-datafusion/src/convert/decimal.rs b/vortex-datafusion/src/convert/decimal.rs new file mode 100644 index 00000000000..d3ac9bf2630 --- /dev/null +++ b/vortex-datafusion/src/convert/decimal.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Pushdown planning for DataFusion decimal arithmetic. +//! +//! DataFusion evaluates decimal `+`/`-` with the arrow-rs `numeric` kernels, which rescale both +//! operands to the wider scale and derive the result type with the Hive rule: +//! +//! ```text +//! s = max(s1, s2) +//! p = max(p1 - s1, p2 - s2) + s + 1 (saturated at the physical type's max precision) +//! ``` +//! +//! Vortex's decimal kernels instead require both operands to share one decimal dtype and widen +//! the result by a single carry digit, capped at Vortex's maximum precision. Vortex's own +//! coercion pass ([`coerce_expression`]) aligns mismatched operands to their least supertype — +//! the same `(p - 1, s)` dtype arrow rescales to internally — so a converted expression run +//! through that pass reproduces DataFusion's result type and values exactly, as long as arrow +//! does not saturate. +//! +//! When arrow saturates (e.g. two `Decimal128(38, 0)` inputs keep precision 38 while Vortex +//! would widen to 39), the result types and overflow behaviour diverge, so such expressions — +//! and decimal Mul/Div, which have no Vortex kernels — must not be pushed down and are instead +//! evaluated by DataFusion. + +use arrow_schema::DECIMAL32_MAX_PRECISION; +use arrow_schema::DECIMAL64_MAX_PRECISION; +use arrow_schema::DECIMAL128_MAX_PRECISION; +use arrow_schema::DECIMAL256_MAX_PRECISION; +use arrow_schema::DataType; +use datafusion_common::Result as DFResult; +use datafusion_common::exec_datafusion_err; +use datafusion_expr::Operator as DFOperator; +use vortex::dtype::DType; +use vortex::expr::Expression; +use vortex::expr::transform::coerce_expression; +use vortex::scalar_fn::fns::binary::Binary; + +/// Decide whether a DataFusion decimal arithmetic expression can be pushed down to Vortex. +/// +/// Returns `false` when the expression must stay in DataFusion: non-`+`/`-` operators (Vortex +/// has no decimal Mul/Div kernels yet), operands of different physical decimal widths, or a +/// result precision that arrow saturates at the physical type's ceiling. +pub(crate) fn decimal_arithmetic_is_pushable( + op: &DFOperator, + lhs: &DataType, + rhs: &DataType, +) -> bool { + if !matches!(op, DFOperator::Plus | DFOperator::Minus) { + return false; + } + + let (Some((p1, s1, lhs_max_precision)), Some((p2, s2, rhs_max_precision))) = + (decimal_parts(lhs), decimal_parts(rhs)) + else { + return false; + }; + // DataFusion coerces mixed physical widths before execution; reject them if they show up + // here rather than guessing at arrow's cross-width semantics. + if lhs_max_precision != rhs_max_precision { + return false; + } + + let scale = s1.max(s2); + let integer_digits = (i32::from(p1) - i32::from(s1)).max(i32::from(p2) - i32::from(s2)); + let precision = integer_digits + i32::from(scale) + 1; + // Arrow saturates the result precision at the ceiling; Vortex would widen instead. + precision <= i32::from(lhs_max_precision) +} + +/// Align decimal arithmetic operands in a converted expression by running Vortex's coercion +/// pass ([`coerce_expression`]) over it. +/// +/// The pass is applied only when the expression contains decimal arithmetic: for everything +/// else the converted operand dtypes already match (DataFusion coerces them), and skipping the +/// pass avoids inserting nullability-only casts into unrelated expressions. +pub(crate) fn coerce_decimal_arithmetic(expr: Expression, scope: &DType) -> DFResult { + if !contains_decimal_arithmetic(&expr, scope) { + return Ok(expr); + } + coerce_expression(expr, scope) + .map_err(|e| exec_datafusion_err!("Failed to coerce decimal arithmetic expression: {e}")) +} + +/// Whether any arithmetic [`Binary`] node in the tree has a decimal-typed child. A child whose +/// dtype cannot be resolved yet (e.g. arithmetic over still-unaligned operands) counts as +/// decimal, since only coercion can make it resolvable. +fn contains_decimal_arithmetic(expr: &Expression, scope: &DType) -> bool { + if expr.is::() + && expr.as_::().is_arithmetic() + && expr.children().iter().any(|child| { + child + .return_dtype(scope) + .map(|dtype| matches!(dtype, DType::Decimal(..))) + .unwrap_or(true) + }) + { + return true; + } + expr.children() + .iter() + .any(|child| contains_decimal_arithmetic(child, scope)) +} + +fn decimal_parts(dt: &DataType) -> Option<(u8, i8, u8)> { + match dt { + DataType::Decimal32(p, s) => Some((*p, *s, DECIMAL32_MAX_PRECISION)), + DataType::Decimal64(p, s) => Some((*p, *s, DECIMAL64_MAX_PRECISION)), + DataType::Decimal128(p, s) => Some((*p, *s, DECIMAL128_MAX_PRECISION)), + DataType::Decimal256(p, s) => Some((*p, *s, DECIMAL256_MAX_PRECISION)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::same_dtype(DataType::Decimal128(10, 2), DataType::Decimal128(10, 2), true)] + #[case::mixed_scale(DataType::Decimal128(10, 2), DataType::Decimal128(12, 4), true)] + #[case::mixed_precision(DataType::Decimal128(10, 2), DataType::Decimal128(20, 2), true)] + #[case::decimal128_ceiling(DataType::Decimal128(38, 0), DataType::Decimal128(38, 0), false)] + #[case::decimal128_ceiling_by_rescale( + DataType::Decimal128(38, 0), + DataType::Decimal128(38, 10), + false + )] + #[case::decimal256_below_ceiling( + DataType::Decimal256(38, 0), + DataType::Decimal256(38, 0), + true + )] + #[case::decimal256_ceiling(DataType::Decimal256(76, 0), DataType::Decimal256(76, 0), false)] + #[case::negative_scale(DataType::Decimal128(10, -2), DataType::Decimal128(10, -2), true)] + #[case::mixed_width(DataType::Decimal128(10, 2), DataType::Decimal256(10, 2), false)] + #[case::non_decimal(DataType::Decimal128(10, 2), DataType::Int64, false)] + fn test_decimal_arithmetic_is_pushable( + #[case] lhs: DataType, + #[case] rhs: DataType, + #[case] expected: bool, + ) { + assert_eq!( + decimal_arithmetic_is_pushable(&DFOperator::Plus, &lhs, &rhs), + expected + ); + assert_eq!( + decimal_arithmetic_is_pushable(&DFOperator::Minus, &lhs, &rhs), + expected + ); + } + + #[rstest] + #[case::mul(DFOperator::Multiply)] + #[case::div(DFOperator::Divide)] + #[case::modulo(DFOperator::Modulo)] + fn test_decimal_arithmetic_unsupported_ops(#[case] op: DFOperator) { + let dt = DataType::Decimal128(10, 2); + assert!(!decimal_arithmetic_is_pushable(&op, &dt, &dt)); + } +} diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 847a020b887..ee1e6b7ec6c 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -24,6 +24,7 @@ use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; use datafusion_physical_plan::expressions as df_expr; use itertools::Itertools; use vortex::VortexSessionDefault; +use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::expr::Expression; use vortex::expr::and_collect; @@ -49,6 +50,8 @@ use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; use crate::convert::FromDataFusion; +use crate::convert::decimal::coerce_decimal_arithmetic; +use crate::convert::decimal::decimal_arithmetic_is_pushable; /// Result of splitting a projection into Vortex expressions and leftover DataFusion projections. pub struct ProcessedProjection { @@ -57,13 +60,20 @@ pub struct ProcessedProjection { } /// Tries to convert the expressions into a vortex conjunction. Will return Ok(None) iff the input conjunction is empty. +/// +/// `scope` is the dtype the predicate will be evaluated against; converted expressions +/// containing decimal arithmetic are aligned to it with Vortex's coercion pass. pub(crate) fn make_vortex_predicate( expr_convertor: &dyn ExpressionConvertor, predicate: &[Arc], + scope: &DType, ) -> DFResult> { let exprs = predicate .iter() - .map(|e| expr_convertor.convert(e.as_ref())) + .map(|e| { + let expr = expr_convertor.convert(e.as_ref())?; + coerce_decimal_arithmetic(expr, scope) + }) .collect::>>()?; Ok(and_collect(exprs)) @@ -374,6 +384,14 @@ impl ExpressionConvertor for DefaultExpressionConvertor { let mut scan_projection = vec![]; let mut leftover_projection: Vec = vec![]; + // The dtype scope pushed-down projections are evaluated against, used to align any + // decimal arithmetic they contain. + let scope = self + .session + .arrow() + .from_arrow_schema(input_schema) + .map_err(|e| exec_datafusion_err!("Failed to convert input schema to dtype: {e}"))?; + for projection_expr in source_projection.iter() { let r = projection_expr.expr.apply(|node| { // We only pull column children of scalar functions that we can't push into the scan. @@ -390,12 +408,11 @@ impl ExpressionConvertor for DefaultExpressionConvertor { return Ok(TreeNodeRecursion::Stop); } - // DataFusion assumes different decimal types can be coerced. - // Vortex expects a perfect match so we don't push it down. + // Vortex reproduces DataFusion's decimal arithmetic exactly only for Add/Sub + // below the physical type's precision ceiling; keep everything else in + // DataFusion. See `convert::decimal`. if let Some(binary_expr) = node.downcast_ref::() - && binary_expr.op().is_numerical_operators() - && binary_expr.left().data_type(input_schema)?.is_decimal() - && binary_expr.right().data_type(input_schema)?.is_decimal() + && !decimal_arithmetic_can_be_pushed_down(binary_expr, input_schema) { scan_projection.extend( collect_columns(node) @@ -412,9 +429,10 @@ impl ExpressionConvertor for DefaultExpressionConvertor { // if we didn't stop early if matches!(r, TreeNodeRecursion::Continue) { + let converted = self.convert(projection_expr.expr.as_ref())?; scan_projection.push(( projection_expr.alias.clone(), - self.convert(projection_expr.expr.as_ref())?, + coerce_decimal_arithmetic(converted, &scope)?, )); leftover_projection.push(ProjectionExpr { expr: Arc::new(df_expr::Column::new_with_schema( @@ -507,7 +525,7 @@ fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> boo supported_data_types(&lit.value().data_type()) } else if let Some(cast_expr) = expr.downcast_ref::() { // CastExpr child must be an expression type that convert() can handle - is_convertible_expr(cast_expr.expr()) + is_convertible_expr(cast_expr.expr(), schema) } else if let Some(is_null) = expr.downcast_ref::() { can_be_pushed_down_impl(is_null.arg(), schema) } else if let Some(is_not_null) = expr.downcast_ref::() { @@ -531,15 +549,16 @@ fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> boo /// Checks if an expression type is one that convert() can handle. /// This is less restrictive than can_be_pushed_down since it only checks /// expression types, not data type support. -fn is_convertible_expr(expr: &Arc) -> bool { +fn is_convertible_expr(expr: &Arc, schema: &Schema) -> bool { // Expression types that convert() handles - expr.downcast_ref::().is_some() + expr.downcast_ref::() + .is_some_and(|binary| binary_tree_decimal_arithmetic_ok(binary, schema)) || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() || expr .downcast_ref::() - .is_some_and(|e| is_convertible_expr(e.expr())) + .is_some_and(|e| is_convertible_expr(e.expr(), schema)) || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() @@ -550,9 +569,41 @@ fn is_convertible_expr(expr: &Arc) -> bool { }) } +/// Walk a [`df_expr::BinaryExpr`] tree checking that every decimal arithmetic node can be +/// pushed down. Non-binary children are accepted structurally, matching +/// [`is_convertible_expr`]'s shallow contract. +fn binary_tree_decimal_arithmetic_ok(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { + decimal_arithmetic_can_be_pushed_down(binary, schema) + && [binary.left(), binary.right()].into_iter().all(|child| { + child + .downcast_ref::() + .is_none_or(|b| binary_tree_decimal_arithmetic_ok(b, schema)) + }) +} + +/// A numeric [`df_expr::BinaryExpr`] over decimals is pushable only when Vortex can reproduce +/// DataFusion's arithmetic exactly; see `convert::decimal`. Non-numeric operators and +/// non-decimal operands are unaffected. +fn decimal_arithmetic_can_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { + if !binary.op().is_numerical_operators() { + return true; + } + let (Ok(lhs), Ok(rhs)) = ( + binary.left().data_type(schema), + binary.right().data_type(schema), + ) else { + return false; + }; + if !lhs.is_decimal() && !rhs.is_decimal() { + return true; + } + decimal_arithmetic_is_pushable(binary.op(), &lhs, &rhs) +} + fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { let is_op_supported = try_operator_from_df(binary.op()).is_ok(); is_op_supported + && decimal_arithmetic_can_be_pushed_down(binary, schema) && can_be_pushed_down_impl(binary.left(), schema) && can_be_pushed_down_impl(binary.right(), schema) } @@ -653,7 +704,7 @@ fn can_array_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Sche data_type, DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) ) - }) && is_convertible_expr(input) + }) && is_convertible_expr(input, schema) } /// Returns the list argument of an `array_length` call if the call is a form we can rewrite to @@ -715,9 +766,33 @@ mod tests { DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), true, ), + Field::new("price", DataType::Decimal128(10, 2), true), + Field::new("cost", DataType::Decimal128(12, 4), true), + Field::new("big", DataType::Decimal128(38, 0), true), ]) } + /// The dtype scope for expressions over `test_schema`. + fn scope(schema: &Schema) -> DType { + VortexSession::default() + .arrow() + .from_arrow_schema(schema) + .unwrap() + } + + /// Build `left right` over columns of `test_schema`. + fn binary_col_expr( + left: (&str, usize), + op: DFOperator, + right: (&str, usize), + ) -> Arc { + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new(left.0, left.1)) as Arc, + op, + Arc::new(df_expr::Column::new(right.0, right.1)) as Arc, + )) + } + fn octet_length_expr(input: Arc, schema: &Schema) -> Arc { Arc::new( ScalarFunctionExpr::try_new( @@ -748,7 +823,7 @@ mod tests { #[test] fn test_make_vortex_predicate_empty() { let expr_convertor = DefaultExpressionConvertor::default(); - let result = make_vortex_predicate(&expr_convertor, &[]).unwrap(); + let result = make_vortex_predicate(&expr_convertor, &[], &scope(&Schema::empty())).unwrap(); assert!(result.is_none()); } @@ -756,7 +831,8 @@ mod tests { fn test_make_vortex_predicate_single() { let expr_convertor = DefaultExpressionConvertor::default(); let col_expr = Arc::new(df_expr::Column::new("test", 0)) as Arc; - let result = make_vortex_predicate(&expr_convertor, &[col_expr]).unwrap(); + let result = + make_vortex_predicate(&expr_convertor, &[col_expr], &scope(&Schema::empty())).unwrap(); assert!(result.is_some()); } @@ -765,7 +841,9 @@ mod tests { let expr_convertor = DefaultExpressionConvertor::default(); let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc; let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc; - let result = make_vortex_predicate(&expr_convertor, &[col1, col2]).unwrap(); + let result = + make_vortex_predicate(&expr_convertor, &[col1, col2], &scope(&Schema::empty())) + .unwrap(); assert!(result.is_some()); // Result should be an AND expression combining the two columns } @@ -1037,6 +1115,101 @@ mod tests { assert!(!can_be_pushed_down_impl(&binary_expr, &test_schema)); } + #[rstest] + #[case::add_same_dtype("price", 6, DFOperator::Plus, true)] + #[case::sub_same_dtype("price", 6, DFOperator::Minus, true)] + #[case::add_mixed_scale("cost", 7, DFOperator::Plus, true)] + #[case::mul_unsupported("price", 6, DFOperator::Multiply, false)] + #[case::div_unsupported("price", 6, DFOperator::Divide, false)] + #[case::add_at_precision_ceiling("big", 8, DFOperator::Plus, false)] + fn test_can_be_pushed_down_decimal_arithmetic( + test_schema: Schema, + #[case] rhs: &str, + #[case] rhs_idx: usize, + #[case] op: DFOperator, + #[case] expected: bool, + ) { + let binary = binary_col_expr(("price", 6), op, (rhs, rhs_idx)); + assert_eq!(can_be_pushed_down_impl(&binary, &test_schema), expected); + } + + /// A comparison whose operand contains un-pushable decimal arithmetic must not be pushed + /// down as a whole, so DataFusion keeps evaluating the filter itself. + #[rstest] + fn test_can_be_pushed_down_decimal_arithmetic_in_comparison(test_schema: Schema) { + let ceiling_sum = binary_col_expr(("big", 8), DFOperator::Plus, ("big", 8)); + let literal = Arc::new(df_expr::Literal::new(ScalarValue::Decimal128( + Some(100), + 38, + 0, + ))) as Arc; + let comparison = Arc::new(df_expr::BinaryExpr::new( + ceiling_sum, + DFOperator::Eq, + literal, + )) as Arc; + + assert!(!can_be_pushed_down_impl(&comparison, &test_schema)); + } + + /// Un-pushable decimal arithmetic hiding under a cast is also rejected. + #[rstest] + fn test_can_be_pushed_down_decimal_arithmetic_under_cast(test_schema: Schema) { + let product = binary_col_expr(("price", 6), DFOperator::Multiply, ("price", 6)); + let cast = Arc::new(df_expr::CastExpr::new( + product, + DataType::Decimal128(21, 4), + None, + )) as Arc; + + assert!(!can_be_pushed_down_impl(&cast, &test_schema)); + } + + #[rstest] + fn test_expr_from_df_decimal_add_same_dtype_has_no_casts(test_schema: Schema) { + let binary = binary_col_expr(("price", 6), DFOperator::Plus, ("price", 6)); + + let result = make_vortex_predicate( + &DefaultExpressionConvertor::default(), + &[binary], + &scope(&test_schema), + ) + .unwrap() + .unwrap(); + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.binary(+) + ├── lhs: vortex.get_item(price) + │ └── input: vortex.root() + └── rhs: vortex.get_item(price) + └── input: vortex.root() + "); + } + + /// Mixed decimal operand dtypes are aligned by the coercion pass to their least supertype, + /// here `decimal(12, 4)` for `Decimal128(10, 2) + Decimal128(12, 4)`. + #[rstest] + fn test_expr_from_df_decimal_add_mixed_scale_aligns_operands(test_schema: Schema) { + let binary = binary_col_expr(("price", 6), DFOperator::Plus, ("cost", 7)); + + let result = make_vortex_predicate( + &DefaultExpressionConvertor::default(), + &[binary], + &scope(&test_schema), + ) + .unwrap() + .unwrap(); + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.binary(+) + ├── lhs: vortex.cast(decimal(12,4)?) + │ └── input: vortex.get_item(price) + │ └── input: vortex.root() + └── rhs: vortex.get_item(cost) + └── input: vortex.root() + "); + } + #[rstest] fn test_can_be_pushed_down_like_supported(test_schema: Schema) { let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 6a6fb8aef08..84430643b05 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -11,6 +11,7 @@ use vortex::error::VortexResult; +pub(crate) mod decimal; pub(crate) mod exprs; mod scalars; pub(crate) mod schema; diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index d0ffb472ebe..251ad0d69b2 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -416,7 +416,7 @@ impl FileOpener for VortexOpener { ))); } - make_vortex_predicate(expr_convertor.as_ref(), &pushed).transpose() + make_vortex_predicate(expr_convertor.as_ref(), &pushed, vxf.dtype()).transpose() }) .transpose()?; diff --git a/vortex-datafusion/src/tests/decimal_arithmetic.rs b/vortex-datafusion/src/tests/decimal_arithmetic.rs new file mode 100644 index 00000000000..11200351bc3 --- /dev/null +++ b/vortex-datafusion/src/tests/decimal_arithmetic.rs @@ -0,0 +1,264 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! End-to-end tests for decimal arithmetic in filters and projections. +//! +//! Decimal Add/Sub below the physical type's precision ceiling is pushed down into the Vortex +//! scan with operand-aligning casts (see `convert::decimal`). Decimal Mul/Div and arithmetic at +//! the precision ceiling stay in DataFusion, so those queries must still work — just without +//! pushdown. + +use std::sync::Arc; + +use datafusion::arrow::array::Decimal128Array; +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion_common::assert_batches_eq; +use rstest::rstest; + +use crate::common_tests::TestSessionContext; + +/// Schema: {a: Decimal128(9, 2), b: Decimal128(9, 2), c: Decimal128(9, 4), big1/big2: +/// Decimal128(38, 0)} +fn make_decimal_batch() -> RecordBatch { + // a = [10.50, 20.00, 5.25], b = [19.50, 10.00, 4.75] + let col_a = Decimal128Array::from(vec![1050i128, 2000, 525]) + .with_precision_and_scale(9, 2) + .unwrap(); + let col_b = Decimal128Array::from(vec![1950i128, 1000, 475]) + .with_precision_and_scale(9, 2) + .unwrap(); + // c = [0.0001, 2.5000, 1.0000] + let col_c = Decimal128Array::from(vec![1i128, 25000, 10000]) + .with_precision_and_scale(9, 4) + .unwrap(); + let big1 = Decimal128Array::from(vec![1i128, 2, 3]) + .with_precision_and_scale(38, 0) + .unwrap(); + let big2 = Decimal128Array::from(vec![10i128, 20, 30]) + .with_precision_and_scale(38, 0) + .unwrap(); + + RecordBatch::try_from_iter(vec![ + ("a", Arc::new(col_a) as _), + ("b", Arc::new(col_b) as _), + ("c", Arc::new(col_c) as _), + ("big1", Arc::new(big1) as _), + ("big2", Arc::new(big2) as _), + ]) + .unwrap() +} + +async fn decimal_table(ctx: &TestSessionContext) -> anyhow::Result<()> { + let batch = make_decimal_batch(); + ctx.write_arrow_batch("files/decimals.vortex", &batch) + .await?; + let provider = ctx + .table_provider("tbl", "/files/", batch.schema().as_ref().clone()) + .await?; + ctx.session.register_table("tbl", provider)?; + Ok(()) +} + +/// Collect the physical plan of `sql` as a string. +async fn explain(ctx: &TestSessionContext, sql: &str) -> anyhow::Result { + let batches = ctx + .session + .sql(&format!("EXPLAIN {sql}")) + .await? + .collect() + .await?; + Ok(pretty_format_batches(&batches)?.to_string()) +} + +#[tokio::test] +async fn test_decimal_add_filter_is_pushed_down() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + decimal_table(&ctx).await?; + + let sql = "SELECT a FROM tbl WHERE a + b = 30.00 ORDER BY a"; + let result = ctx.session.sql(sql).await?.collect().await?; + + assert_batches_eq!( + [ + "+-------+", + "| a |", + "+-------+", + "| 10.50 |", + "| 20.00 |", + "+-------+", + ], + &result + ); + + // The filter runs inside the Vortex scan, so no FilterExec remains in the plan. + let plan = explain(&ctx, sql).await?; + assert!(!plan.contains("FilterExec"), "expected pushdown:\n{plan}"); + Ok(()) +} + +#[tokio::test] +async fn test_decimal_sub_filter_is_pushed_down() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + decimal_table(&ctx).await?; + + let sql = "SELECT a FROM tbl WHERE a - b > 0 ORDER BY a"; + let result = ctx.session.sql(sql).await?.collect().await?; + + assert_batches_eq!( + [ + "+-------+", + "| a |", + "+-------+", + "| 5.25 |", + "| 20.00 |", + "+-------+", + ], + &result + ); + + let plan = explain(&ctx, sql).await?; + assert!(!plan.contains("FilterExec"), "expected pushdown:\n{plan}"); + Ok(()) +} + +/// Operands with different scales are rescaled inside the pushed-down expression, matching +/// DataFusion's arrow kernel semantics exactly. +#[tokio::test] +async fn test_decimal_mixed_scale_filter_is_pushed_down() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + decimal_table(&ctx).await?; + + let sql = "SELECT a FROM tbl WHERE a + c > 22.0"; + let result = ctx.session.sql(sql).await?.collect().await?; + + assert_batches_eq!( + [ + "+-------+", + "| a |", + "+-------+", + "| 20.00 |", + "+-------+", + ], + &result + ); + + let plan = explain(&ctx, sql).await?; + assert!(!plan.contains("FilterExec"), "expected pushdown:\n{plan}"); + Ok(()) +} + +/// At the Decimal128 precision ceiling arrow saturates the result precision while Vortex would +/// widen it, so the filter must stay in DataFusion — and still produce correct results. +#[tokio::test] +async fn test_decimal_add_at_precision_ceiling_falls_back() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + decimal_table(&ctx).await?; + + let sql = "SELECT big1 FROM tbl WHERE big1 + big2 = 33"; + let result = ctx.session.sql(sql).await?.collect().await?; + + assert_batches_eq!( + ["+------+", "| big1 |", "+------+", "| 3 |", "+------+"], + &result + ); + + let plan = explain(&ctx, sql).await?; + assert!( + plan.contains("FilterExec"), + "expected DataFusion fallback:\n{plan}" + ); + Ok(()) +} + +/// Vortex has no decimal Mul/Div kernels; the filter stays in DataFusion instead of failing the +/// query at scan time. +#[tokio::test] +async fn test_decimal_mul_filter_falls_back() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + decimal_table(&ctx).await?; + + let sql = "SELECT a FROM tbl WHERE a * b > 100.0 ORDER BY a"; + let result = ctx.session.sql(sql).await?.collect().await?; + + assert_batches_eq!( + [ + "+-------+", + "| a |", + "+-------+", + "| 10.50 |", + "| 20.00 |", + "+-------+", + ], + &result + ); + + let plan = explain(&ctx, sql).await?; + assert!( + plan.contains("FilterExec"), + "expected DataFusion fallback:\n{plan}" + ); + Ok(()) +} + +/// Decimal Add/Sub in projections, with and without projection pushdown. +#[rstest] +#[tokio::test] +async fn test_decimal_add_projection( + #[values(false, true)] projection_pushdown: bool, +) -> anyhow::Result<()> { + let ctx = TestSessionContext::new(projection_pushdown); + decimal_table(&ctx).await?; + + let result = ctx + .session + .sql("SELECT a + b AS total, a - c AS diff FROM tbl ORDER BY total") + .await? + .collect() + .await?; + + assert_batches_eq!( + [ + "+-------+---------+", + "| total | diff |", + "+-------+---------+", + "| 10.00 | 4.2500 |", + "| 30.00 | 10.4999 |", + "| 30.00 | 17.5000 |", + "+-------+---------+", + ], + &result + ); + Ok(()) +} + +/// Decimal Mul in a projection must fall back to DataFusion evaluation. +#[rstest] +#[tokio::test] +async fn test_decimal_mul_projection_falls_back( + #[values(false, true)] projection_pushdown: bool, +) -> anyhow::Result<()> { + let ctx = TestSessionContext::new(projection_pushdown); + decimal_table(&ctx).await?; + + let result = ctx + .session + .sql("SELECT a * b AS product FROM tbl ORDER BY product") + .await? + .collect() + .await?; + + assert_batches_eq!( + [ + "+----------+", + "| product |", + "+----------+", + "| 24.9375 |", + "| 200.0000 |", + "| 204.7500 |", + "+----------+", + ], + &result + ); + Ok(()) +} diff --git a/vortex-datafusion/src/tests/mod.rs b/vortex-datafusion/src/tests/mod.rs index e3421f63c67..f4d3f5ee9f9 100644 --- a/vortex-datafusion/src/tests/mod.rs +++ b/vortex-datafusion/src/tests/mod.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod decimal_arithmetic; mod nested_projection; mod schema_evolution; diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 325f8eae92f..f3656dc636a 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -635,7 +635,7 @@ impl DataSource for VortexDataSource { .collect(); // Convert to Vortex conjunction. - let vortex_pred = make_vortex_predicate(&convertor, &pushable)?; + let vortex_pred = make_vortex_predicate(&convertor, &pushable, self.data_source.dtype())?; // Combine with existing filter. let new_filter = match (&self.filter, vortex_pred) { From 6e2fae688ec533b59edf540090239d3f34fcf359 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 23 Jul 2026 16:03:43 -0400 Subject: [PATCH 2/3] slt: decimal arithmetic pushdown Signed-off-by: Matt Katz --- .../decimal_arithmetic_pushdown.slt | 174 ++++++++++++++++++ 1 file changed, 174 insertions(+) create mode 100644 vortex-sqllogictest/slt/datafusion/decimal_arithmetic_pushdown.slt diff --git a/vortex-sqllogictest/slt/datafusion/decimal_arithmetic_pushdown.slt b/vortex-sqllogictest/slt/datafusion/decimal_arithmetic_pushdown.slt new file mode 100644 index 00000000000..1a159cb0f1e --- /dev/null +++ b/vortex-sqllogictest/slt/datafusion/decimal_arithmetic_pushdown.slt @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +# Decimal Add/Sub in filters and projections is pushed into the Vortex scan (aligned through +# Vortex's coercion pass). Decimal Mul/Div (no Vortex kernels) and arithmetic at the physical +# type's precision ceiling (where arrow saturates the result precision while Vortex widens it) +# stay in DataFusion. + +include ../setup.slt.no + +statement ok +CREATE EXTERNAL TABLE decimal_arith ( + id INT NOT NULL, + a DECIMAL(9, 2) NOT NULL, + b DECIMAL(9, 2) NOT NULL, + c DECIMAL(9, 4) NOT NULL, + big1 DECIMAL(38, 0) NOT NULL, + big2 DECIMAL(38, 0) NOT NULL +) +STORED AS vortex +LOCATION '${WORK_DIR}/decimal_arith/'; + +statement ok +INSERT INTO decimal_arith VALUES + (1, 10.50, 19.50, 0.0001, 1, 10), + (2, 20.00, 10.00, 2.5000, 2, 20), + (3, 5.25, 4.75, 1.0000, 3, 30); + +# Decimal Add over same-typed operands is pushed into the scan: no FilterExec remains. + +query TT +EXPLAIN SELECT a FROM decimal_arith WHERE a + b = 30.00 ORDER BY a; +---- +logical_plan +01)Sort: decimal_arith.a ASC NULLS LAST +02)--Projection: decimal_arith.a +03)----Filter: decimal_arith.a + decimal_arith.b = Decimal128(Some(3000),10,2) +04)------TableScan: decimal_arith projection=[a, b], partial_filters=[decimal_arith.a + decimal_arith.b = Decimal128(Some(3000),10,2)] +physical_plan +01)SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={}, projection=[a], file_type=vortex, predicate: a@1 + b@2 = Some(3000),10,2 + +query R +SELECT a FROM decimal_arith WHERE a + b = 30.00 ORDER BY a; +---- +10.5 +20 + +# Decimal Sub is pushed down too. + +query TT +EXPLAIN SELECT a FROM decimal_arith WHERE a - b > 0 ORDER BY a; +---- +logical_plan +01)Sort: decimal_arith.a ASC NULLS LAST +02)--Projection: decimal_arith.a +03)----Filter: decimal_arith.a - decimal_arith.b > Decimal128(Some(0),10,2) +04)------TableScan: decimal_arith projection=[a, b], partial_filters=[decimal_arith.a - decimal_arith.b > Decimal128(Some(0),10,2)] +physical_plan +01)SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={}, projection=[a], file_type=vortex, predicate: a@1 - b@2 > Some(0),10,2 + +query R +SELECT a FROM decimal_arith WHERE a - b > 0 ORDER BY a; +---- +5.25 +20 + +# Operands with different scales are rescaled inside the pushed-down expression. + +query TT +EXPLAIN SELECT a FROM decimal_arith WHERE a + c > 22.0; +---- +logical_plan +01)Projection: decimal_arith.a +02)--Filter: decimal_arith.a + decimal_arith.c > Decimal128(Some(220000),12,4) +03)----TableScan: decimal_arith projection=[a, c], partial_filters=[decimal_arith.a + decimal_arith.c > Decimal128(Some(220000),12,4)] +physical_plan DataSourceExec: file_groups={}, projection=[a], file_type=vortex, predicate: a@1 + c@3 > Some(220000),12,4 + +query R +SELECT a FROM decimal_arith WHERE a + c > 22.0; +---- +20 + +# At the Decimal128 precision ceiling (38) the filter stays in DataFusion (FilterExec remains) +# and still produces correct results. + +query TT +EXPLAIN SELECT big1 FROM decimal_arith WHERE big1 + big2 = 33; +---- +logical_plan +01)Projection: decimal_arith.big1 +02)--Filter: decimal_arith.big1 + decimal_arith.big2 = Decimal128(Some(33),38,0) +03)----TableScan: decimal_arith projection=[big1, big2], partial_filters=[decimal_arith.big1 + decimal_arith.big2 = Decimal128(Some(33),38,0)] +physical_plan +01)FilterExec: big1@0 + big2@1 = Some(33),38,0, projection=[big1@0] +02)--DataSourceExec: file_groups={}, projection=[big1, big2], file_type=vortex + +query R +SELECT big1 FROM decimal_arith WHERE big1 + big2 = 33; +---- +3 + +# Decimal Mul has no Vortex kernel; the filter stays in DataFusion. + +query TT +EXPLAIN SELECT a FROM decimal_arith WHERE a * b > 100.0 ORDER BY a; +---- +logical_plan +01)Sort: decimal_arith.a ASC NULLS LAST +02)--Projection: decimal_arith.a +03)----Filter: decimal_arith.a * decimal_arith.b > Decimal128(Some(1000000),19,4) +04)------TableScan: decimal_arith projection=[a, b], partial_filters=[decimal_arith.a * decimal_arith.b > Decimal128(Some(1000000),19,4)] +physical_plan +01)SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--FilterExec: a@0 * b@1 > Some(1000000),19,4, projection=[a@0] +03)----DataSourceExec: file_groups={}, projection=[a, b], file_type=vortex + +query R +SELECT a FROM decimal_arith WHERE a * b > 100.0 ORDER BY a; +---- +10.5 +20 + +# Decimal Add/Sub in projections is pushed into the scan when projection pushdown is enabled. +# Note the Vortex source splits pushed vs leftover projections internally (in the file opener), +# so DataFusion's plan shows the full projection on the DataSourceExec either way; the queries +# below are the regression tests that both variants evaluate correctly. + +statement ok +SET vortex.projection_pushdown = true; + +query TT +EXPLAIN SELECT a + b AS total, a - c AS diff FROM decimal_arith ORDER BY total; +---- +logical_plan +01)Sort: total ASC NULLS LAST +02)--Projection: decimal_arith.a + decimal_arith.b AS total, decimal_arith.a - decimal_arith.c AS diff +03)----TableScan: decimal_arith projection=[a, b, c] +physical_plan +01)SortExec: expr=[total@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={}, projection=[a@1 + b@2 as total, a@1 - c@3 as diff], file_type=vortex + +query RR +SELECT a + b AS total, a - c AS diff FROM decimal_arith ORDER BY total; +---- +10 4.25 +30 10.4999 +30 17.5 + +# Decimal Mul in a projection is kept out of the Vortex scan and evaluated with DataFusion's +# expression machinery (hosted in the source's post-scan projector). If the split_projection +# guard ever leaked decimal Mul into the Vortex scan, this query would error. + +query TT +EXPLAIN SELECT a * b AS product FROM decimal_arith ORDER BY product; +---- +logical_plan +01)Sort: product ASC NULLS LAST +02)--Projection: decimal_arith.a * decimal_arith.b AS product +03)----TableScan: decimal_arith projection=[a, b] +physical_plan +01)SortExec: expr=[product@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={}, projection=[a@1 * b@2 as product], file_type=vortex + +query R +SELECT a * b AS product FROM decimal_arith ORDER BY product; +---- +24.9375 +200 +204.75 + +statement ok +SET vortex.projection_pushdown = false; From b5d5d2ffc27b578d4986ac97406acc06e7b9b0a4 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 24 Jul 2026 11:24:40 -0400 Subject: [PATCH 3/3] fixes Signed-off-by: Matt Katz --- vortex-datafusion/src/convert/exprs.rs | 29 +++++----- .../src/tests/decimal_arithmetic.rs | 58 +++++++++++++++++++ 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index ee1e6b7ec6c..d3312274b3d 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -399,7 +399,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { && !can_scalar_fn_be_pushed_down(scalar_fn_expr, input_schema) { scan_projection.extend( - collect_columns(node) + collect_columns(&projection_expr.expr) .into_iter() .map(|c| (c.name().to_string(), get_item(c.name(), root()))), ); @@ -415,7 +415,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { && !decimal_arithmetic_can_be_pushed_down(binary_expr, input_schema) { scan_projection.extend( - collect_columns(node) + collect_columns(&projection_expr.expr) .into_iter() .map(|c| (c.name().to_string(), get_item(c.name(), root()))), ); @@ -550,9 +550,12 @@ fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> boo /// This is less restrictive than can_be_pushed_down since it only checks /// expression types, not data type support. fn is_convertible_expr(expr: &Arc, schema: &Schema) -> bool { + if !expression_tree_decimal_arithmetic_ok(expr, schema) { + return false; + } + // Expression types that convert() handles - expr.downcast_ref::() - .is_some_and(|binary| binary_tree_decimal_arithmetic_ok(binary, schema)) + expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() @@ -569,16 +572,14 @@ fn is_convertible_expr(expr: &Arc, schema: &Schema) -> bool { }) } -/// Walk a [`df_expr::BinaryExpr`] tree checking that every decimal arithmetic node can be -/// pushed down. Non-binary children are accepted structurally, matching -/// [`is_convertible_expr`]'s shallow contract. -fn binary_tree_decimal_arithmetic_ok(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { - decimal_arithmetic_can_be_pushed_down(binary, schema) - && [binary.left(), binary.right()].into_iter().all(|child| { - child - .downcast_ref::() - .is_none_or(|b| binary_tree_decimal_arithmetic_ok(b, schema)) - }) +/// Checks that every decimal arithmetic node in an expression tree can be pushed down. +fn expression_tree_decimal_arithmetic_ok(expr: &Arc, schema: &Schema) -> bool { + expr.downcast_ref::() + .is_none_or(|binary| decimal_arithmetic_can_be_pushed_down(binary, schema)) + && expr + .children() + .into_iter() + .all(|child| expression_tree_decimal_arithmetic_ok(child, schema)) } /// A numeric [`df_expr::BinaryExpr`] over decimals is pushable only when Vortex can reproduce diff --git a/vortex-datafusion/src/tests/decimal_arithmetic.rs b/vortex-datafusion/src/tests/decimal_arithmetic.rs index 11200351bc3..97082c7b92b 100644 --- a/vortex-datafusion/src/tests/decimal_arithmetic.rs +++ b/vortex-datafusion/src/tests/decimal_arithmetic.rs @@ -201,6 +201,35 @@ async fn test_decimal_mul_filter_falls_back() -> anyhow::Result<()> { Ok(()) } +/// Unsupported decimal arithmetic nested below casts must keep the complete filter in +/// DataFusion instead of being accepted and failing during Vortex expression coercion. +#[tokio::test] +async fn test_decimal_mul_under_nested_casts_falls_back() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + decimal_table(&ctx).await?; + + let sql = "SELECT a FROM tbl WHERE CAST(a + CAST(a * b AS DECIMAL(24, 4)) AS DOUBLE) = 215.25"; + let result = ctx.session.sql(sql).await?.collect().await?; + + assert_batches_eq!( + [ + "+-------+", + "| a |", + "+-------+", + "| 10.50 |", + "+-------+" + ], + &result + ); + + let plan = explain(&ctx, sql).await?; + assert!( + plan.contains("FilterExec"), + "expected DataFusion fallback:\n{plan}" + ); + Ok(()) +} + /// Decimal Add/Sub in projections, with and without projection pushdown. #[rstest] #[tokio::test] @@ -262,3 +291,32 @@ async fn test_decimal_mul_projection_falls_back( ); Ok(()) } + +/// Falling back on an unsupported subtree must retain columns used by the rest of the +/// projection expression. +#[tokio::test] +async fn test_decimal_mul_nested_in_add_projection_falls_back() -> anyhow::Result<()> { + let ctx = TestSessionContext::new(true); + decimal_table(&ctx).await?; + + let result = ctx + .session + .sql("SELECT a + b * c AS mixed FROM tbl ORDER BY mixed") + .await? + .collect() + .await?; + + assert_batches_eq!( + [ + "+-----------+", + "| mixed |", + "+-----------+", + "| 10.000000 |", + "| 10.501950 |", + "| 45.000000 |", + "+-----------+", + ], + &result + ); + Ok(()) +}