From 60e5b763f2ef2a5f636ce825ed41089fbb64e74a Mon Sep 17 00:00:00 2001 From: Daipayan Date: Tue, 8 Sep 2026 17:17:37 +0100 Subject: [PATCH 1/2] fix(substrait): reject incompatible decimal output types --- .../consumer/expr/scalar_function.rs | 269 +++++++++++++++++- .../tests/cases/consumer_integration.rs | 219 ++++---------- .../tests/cases/roundtrip_logical_plan.rs | 86 ++++++ 3 files changed, 401 insertions(+), 173 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs index 13f86997f31ca..a65be3052adbc 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs @@ -15,15 +15,23 @@ // specific language governing permissions and limitations // under the License. -use crate::logical_plan::consumer::{SubstraitConsumer, from_substrait_func_args}; +use crate::logical_plan::consumer::{ + SubstraitConsumer, from_substrait_func_args, from_substrait_type_without_names, +}; +use datafusion::arrow::datatypes::{ + DataType, Decimal128Type, Decimal256Type, validate_decimal_precision_and_scale, +}; use datafusion::common::Result; use datafusion::common::{ DFSchema, DataFusionError, ScalarValue, not_impl_err, plan_err, substrait_err, }; use datafusion::execution::FunctionRegistry; -use datafusion::logical_expr::{Between, BinaryExpr, Expr, Like, Operator, expr}; +use datafusion::logical_expr::{ + Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, expr, +}; use std::vec::Drain; use substrait::proto::expression::ScalarFunction; +use substrait::proto::r#type::Kind; pub async fn from_scalar_function( consumer: &impl SubstraitConsumer, @@ -86,7 +94,18 @@ pub async fn from_scalar_function( ); } // In those cases we build a balanced tree of BinaryExprs - arg_list_to_binary_op_tree(op, args) + let expr = arg_list_to_binary_op_tree(op, args)?; + if matches!( + op, + Operator::Plus + | Operator::Minus + | Operator::Multiply + | Operator::Divide + | Operator::Modulo + ) { + validate_decimal_output_type(consumer, f, fn_signature, &expr, input_schema)?; + } + Ok(expr) } else if let Some(builder) = BuiltinExprBuilder::try_from_name(fn_name) { builder.build(consumer, f, args) } else { @@ -94,6 +113,57 @@ pub async fn from_scalar_function( } } +/// Native decimal arithmetic can derive different precision and scale from the +/// referenced Substrait function. Reject mismatches rather than silently changing +/// the result type. An output cast cannot recover lost digits or prevent an +/// intermediate arithmetic overflow. +fn validate_decimal_output_type( + consumer: &impl SubstraitConsumer, + f: &ScalarFunction, + fn_signature: &str, + expr: &Expr, + input_schema: &DFSchema, +) -> Result<()> { + // Preserve compatibility with existing plans that omit the output type. + let Some(output_type) = &f.output_type else { + return Ok(()); + }; + let derived_type = expr.get_type(input_schema)?; + if !derived_type.is_decimal() && !matches!(output_type.kind, Some(Kind::Decimal(_))) { + return Ok(()); + } + + // The type decoder narrows these protobuf integers with `as`. Check them + // first so an invalid declaration cannot wrap into a matching type. + if let Some(Kind::Decimal(decimal)) = &output_type.kind + && (u8::try_from(decimal.precision).is_err() + || i8::try_from(decimal.scale).is_err()) + { + return substrait_err!( + "Invalid decimal output type for {fn_signature}: precision {}, scale {}", + decimal.precision, + decimal.scale + ); + } + let declared_type = from_substrait_type_without_names(consumer, output_type)?; + match &declared_type { + DataType::Decimal128(p, s) => { + validate_decimal_precision_and_scale::(*p, *s)?; + } + DataType::Decimal256(p, s) => { + validate_decimal_precision_and_scale::(*p, *s)?; + } + _ => {} + } + if declared_type != derived_type { + return substrait_err!( + "Decimal return type mismatch for {fn_signature} (function reference {}): declared {declared_type:?}, but native expression {expr} derives {derived_type:?}; this conversion is unsupported", + f.function_reference + ); + } + Ok(()) +} + pub fn substrait_fun_name(name: &str) -> &str { (match name.rsplit_once(':') { // Since 0.32.0, Substrait requires the function names to be in a compound format @@ -378,15 +448,206 @@ mod tests { use crate::extensions::Extensions; use crate::logical_plan::consumer::tests::TEST_SESSION_STATE; use crate::logical_plan::consumer::{DefaultSubstraitConsumer, SubstraitConsumer}; + use crate::logical_plan::producer::{ + DefaultSubstraitProducer, substrait_field_ref, to_substrait_type, + }; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::{DFSchema, Result, ScalarValue}; - use datafusion::logical_expr::{Expr, Operator}; + use datafusion::logical_expr::{Expr, ExprSchemable, Operator}; use insta::assert_snapshot; use substrait::proto::expression::literal::LiteralType; use substrait::proto::expression::{Literal, RexType, ScalarFunction}; use substrait::proto::function_argument::ArgType; use substrait::proto::{Expression, FunctionArgument}; + fn decimal_function( + name: &str, + left_type: DataType, + right_type: DataType, + output_type: Option, + ) -> Result<(Extensions, ScalarFunction, DFSchema)> { + let mut extensions = Extensions::default(); + extensions.functions.insert(0, format!("{name}:dec_dec")); + let mut producer = DefaultSubstraitProducer::new(&TEST_SESSION_STATE); + let func = ScalarFunction { + function_reference: 0, + arguments: (0..2) + .map(|index| { + Ok(FunctionArgument { + arg_type: Some(ArgType::Value(substrait_field_ref(index)?)), + }) + }) + .collect::>()?, + output_type: output_type + .map(|dt| to_substrait_type(&mut producer, &dt, false)) + .transpose()?, + ..Default::default() + }; + let schema = DFSchema::try_from(Schema::new(vec![ + Field::new("a", left_type, false), + Field::new("b", right_type, false), + ]))?; + Ok((extensions, func, schema)) + } + + #[tokio::test] + async fn test_decimal_arithmetic_output_types() -> Result<()> { + use DataType::Decimal128 as D; + + // The five cases from #25043, plus subtraction and both remainder names. + for (name, left, right, declared, derived) in [ + ("add", D(10, 2), D(5, 1), D(11, 2), D(11, 2)), + ("add", D(38, 10), D(38, 10), D(38, 9), D(38, 10)), + ("multiply", D(10, 2), D(5, 1), D(16, 3), D(16, 3)), + ("multiply", D(38, 10), D(38, 10), D(38, 6), D(38, 20)), + ("divide", D(10, 2), D(5, 1), D(21, 8), D(15, 6)), + ("subtract", D(38, 10), D(38, 10), D(38, 9), D(38, 10)), + ("modulus", D(10, 2), D(5, 1), D(6, 2), D(6, 2)), + ("mod", D(10, 2), D(5, 1), D(7, 2), D(6, 2)), + ( + "add", + DataType::Decimal256(10, 2), + DataType::Decimal256(5, 1), + DataType::Decimal256(11, 2), + DataType::Decimal256(11, 2), + ), + ] { + let (extensions, func, schema) = + decimal_function(name, left, right, Some(declared.clone()))?; + let consumer = + DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); + let result = consumer.consume_scalar_function(&func, &schema).await; + if declared == derived { + let expr = result?; + assert_eq!(expr.get_type(&schema)?, declared); + assert!(matches!(expr, Expr::BinaryExpr(_))); + } else { + let err = result.unwrap_err().to_string(); + assert!( + err.contains(&format!( + "Decimal return type mismatch for {name}:dec_dec" + )), + "{err}" + ); + assert!(err.contains(&format!("declared {declared:?}")), "{err}"); + assert!(err.contains(&format!("derives {derived:?}")), "{err}"); + } + } + Ok(()) + } + + #[tokio::test] + async fn test_decimal_output_type_compatibility() -> Result<()> { + for output_type in [None, Some(DataType::Decimal128(15, 6))] { + let (extensions, func, schema) = decimal_function( + "divide", + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + output_type, + )?; + let consumer = + DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); + let expr = consumer.consume_scalar_function(&func, &schema).await?; + assert!(matches!(expr, Expr::BinaryExpr(_))); + assert_eq!(expr.get_type(&schema)?, DataType::Decimal128(15, 6)); + } + + for (left, right, declared) in [ + ( + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + DataType::Decimal128(12, 2), + ), + ( + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + DataType::Decimal256(11, 2), + ), + ( + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + DataType::Float64, + ), + ( + DataType::Int64, + DataType::Int64, + DataType::Decimal128(20, 0), + ), + ] { + let (extensions, func, schema) = + decimal_function("add", left, right, Some(declared))?; + let consumer = + DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); + let err = consumer + .consume_scalar_function(&func, &schema) + .await + .unwrap_err(); + assert!( + err.to_string().contains("Decimal return type mismatch"), + "{err}" + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_invalid_decimal_output_type() -> Result<()> { + for (precision, scale) in + [(267, 2), (11, 258), (-245, 2), (0, 0), (39, 2), (11, 12)] + { + let (extensions, mut func, schema) = decimal_function( + "add", + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + Some(DataType::Decimal128(11, 2)), + )?; + let Some(super::Kind::Decimal(decimal)) = + &mut func.output_type.as_mut().unwrap().kind + else { + unreachable!() + }; + decimal.precision = precision; + decimal.scale = scale; + let consumer = + DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); + assert!( + consumer + .consume_scalar_function(&func, &schema) + .await + .is_err() + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_nested_decimal_output_type_mismatch() -> Result<()> { + let (mut extensions, mut func, schema) = decimal_function( + "divide", + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + Some(DataType::Decimal128(21, 8)), + )?; + extensions.functions.insert(1, "add:dec_dec".to_string()); + let inner = Expression { + rex_type: Some(RexType::ScalarFunction(func.clone())), + }; + func.function_reference = 1; + func.output_type = None; + func.arguments[0].arg_type = Some(ArgType::Value(inner)); + let consumer = DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); + let err = consumer + .consume_scalar_function(&func, &schema) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Decimal return type mismatch for divide:dec_dec"), + "{err}" + ); + Ok(()) + } + /// Test that large argument lists for binary operations do not crash the consumer #[tokio::test] async fn test_binary_op_large_argument_list() -> Result<()> { diff --git a/datafusion/substrait/tests/cases/consumer_integration.rs b/datafusion/substrait/tests/cases/consumer_integration.rs index adfbb4d7192bf..41ced2a740bc2 100644 --- a/datafusion/substrait/tests/cases/consumer_integration.rs +++ b/datafusion/substrait/tests/cases/consumer_integration.rs @@ -79,18 +79,12 @@ mod tests { #[tokio::test] async fn tpch_test_01() -> Result<()> { - let plan_str = tpch_plan_to_string(1).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(1).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, sum(LINEITEM.L_QUANTITY) AS SUM_QTY, sum(LINEITEM.L_EXTENDEDPRICE) AS SUM_BASE_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS SUM_DISC_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT * Int32(1) + LINEITEM.L_TAX) AS SUM_CHARGE, avg(LINEITEM.L_QUANTITY) AS AVG_QTY, avg(LINEITEM.L_EXTENDEDPRICE) AS AVG_PRICE, avg(LINEITEM.L_DISCOUNT) AS AVG_DISC, count(Int64(1)) AS COUNT_ORDER - Sort: LINEITEM.L_RETURNFLAG ASC NULLS LAST, LINEITEM.L_LINESTATUS ASC NULLS LAST - Aggregate: groupBy=[[LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS]], aggr=[[sum(LINEITEM.L_QUANTITY), sum(LINEITEM.L_EXTENDEDPRICE), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT * Int32(1) + LINEITEM.L_TAX), avg(LINEITEM.L_QUANTITY), avg(LINEITEM.L_EXTENDEDPRICE), avg(LINEITEM.L_DISCOUNT), count(Int64(1))]] - Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, LINEITEM.L_QUANTITY, LINEITEM.L_EXTENDEDPRICE, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT), LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) * (CAST(Int32(1) AS Decimal128(15, 2)) + LINEITEM.L_TAX), LINEITEM.L_DISCOUNT - Filter: LINEITEM.L_SHIPDATE <= Date32("1998-12-01") - IntervalDayTime("IntervalDayTime { days: 0, milliseconds: 10368000 }") - TableScan: LINEITEM - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 2): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" + ); Ok(()) } @@ -131,24 +125,12 @@ mod tests { #[tokio::test] async fn tpch_test_03() -> Result<()> { - let plan_str = tpch_plan_to_string(3).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(3).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Projection: LINEITEM.L_ORDERKEY, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS REVENUE, ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY - Limit: skip=0, fetch=10 - Sort: sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) DESC NULLS FIRST, ORDERS.O_ORDERDATE ASC NULLS LAST - Projection: LINEITEM.L_ORDERKEY, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT), ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY - Aggregate: groupBy=[[LINEITEM.L_ORDERKEY, ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT)]] - Projection: LINEITEM.L_ORDERKEY, ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) - Filter: CUSTOMER.C_MKTSEGMENT = Utf8("BUILDING") AND CUSTOMER.C_CUSTKEY = ORDERS.O_CUSTKEY AND LINEITEM.L_ORDERKEY = ORDERS.O_ORDERKEY AND ORDERS.O_ORDERDATE < CAST(Utf8("1995-03-15") AS Date32) AND LINEITEM.L_SHIPDATE > CAST(Utf8("1995-03-15") AS Date32) - Cross Join: - Cross Join: - TableScan: LINEITEM - TableScan: CUSTOMER - TableScan: ORDERS - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" + ); Ok(()) } @@ -174,43 +156,23 @@ mod tests { #[tokio::test] async fn tpch_test_05() -> Result<()> { - let plan_str = tpch_plan_to_string(5).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(5).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Projection: NATION.N_NAME, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS REVENUE - Sort: sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) DESC NULLS FIRST - Aggregate: groupBy=[[NATION.N_NAME]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT)]] - Projection: NATION.N_NAME, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) - Filter: CUSTOMER.C_CUSTKEY = ORDERS.O_CUSTKEY AND LINEITEM.L_ORDERKEY = ORDERS.O_ORDERKEY AND LINEITEM.L_SUPPKEY = SUPPLIER.S_SUPPKEY AND CUSTOMER.C_NATIONKEY = SUPPLIER.S_NATIONKEY AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_REGIONKEY = REGION.R_REGIONKEY AND REGION.R_NAME = Utf8("ASIA") AND ORDERS.O_ORDERDATE >= CAST(Utf8("1994-01-01") AS Date32) AND ORDERS.O_ORDERDATE < CAST(Utf8("1995-01-01") AS Date32) - Cross Join: - Cross Join: - Cross Join: - Cross Join: - Cross Join: - TableScan: CUSTOMER - TableScan: ORDERS - TableScan: LINEITEM - TableScan: SUPPLIER - TableScan: NATION - TableScan: REGION - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" + ); Ok(()) } #[tokio::test] async fn tpch_test_06() -> Result<()> { - let plan_str = tpch_plan_to_string(6).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(6).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE * LINEITEM.L_DISCOUNT) AS REVENUE]] - Projection: LINEITEM.L_EXTENDEDPRICE * LINEITEM.L_DISCOUNT - Filter: LINEITEM.L_SHIPDATE >= CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-01-01") AS Date32) AND LINEITEM.L_DISCOUNT >= Decimal128(0.05,3,2) AND LINEITEM.L_DISCOUNT <= Decimal128(0.07,3,2) AND LINEITEM.L_QUANTITY < CAST(Int32(24) AS Decimal128(15, 2)) - TableScan: LINEITEM - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 6): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * LINEITEM.L_DISCOUNT derives Decimal128(31, 4); this conversion is unsupported" + ); Ok(()) } @@ -240,58 +202,23 @@ mod tests { #[tokio::test] async fn tpch_test_10() -> Result<()> { - let plan_str = tpch_plan_to_string(10).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(10).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Projection: CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS REVENUE, CUSTOMER.C_ACCTBAL, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_PHONE, CUSTOMER.C_COMMENT - Limit: skip=0, fetch=20 - Sort: sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) DESC NULLS FIRST - Projection: CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT), CUSTOMER.C_ACCTBAL, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_PHONE, CUSTOMER.C_COMMENT - Aggregate: groupBy=[[CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, CUSTOMER.C_ACCTBAL, CUSTOMER.C_PHONE, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_COMMENT]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT)]] - Projection: CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, CUSTOMER.C_ACCTBAL, CUSTOMER.C_PHONE, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_COMMENT, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) - Filter: CUSTOMER.C_CUSTKEY = ORDERS.O_CUSTKEY AND LINEITEM.L_ORDERKEY = ORDERS.O_ORDERKEY AND ORDERS.O_ORDERDATE >= CAST(Utf8("1993-10-01") AS Date32) AND ORDERS.O_ORDERDATE < CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_RETURNFLAG = Utf8("R") AND CUSTOMER.C_NATIONKEY = NATION.N_NATIONKEY - Cross Join: - Cross Join: - Cross Join: - TableScan: CUSTOMER - TableScan: ORDERS - TableScan: LINEITEM - TableScan: NATION - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" + ); Ok(()) } #[tokio::test] async fn tpch_test_11() -> Result<()> { - let plan_str = tpch_plan_to_string(11).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(11).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Projection: PARTSUPP.PS_PARTKEY, sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) AS value - Sort: sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) DESC NULLS FIRST - Filter: sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) > () - Subquery: - Projection: sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) * Decimal128(0.0001000000,11,10) - Aggregate: groupBy=[[]], aggr=[[sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY)]] - Projection: PARTSUPP.PS_SUPPLYCOST * CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0)) - Filter: PARTSUPP.PS_SUPPKEY = SUPPLIER.S_SUPPKEY AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8("JAPAN") - Cross Join: - Cross Join: - TableScan: PARTSUPP - TableScan: SUPPLIER - TableScan: NATION - Aggregate: groupBy=[[PARTSUPP.PS_PARTKEY]], aggr=[[sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY)]] - Projection: PARTSUPP.PS_PARTKEY, PARTSUPP.PS_SUPPLYCOST * CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0)) - Filter: PARTSUPP.PS_SUPPKEY = SUPPLIER.S_SUPPKEY AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8("JAPAN") - Cross Join: - Cross Join: - TableScan: PARTSUPP - TableScan: SUPPLIER - TableScan: NATION - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 2): declared Decimal128(19, 2), but native expression PARTSUPP.PS_SUPPLYCOST * CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0)) derives Decimal128(35, 2); this conversion is unsupported" + ); Ok(()) } @@ -336,19 +263,12 @@ mod tests { #[tokio::test] async fn tpch_test_14() -> Result<()> { - let plan_str = tpch_plan_to_string(14).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(14).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Projection: Decimal128(100.00,5,2) * sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT ELSE Decimal128(0.0000,19,4) END) / sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS PROMO_REVENUE - Aggregate: groupBy=[[]], aggr=[[sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT ELSE Decimal128(0.0000,19,4) END), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT)]] - Projection: CASE WHEN PART.P_TYPE LIKE CAST(Utf8("PROMO%") AS Utf8) THEN LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) ELSE Decimal128(0.0000,19,4) END, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) - Filter: LINEITEM.L_PARTKEY = PART.P_PARTKEY AND LINEITEM.L_SHIPDATE >= Date32("1995-09-01") AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-10-01") AS Date32) - Cross Join: - TableScan: LINEITEM - TableScan: PART - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 5): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" + ); Ok(()) } @@ -385,25 +305,12 @@ mod tests { #[tokio::test] async fn tpch_test_17() -> Result<()> { - let plan_str = tpch_plan_to_string(17).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(17).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Projection: sum(LINEITEM.L_EXTENDEDPRICE) / Decimal128(7.0,2,1) AS AVG_YEARLY - Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE)]] - Projection: LINEITEM.L_EXTENDEDPRICE - Filter: PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#23") AND PART.P_CONTAINER = Utf8("MED BOX") AND LINEITEM.L_QUANTITY < () - Subquery: - Projection: Decimal128(0.2,2,1) * avg(LINEITEM.L_QUANTITY) - Aggregate: groupBy=[[]], aggr=[[avg(LINEITEM.L_QUANTITY)]] - Projection: LINEITEM.L_QUANTITY - Filter: LINEITEM.L_PARTKEY = outer_ref(PART.P_PARTKEY) - TableScan: LINEITEM - Cross Join: - TableScan: LINEITEM - TableScan: PART - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(17, 3), but native expression Decimal128(0.2,2,1) * avg(LINEITEM.L_QUANTITY) derives Decimal128(22, 7); this conversion is unsupported" + ); Ok(()) } @@ -436,49 +343,23 @@ mod tests { } #[tokio::test] async fn tpch_test_19() -> Result<()> { - let plan_str = tpch_plan_to_string(19).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(19).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS REVENUE]] - Projection: LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) - Filter: PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#12") AND (PART.P_CONTAINER = CAST(Utf8("SM CASE") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("SM BOX") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("SM PACK") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("SM PKG") AS Utf8)) AND LINEITEM.L_QUANTITY >= CAST(Int32(1) AS Decimal128(15, 2)) AND LINEITEM.L_QUANTITY <= CAST(Int32(1) + Int32(10) AS Decimal128(15, 2)) AND PART.P_SIZE >= Int32(1) AND PART.P_SIZE <= Int32(5) AND (LINEITEM.L_SHIPMODE = CAST(Utf8("AIR") AS Utf8) OR LINEITEM.L_SHIPMODE = CAST(Utf8("AIR REG") AS Utf8)) AND LINEITEM.L_SHIPINSTRUCT = Utf8("DELIVER IN PERSON") OR PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#23") AND (PART.P_CONTAINER = CAST(Utf8("MED BAG") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("MED BOX") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("MED PKG") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("MED PACK") AS Utf8)) AND LINEITEM.L_QUANTITY >= CAST(Int32(10) AS Decimal128(15, 2)) AND LINEITEM.L_QUANTITY <= CAST(Int32(10) + Int32(10) AS Decimal128(15, 2)) AND PART.P_SIZE >= Int32(1) AND PART.P_SIZE <= Int32(10) AND (LINEITEM.L_SHIPMODE = CAST(Utf8("AIR") AS Utf8) OR LINEITEM.L_SHIPMODE = CAST(Utf8("AIR REG") AS Utf8)) AND LINEITEM.L_SHIPINSTRUCT = Utf8("DELIVER IN PERSON") OR PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#34") AND (PART.P_CONTAINER = CAST(Utf8("LG CASE") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("LG BOX") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("LG PACK") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("LG PKG") AS Utf8)) AND LINEITEM.L_QUANTITY >= CAST(Int32(20) AS Decimal128(15, 2)) AND LINEITEM.L_QUANTITY <= CAST(Int32(20) + Int32(10) AS Decimal128(15, 2)) AND PART.P_SIZE >= Int32(1) AND PART.P_SIZE <= Int32(15) AND (LINEITEM.L_SHIPMODE = CAST(Utf8("AIR") AS Utf8) OR LINEITEM.L_SHIPMODE = CAST(Utf8("AIR REG") AS Utf8)) AND LINEITEM.L_SHIPINSTRUCT = Utf8("DELIVER IN PERSON") - Cross Join: - TableScan: LINEITEM - TableScan: PART - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 6): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" + ); Ok(()) } #[tokio::test] async fn tpch_test_20() -> Result<()> { - let plan_str = tpch_plan_to_string(20).await?; + // This fixture declares a decimal result type the native operator cannot provide. + let err = tpch_plan_to_string(20).await.unwrap_err(); assert_snapshot!( - plan_str, - @r#" - Sort: SUPPLIER.S_NAME ASC NULLS LAST - Projection: SUPPLIER.S_NAME, SUPPLIER.S_ADDRESS - Filter: SUPPLIER.S_SUPPKEY IN () AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8("CANADA") - Subquery: - Projection: PARTSUPP.PS_SUPPKEY - Filter: PARTSUPP.PS_PARTKEY IN () AND CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0)) > () - Subquery: - Projection: PART.P_PARTKEY - Filter: PART.P_NAME LIKE CAST(Utf8("forest%") AS Utf8) - TableScan: PART - Subquery: - Projection: Decimal128(0.5,2,1) * sum(LINEITEM.L_QUANTITY) - Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_QUANTITY)]] - Projection: LINEITEM.L_QUANTITY - Filter: LINEITEM.L_PARTKEY = outer_ref(PARTSUPP.PS_PARTKEY) AND LINEITEM.L_SUPPKEY = outer_ref(PARTSUPP.PS_SUPPKEY) AND LINEITEM.L_SHIPDATE >= CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-01-01") AS Date32) - TableScan: LINEITEM - TableScan: PARTSUPP - Cross Join: - TableScan: SUPPLIER - TableScan: NATION - "# - ); + err, + @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 7): declared Decimal128(17, 3), but native expression Decimal128(0.5,2,1) * sum(LINEITEM.L_QUANTITY) derives Decimal128(28, 3); this conversion is unsupported" + ); Ok(()) } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 3716b0feba3cc..5ee5c90091fc5 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -464,6 +464,92 @@ async fn decimal_literal() -> Result<()> { roundtrip("SELECT * FROM data WHERE b > 2.5").await } +#[tokio::test] +async fn decimal_arithmetic_output_type_contract() -> Result<()> { + use datafusion::arrow::array::Decimal128Array; + use datafusion::arrow::record_batch::RecordBatch; + use substrait::proto::expression::RexType; + use substrait::proto::r#type::Kind; + + for (left, right, op, native, declared, expected) in [ + ((10, 2), (5, 1), "+", (11, 2), (11, 2), 400_i128), + ((38, 10), (38, 10), "+", (38, 10), (38, 9), 40_000_000_000), + ((10, 2), (5, 1), "*", (16, 3), (16, 3), 3000), + ( + (38, 10), + (38, 10), + "*", + (38, 20), + (38, 6), + 300_000_000_000_000_000_000, + ), + ((10, 2), (5, 1), "/", (15, 6), (21, 8), 333_333), + ] { + let ctx = SessionContext::new(); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Decimal128(left.0, left.1), false), + Field::new("b", DataType::Decimal128(right.0, right.1), false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new( + Decimal128Array::from(vec![10_i128.pow(left.1 as u32)]) + .with_precision_and_scale(left.0, left.1)?, + ), + Arc::new( + Decimal128Array::from(vec![3 * 10_i128.pow(right.1 as u32)]) + .with_precision_and_scale(right.0, right.1)?, + ), + ], + )?; + ctx.register_batch("decimals", batch)?; + let df = ctx + .sql(&format!("SELECT a {op} b AS result FROM decimals")) + .await?; + let mut proto = to_substrait_plan(df.logical_plan(), &ctx.state())?; + + // DataFusion-produced declarations retain their schema and execution. + let consumed = from_substrait_plan(&ctx.state(), &proto).await?; + let expected_type = DataType::Decimal128(native.0, native.1); + assert_eq!(consumed.schema().field(0).data_type(), &expected_type); + let batches = ctx.execute_logical_plan(consumed).await?.collect().await?; + assert_eq!( + ScalarValue::try_from_array(batches[0].column(0), 0)?, + ScalarValue::Decimal128(Some(expected), native.0, native.1) + ); + + // Replace the producer's return type with the Substrait decimal rule. + let Some(plan_rel::RelType::Root(root)) = &mut proto.relations[0].rel_type else { + panic!("expected root relation") + }; + let Some(RelType::Project(project)) = &mut root.input.as_mut().unwrap().rel_type + else { + panic!("expected projection") + }; + let Some(RexType::ScalarFunction(function)) = + &mut project.expressions[0].rex_type + else { + panic!("expected scalar function") + }; + let Some(Kind::Decimal(decimal)) = + &mut function.output_type.as_mut().unwrap().kind + else { + panic!("expected decimal output type") + }; + decimal.precision = declared.0; + decimal.scale = declared.1; + let result = from_substrait_plan(&ctx.state(), &proto).await; + if native == (declared.0 as u8, declared.1 as i8) { + assert_eq!(result?.schema().field(0).data_type(), &expected_type); + } else { + let err = result.unwrap_err().to_string(); + assert!(err.contains("Decimal return type mismatch"), "{err}"); + } + } + Ok(()) +} + #[tokio::test] async fn null_decimal_literal() -> Result<()> { roundtrip("SELECT *, CAST(NULL AS decimal(10, 2)) FROM data").await From 555b760d429d9753ef3a13f2cff8fd9c2c0f441d Mon Sep 17 00:00:00 2001 From: Daipayan Date: Tue, 8 Sep 2026 23:21:31 +0100 Subject: [PATCH 2/2] fix(substrait): honor declared decimal arithmetic types --- .../consumer/expr/scalar_function.rs | 395 ++++++++++++++-- .../substrait/src/logical_plan/decimal.rs | 420 ++++++++++++++++++ datafusion/substrait/src/logical_plan/mod.rs | 1 + .../producer/expr/scalar_function.rs | 20 + .../tests/cases/consumer_integration.rs | 219 ++++++--- .../tests/cases/roundtrip_logical_plan.rs | 28 +- 6 files changed, 983 insertions(+), 100 deletions(-) create mode 100644 datafusion/substrait/src/logical_plan/decimal.rs diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs index a65be3052adbc..c35212eddd7b1 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs @@ -18,6 +18,7 @@ use crate::logical_plan::consumer::{ SubstraitConsumer, from_substrait_func_args, from_substrait_type_without_names, }; +use crate::logical_plan::decimal::DecimalArithmetic; use datafusion::arrow::datatypes::{ DataType, Decimal128Type, Decimal256Type, validate_decimal_precision_and_scale, }; @@ -27,7 +28,7 @@ use datafusion::common::{ }; use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::{ - Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, expr, + Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, ScalarUDF, expr, }; use std::vec::Drain; use substrait::proto::expression::ScalarFunction; @@ -103,7 +104,13 @@ pub async fn from_scalar_function( | Operator::Divide | Operator::Modulo ) { - validate_decimal_output_type(consumer, f, fn_signature, &expr, input_schema)?; + return apply_decimal_output_type( + consumer, + f, + fn_signature, + expr, + input_schema, + ); } Ok(expr) } else if let Some(builder) = BuiltinExprBuilder::try_from_name(fn_name) { @@ -113,24 +120,28 @@ pub async fn from_scalar_function( } } -/// Native decimal arithmetic can derive different precision and scale from the -/// referenced Substrait function. Reject mismatches rather than silently changing -/// the result type. An output cast cannot recover lost digits or prevent an -/// intermediate arithmetic overflow. -fn validate_decimal_output_type( +/// Preserve the declared decimal type with arithmetic evaluated at that scale. +/// Keep native expressions when their result type matches and no explicit +/// function options require different execution behavior. +fn apply_decimal_output_type( consumer: &impl SubstraitConsumer, f: &ScalarFunction, fn_signature: &str, - expr: &Expr, + expr: Expr, input_schema: &DFSchema, -) -> Result<()> { +) -> Result { // Preserve compatibility with existing plans that omit the output type. let Some(output_type) = &f.output_type else { - return Ok(()); + return Ok(expr); }; - let derived_type = expr.get_type(input_schema)?; - if !derived_type.is_decimal() && !matches!(output_type.kind, Some(Kind::Decimal(_))) { - return Ok(()); + // Native inference can itself fail (for example, multiplication with an + // intermediate scale over 38). A valid declared decimal type can still be + // implemented by DecimalArithmetic in that case. + let derived_type = expr.get_type(input_schema); + if !matches!(output_type.kind, Some(Kind::Decimal(_))) + && !derived_type.as_ref().is_ok_and(|dt| dt.is_decimal()) + { + return Ok(expr); } // The type decoder narrows these protobuf integers with `as`. Check them @@ -155,13 +166,26 @@ fn validate_decimal_output_type( } _ => {} } - if declared_type != derived_type { + if derived_type.as_ref().is_ok_and(|dt| *dt == declared_type) && f.options.is_empty() + { + return Ok(expr); + } + let Expr::BinaryExpr(BinaryExpr { left, op, right }) = expr else { + unreachable!("called only for native binary expressions") + }; + if f.arguments.len() != 2 { return substrait_err!( - "Decimal return type mismatch for {fn_signature} (function reference {}): declared {declared_type:?}, but native expression {expr} derives {derived_type:?}; this conversion is unsupported", - f.function_reference + "Declared decimal arithmetic requires two arguments for {fn_signature}" ); } - Ok(()) + let function = DecimalArithmetic::try_new( + fn_signature, + op, + [left.get_type(input_schema)?, right.get_type(input_schema)?], + &declared_type, + &f.options, + )?; + Ok(ScalarUDF::from(function).call(vec![*left, *right])) } pub fn substrait_fun_name(name: &str) -> &str { @@ -516,21 +540,12 @@ mod tests { decimal_function(name, left, right, Some(declared.clone()))?; let consumer = DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); - let result = consumer.consume_scalar_function(&func, &schema).await; + let expr = consumer.consume_scalar_function(&func, &schema).await?; + assert_eq!(expr.get_type(&schema)?, declared); if declared == derived { - let expr = result?; - assert_eq!(expr.get_type(&schema)?, declared); assert!(matches!(expr, Expr::BinaryExpr(_))); } else { - let err = result.unwrap_err().to_string(); - assert!( - err.contains(&format!( - "Decimal return type mismatch for {name}:dec_dec" - )), - "{err}" - ); - assert!(err.contains(&format!("declared {declared:?}")), "{err}"); - assert!(err.contains(&format!("derives {derived:?}")), "{err}"); + assert!(matches!(expr, Expr::ScalarFunction(_))); } } Ok(()) @@ -553,11 +568,6 @@ mod tests { } for (left, right, declared) in [ - ( - DataType::Decimal128(10, 2), - DataType::Decimal128(5, 1), - DataType::Decimal128(12, 2), - ), ( DataType::Decimal128(10, 2), DataType::Decimal128(5, 1), @@ -583,7 +593,8 @@ mod tests { .await .unwrap_err(); assert!( - err.to_string().contains("Decimal return type mismatch"), + err.to_string() + .contains("Unsupported decimal arithmetic type"), "{err}" ); } @@ -621,7 +632,7 @@ mod tests { } #[tokio::test] - async fn test_nested_decimal_output_type_mismatch() -> Result<()> { + async fn test_nested_decimal_output_type() -> Result<()> { let (mut extensions, mut func, schema) = decimal_function( "divide", DataType::Decimal128(10, 2), @@ -636,18 +647,314 @@ mod tests { func.output_type = None; func.arguments[0].arg_type = Some(ArgType::Value(inner)); let consumer = DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); - let err = consumer - .consume_scalar_function(&func, &schema) - .await - .unwrap_err(); - assert!( - err.to_string() - .contains("Decimal return type mismatch for divide:dec_dec"), - "{err}" + let expr = consumer.consume_scalar_function(&func, &schema).await?; + assert_eq!(expr.get_type(&schema)?, DataType::Decimal128(22, 8)); + Ok(()) + } + + #[tokio::test] + async fn test_declared_decimal_execution() -> Result<()> { + use datafusion::arrow::array::Decimal128Array; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use std::sync::Arc; + + // Include cases whose *native type inference* fails because the + // intermediate scale exceeds 38, as well as native execution overflow. + for (name, left_type, right_type, output, left, right, expected) in [ + ("divide", (10, 2), (5, 1), (21, 8), 100, 30, 33_333_333), + ( + "divide", + (10, 2), + (10, 2), + (31, 13), + 100, + 300, + 3_333_333_333_333, + ), + ("divide", (18, 2), (18, 2), (38, 6), 100, 300, 333_333), + ( + "divide", + (38, 38), + (38, 38), + (38, 6), + 10_i128.pow(37), + 10_i128.pow(37), + 1_000_000, + ), + ( + "multiply", + (38, 38), + (38, 38), + (38, 6), + 10_i128.pow(37), + 10_i128.pow(37), + 10_000, + ), + ( + "add", + (38, 10), + (38, 10), + (38, 9), + 9 * 10_i128.pow(37), + 9 * 10_i128.pow(37), + 18 * 10_i128.pow(36), + ), + ( + "subtract", + (38, 10), + (38, 10), + (38, 9), + 9 * 10_i128.pow(37), + 9 * 10_i128.pow(37) - 6, + 1, + ), + ("modulus", (10, 3), (5, 1), (5, 2), 12_345, 20, 35), + ] { + let left_type = DataType::Decimal128(left_type.0, left_type.1); + let right_type = DataType::Decimal128(right_type.0, right_type.1); + let (extensions, mut func, _) = decimal_function( + name, + left_type.clone(), + right_type.clone(), + Some(DataType::Decimal128(output.0, output.1)), + )?; + if let Some(super::Kind::Decimal(decimal)) = + &mut func.output_type.as_mut().unwrap().kind + { + decimal.nullability = + substrait::proto::r#type::Nullability::Nullable as i32; + } + let ctx = SessionContext::new(); + let state = ctx.state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &state); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", left_type.clone(), true), + Field::new("b", right_type.clone(), true), + ])); + let df_schema = DFSchema::try_from(schema.as_ref().clone())?; + let expr = consumer.consume_scalar_function(&func, &df_schema).await?; + let DataType::Decimal128(lp, ls) = left_type else { + unreachable!() + }; + let DataType::Decimal128(rp, rs) = right_type else { + unreachable!() + }; + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new( + Decimal128Array::from(vec![Some(left), None, Some(-left)]) + .with_precision_and_scale(lp, ls)?, + ), + Arc::new( + Decimal128Array::from(vec![Some(right), Some(0), Some(-right)]) + .with_precision_and_scale(rp, rs)?, + ), + ], + )?; + ctx.register_batch("decimals", batch)?; + let batches = ctx + .table("decimals") + .await? + .select(vec![expr])? + .collect() + .await?; + assert_eq!( + ScalarValue::try_from_array(batches[0].column(0), 0)?, + ScalarValue::Decimal128(Some(expected), output.0, output.1), + "{name}" + ); + // NULL propagation must skip arithmetic, including division by 0. + assert_eq!( + ScalarValue::try_from_array(batches[0].column(0), 1)?, + ScalarValue::Decimal128(None, output.0, output.1), + "{name}" + ); + let negative = if matches!(name, "divide" | "multiply") { + expected + } else { + -expected + }; + assert_eq!( + ScalarValue::try_from_array(batches[0].column(0), 2)?, + ScalarValue::Decimal128(Some(negative), output.0, output.1), + "{name}" + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_declared_decimal_scalar_and_array_arguments() -> Result<()> { + use crate::logical_plan::producer::to_substrait_literal_expr; + use datafusion::arrow::array::{Array, Decimal128Array}; + use datafusion::arrow::record_batch::RecordBatch; + use std::sync::Arc; + + for scalar_left in [false, true] { + for scalar_right in [false, true] { + let (extensions, mut func, schema) = decimal_function( + "divide", + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + Some(DataType::Decimal128(21, 8)), + )?; + let mut producer = DefaultSubstraitProducer::new(&TEST_SESSION_STATE); + for (index, scalar, value) in [ + (0, scalar_left, ScalarValue::Decimal128(Some(200), 10, 2)), + (1, scalar_right, ScalarValue::Decimal128(Some(30), 5, 1)), + ] { + if scalar { + func.arguments[index].arg_type = Some(ArgType::Value( + to_substrait_literal_expr(&mut producer, &value)?, + )); + } + } + let consumer = + DefaultSubstraitConsumer::new(&extensions, &TEST_SESSION_STATE); + let expr = consumer.consume_scalar_function(&func, &schema).await?; + let physical = TEST_SESSION_STATE.create_physical_expr(expr, &schema)?; + for rows in [0, 3] { + let batch = RecordBatch::try_new( + Arc::new(schema.as_arrow().clone()), + vec![ + Arc::new( + Decimal128Array::from(vec![200; rows]) + .with_precision_and_scale(10, 2)?, + ), + Arc::new( + Decimal128Array::from(vec![30; rows]) + .with_precision_and_scale(5, 1)?, + ), + ], + )?; + let result = physical.evaluate(&batch)?.into_array(rows)?; + assert_eq!(result.len(), rows); + if rows > 0 { + assert_eq!( + ScalarValue::try_from_array(&result, 2)?, + ScalarValue::Decimal128(Some(66_666_667), 21, 8) + ); + } + } + } + } + Ok(()) + } + + #[tokio::test] + async fn test_decimal_output_types_survive_optimization() -> Result<()> { + use datafusion::arrow::array::{ArrayRef, Decimal128Array}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use std::sync::Arc; + + let ctx = SessionContext::new(); + let state = ctx.state(); + let mut expressions = Vec::new(); + for (scale, alias) in [(8, "eight"), (7, "seven")] { + let (extensions, func, schema) = decimal_function( + "divide", + DataType::Decimal128(10, 2), + DataType::Decimal128(5, 1), + Some(DataType::Decimal128(21, scale)), + )?; + let consumer = DefaultSubstraitConsumer::new(&extensions, &state); + expressions.push( + consumer + .consume_scalar_function(&func, &schema) + .await? + .alias(alias), + ); + } + let batch = RecordBatch::try_from_iter(vec![ + ( + "a", + Arc::new( + Decimal128Array::from(vec![100]).with_precision_and_scale(10, 2)?, + ) as ArrayRef, + ), + ( + "b", + Arc::new(Decimal128Array::from(vec![30]).with_precision_and_scale(5, 1)?), + ), + ])?; + ctx.register_batch("decimals", batch)?; + let batches = ctx + .table("decimals") + .await? + .select(expressions)? + .collect() + .await?; + assert_eq!( + ScalarValue::try_from_array(batches[0].column(0), 0)?, + ScalarValue::Decimal128(Some(33_333_333), 21, 8) + ); + assert_eq!( + ScalarValue::try_from_array(batches[0].column(1), 0)?, + ScalarValue::Decimal128(Some(3_333_333), 21, 7) ); Ok(()) } + #[tokio::test] + async fn test_decimal_overflow_option_survives_roundtrip() -> Result<()> { + use crate::logical_plan::consumer::from_substrait_plan; + use crate::logical_plan::producer::to_substrait_plan; + use datafusion::arrow::array::{ArrayRef, Decimal128Array}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use std::sync::Arc; + use substrait::proto::FunctionOption; + + let ctx = SessionContext::new(); + let state = ctx.state(); + // The declared type matches Arrow's, but the explicit ERROR option + // must still select checked arithmetic and survive serialization. + let (extensions, mut func, schema) = decimal_function( + "add", + DataType::Decimal128(38, 0), + DataType::Decimal128(38, 0), + Some(DataType::Decimal128(38, 0)), + )?; + func.options = vec![FunctionOption { + name: "overflow".into(), + preference: vec!["ERROR".into()], + }]; + let consumer = DefaultSubstraitConsumer::new(&extensions, &state); + let expr = consumer.consume_scalar_function(&func, &schema).await?; + assert!(matches!(expr, Expr::ScalarFunction(_))); + let batch = RecordBatch::try_from_iter(vec![ + ( + "a", + Arc::new( + Decimal128Array::from(vec![9 * 10_i128.pow(37)]) + .with_precision_and_scale(38, 0)?, + ) as ArrayRef, + ), + ( + "b", + Arc::new( + Decimal128Array::from(vec![9 * 10_i128.pow(37)]) + .with_precision_and_scale(38, 0)?, + ), + ), + ])?; + ctx.register_batch("decimals", batch)?; + let df = ctx + .table("decimals") + .await? + .select(vec![expr.alias("result")])?; + let proto = to_substrait_plan(df.logical_plan(), &state)?; + let reimported = from_substrait_plan(&state, &proto).await?; + for df in [df, ctx.execute_logical_plan(reimported).await?] { + let err = df.collect().await.unwrap_err(); + assert!(err.to_string().contains("Decimal overflow"), "{err}"); + } + Ok(()) + } + /// Test that large argument lists for binary operations do not crash the consumer #[tokio::test] async fn test_binary_op_large_argument_list() -> Result<()> { diff --git a/datafusion/substrait/src/logical_plan/decimal.rs b/datafusion/substrait/src/logical_plan/decimal.rs new file mode 100644 index 0000000000000..cf00b38eb4b14 --- /dev/null +++ b/datafusion/substrait/src/logical_plan/decimal.rs @@ -0,0 +1,420 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Arithmetic whose result precision and scale come from a Substrait call. +//! +//! Arrow's native decimal kernels derive their own result types. Casting their +//! output cannot recover discarded digits or avoid an intermediate overflow. +//! This function instead computes at the declared scale using checked i256 +//! arithmetic, without changing DataFusion's native SQL operators. +//! +//! Supports Decimal128 operands and outputs with nonnegative scales, including +//! all standard Substrait decimal precisions (1 through 38). Values are rounded +//! once to nearest, with ties away from zero, and checked against the output +//! precision. NULL inputs propagate; overflow and zero divisors return errors. +//! An explicit overflow option must allow ERROR. Other decimal representations +//! and function options require separate implementations. + +use std::sync::Arc; + +use datafusion::arrow::array::Decimal128Array; +use datafusion::arrow::datatypes::i256; +use datafusion::arrow::datatypes::validate_decimal_precision_and_scale; +use datafusion::arrow::datatypes::{DataType, Decimal128Type, Field, FieldRef}; +use datafusion::common::cast::as_primitive_array; +use datafusion::common::{ + Result, ScalarValue, exec_datafusion_err, exec_err, substrait_err, +}; +use datafusion::logical_expr::{ + ColumnarValue, Operator, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, Volatility, +}; +use substrait::proto::FunctionOption; + +/// A consumer-created function, with the declared output type included in its +/// identity so optimizations cannot merge calls with different decimal types. +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct DecimalArithmetic { + function_signature: String, + signature: Signature, + op: Operator, + input_scales: [i8; 2], + precision: u8, + scale: i8, +} + +impl DecimalArithmetic { + pub(crate) fn try_new( + function_signature: &str, + op: Operator, + input_types: [DataType; 2], + output_type: &DataType, + options: &[FunctionOption], + ) -> Result { + let decimal = |dt: &DataType| -> Result<(u8, i8)> { + if let DataType::Decimal128(p, s) = dt + && *s >= 0 + { + validate_decimal_precision_and_scale::(*p, *s)?; + return Ok((*p, *s)); + } + substrait_err!( + "Unsupported decimal arithmetic type {dt:?} for {function_signature}: \ + expected Decimal128 with nonnegative scale" + ) + }; + let (_, left_scale) = decimal(&input_types[0])?; + let (_, right_scale) = decimal(&input_types[1])?; + let (precision, scale) = decimal(output_type)?; + // An omitted option leaves behavior to the consumer. We choose ERROR. + // For a specified option, use the first supported preference (ERROR is + // currently the only supported overflow behavior), or reject the call. + let mut overflow_seen = false; + for option in options { + if !option.name.eq_ignore_ascii_case("overflow") || overflow_seen { + return substrait_err!( + "Unsupported or duplicate option {} for {function_signature}", + option.name + ); + } + overflow_seen = true; + if !option + .preference + .iter() + .any(|v| v.eq_ignore_ascii_case("ERROR")) + { + return substrait_err!( + "Unsupported overflow preferences {:?} for {function_signature}; supported: ERROR", + option.preference + ); + } + } + Ok(Self { + function_signature: function_signature.to_owned(), + signature: Signature::exact(Vec::from(input_types), Volatility::Immutable), + op, + input_scales: [left_scale, right_scale], + precision, + scale, + }) + } + + pub(crate) fn function_signature(&self) -> &str { + &self.function_signature + } + + fn output_type(&self) -> DataType { + DataType::Decimal128(self.precision, self.scale) + } + + fn evaluate_value(&self, left: i128, right: i128) -> Result { + if matches!(self.op, Operator::Divide | Operator::Modulo) && right == 0 { + return exec_err!("Divide by zero in Substrait {}", self.name()); + } + let left = i256::from_i128(left); + let right = i256::from_i128(right); + let [s1, s2] = self.input_scales.map(i16::from); + let output_scale = i16::from(self.scale); + // Each input has at most 38 digits. Exact products and operands aligned + // to max(s1, s2) fit in i256; rescale only after the operation so neither + // cancellation nor fractional carries are lost. + let result = match self.op { + Operator::Plus | Operator::Minus | Operator::Modulo => { + let scale = s1.max(s2); + left.checked_mul(power_of_ten(scale - s1)).and_then(|l| { + let r = right.checked_mul(power_of_ten(scale - s2))?; + let value = match self.op { + Operator::Plus => l.checked_add(r), + Operator::Minus => l.checked_sub(r), + _ => l.checked_rem(r), + }?; + rescale(value, scale, output_scale) + }) + } + Operator::Multiply => left + .checked_mul(right) + .and_then(|value| rescale(value, s1 + s2, output_scale)), + Operator::Divide => { + // Work directly at the output scale. In particular, do not use + // Arrow's s1 + 4 intermediate scale for high-scale operands. + let exponent = output_scale + s2 - s1; + if exponent >= 0 { + left.checked_mul(power_of_ten(exponent)) + .and_then(|numerator| rounded_div(numerator, right)) + } else { + right + .checked_mul(power_of_ten(-exponent)) + .and_then(|denominator| rounded_div(left, denominator)) + } + // If scaling the numerator overflows i256, dividing by a + // <=38-digit denominator cannot yield a <=38-digit result. + // Thus this failure cannot reject a representable output. + } + _ => unreachable!("only decimal arithmetic operators are constructed"), + }; + let max = power_of_ten(i16::from(self.precision)) - i256::ONE; + result + .filter(|v| *v >= -max && *v <= max) + .and_then(i256::to_i128) + .ok_or_else(|| { + exec_datafusion_err!( + "Decimal overflow in Substrait {}: result does not fit {:?}", + self.name(), + self.output_type() + ) + }) + } +} + +/// Exponents are bounded by 76 by the validated input and output types. +fn power_of_ten(exponent: i16) -> i256 { + i256::from_i128(10).wrapping_pow(exponent as u32) +} + +fn rescale(value: i256, from: i16, to: i16) -> Option { + if to >= from { + value.checked_mul(power_of_ten(to - from)) + } else { + rounded_div(value, power_of_ten(from - to)) + } +} + +/// Round once, to nearest with ties away from zero, as Arrow's decimal casts do. +/// The decimal arithmetic extension does not specify a rounding option. +fn rounded_div(numerator: i256, denominator: i256) -> Option { + let quotient = numerator.checked_div(denominator)?; + let remainder = numerator.checked_rem(denominator)?.checked_abs()?; + let divisor = denominator.checked_abs()?; + // ceil(divisor / 2), without doubling a potentially large remainder. + let half = divisor / i256::from_i128(2) + divisor % i256::from_i128(2); + if remainder >= half { + quotient.checked_add(if (numerator < i256::ZERO) != (denominator < i256::ZERO) { + i256::MINUS_ONE + } else { + i256::ONE + }) + } else { + Some(quotient) + } +} + +impl ScalarUDFImpl for DecimalArithmetic { + fn name(&self) -> &str { + self.function_signature + .split(':') + .next() + .unwrap_or(&self.function_signature) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(self.output_type()) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + Ok(Arc::new(Field::new( + self.name(), + self.output_type(), + args.arg_fields.iter().any(|f| f.is_nullable()), + ))) + } + + fn is_strict(&self) -> bool { + true + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let scalar = args + .args + .iter() + .all(|a| matches!(a, ColumnarValue::Scalar(_))); + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let left = as_primitive_array::(&arrays[0])?; + let right = as_primitive_array::(&arrays[1])?; + let values = left + .iter() + .zip(right.iter()) + .map(|(left, right)| match (left, right) { + (Some(left), Some(right)) => self.evaluate_value(left, right).map(Some), + _ => Ok(None), + }) + .collect::>>()?; + let result = Decimal128Array::from(values) + .with_precision_and_scale(self.precision, self.scale)?; + if scalar { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } else { + Ok(ColumnarValue::Array(Arc::new(result))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn function(op: Operator, scales: [i8; 2], output: (u8, i8)) -> DecimalArithmetic { + DecimalArithmetic::try_new( + "decimal_test:dec_dec", + op, + scales.map(|s| DataType::Decimal128(38, s)), + &DataType::Decimal128(output.0, output.1), + &[], + ) + .unwrap() + } + + #[test] + fn decimal_values_at_declared_scale() -> Result<()> { + use Operator::*; + // Raw integers encode value * 10^scale. Expected answers are rounded + // mathematical results, not results obtained from Arrow's kernels. + for (op, scales, output, left, right, expected) in [ + (Divide, [2, 1], (21, 8), 100, 30, 33_333_333), + (Divide, [2, 1], (21, 8), 200, 30, 66_666_667), + (Divide, [2, 1], (21, 8), -200, 30, -66_666_667), + (Divide, [2, 1], (21, 8), 200, -30, -66_666_667), + (Divide, [2, 1], (21, 8), -200, -30, 66_666_667), + (Divide, [0, 0], (3, 0), 5, 2, 3), + (Divide, [0, 0], (3, 0), -5, 2, -3), + (Divide, [0, 0], (3, 0), 4, 3, 1), + // The half threshold must work for odd denominators too. + (Divide, [0, 0], (3, 0), 5, 3, 2), + (Plus, [2, 2], (3, 1), 4, 4, 1), + (Plus, [2, 2], (3, 1), -4, -4, -1), + (Minus, [2, 2], (3, 1), 104, 96, 1), + (Multiply, [2, 2], (5, 2), 125, 250, 313), + (Multiply, [2, 2], (5, 2), -125, 250, -313), + (Modulo, [3, 1], (5, 2), 12_345, 20, 35), + (Modulo, [3, 1], (5, 2), -12_345, 20, -35), + // Precision-only changes and increasing scale are supported. + (Plus, [2, 1], (20, 2), 100, 30, 400), + (Multiply, [1, 1], (20, 8), 15, 20, 300_000_000), + ] { + assert_eq!( + function(op, scales, output).evaluate_value(left, right)?, + expected, + "{left} {op} {right}, scales={scales:?}, output={output:?}" + ); + } + Ok(()) + } + + #[test] + fn decimal_wide_intermediates() -> Result<()> { + // 0.1 / 0.1 = 1.000000. Using Arrow Decimal256 division at scale + // s1 + 4 would overflow while scaling the numerator for this example. + let tenth = 10_i128.pow(37); + assert_eq!( + function(Operator::Divide, [38, 38], (38, 6)).evaluate_value(tenth, tenth)?, + 1_000_000 + ); + assert_eq!( + function(Operator::Divide, [38, 0], (38, 6)).evaluate_value(tenth, 1)?, + 100_000 + ); + assert_eq!( + function(Operator::Multiply, [38, 38], (38, 6)) + .evaluate_value(tenth, tenth)?, + 10_000 + ); + // The exact sum exceeds i128, but fits after reducing scale. + assert_eq!( + function(Operator::Plus, [10, 10], (38, 9)) + .evaluate_value(9 * tenth, 9 * tenth)?, + 18 * 10_i128.pow(36) + ); + // Keep all digits before cancellation and before multiplication's + // final scale reduction. + assert_eq!( + function(Operator::Minus, [10, 10], (38, 9)) + .evaluate_value(9 * tenth, 9 * tenth - 6)?, + 1 + ); + assert_eq!( + function(Operator::Multiply, [10, 10], (38, 6)) + .evaluate_value(9 * tenth, 10_i128.pow(10))?, + 9 * 10_i128.pow(33) + ); + // Maximum scale adjustment, including zero numerators. + assert_eq!( + function(Operator::Divide, [0, 38], (38, 38)).evaluate_value(0, 1)?, + 0 + ); + Ok(()) + } + + #[test] + fn decimal_overflow_and_zero_divisors() { + for (op, scales, output, left, right) in [ + (Operator::Plus, [0, 0], (2, 0), 99, 1), + // Rounding can itself make the result overflow. + (Operator::Divide, [0, 0], (2, 0), 199, 2), + (Operator::Divide, [0, 0], (2, 0), -199, 2), + ( + Operator::Multiply, + [10, 10], + (38, 6), + 9 * 10_i128.pow(37), + 9 * 10_i128.pow(37), + ), + // The scaled numerator exceeds i256, and the final result really + // does exceed the declared precision as well. + (Operator::Divide, [0, 38], (38, 38), 10_i128.pow(37), 1), + ] { + let err = function(op, scales, output) + .evaluate_value(left, right) + .unwrap_err(); + assert!(err.to_string().contains("Decimal overflow"), "{err}"); + } + for op in [Operator::Divide, Operator::Modulo] { + let err = function(op, [2, 2], (21, 8)) + .evaluate_value(0, 0) + .unwrap_err(); + assert!(err.to_string().contains("Divide by zero"), "{err}"); + } + } + + #[test] + fn decimal_option_negotiation() { + for (name, preferences, supported) in [ + ("overflow", vec!["ERROR"], true), + ("OVERFLOW", vec!["SILENT", "error"], true), + ("overflow", vec!["SATURATE"], false), + ("overflow", vec!["SILENT"], false), + ("overflow", vec![], false), + ("rounding", vec!["HALF_UP"], false), + ] { + let result = DecimalArithmetic::try_new( + "divide:dec_dec", + Operator::Divide, + [DataType::Decimal128(10, 2), DataType::Decimal128(5, 1)], + &DataType::Decimal128(21, 8), + &[FunctionOption { + name: name.into(), + preference: preferences.into_iter().map(str::to_owned).collect(), + }], + ); + assert_eq!(result.is_ok(), supported, "{result:?}"); + } + } +} diff --git a/datafusion/substrait/src/logical_plan/mod.rs b/datafusion/substrait/src/logical_plan/mod.rs index 6f8b8e493f529..5ccd4f1228ba9 100644 --- a/datafusion/substrait/src/logical_plan/mod.rs +++ b/datafusion/substrait/src/logical_plan/mod.rs @@ -16,4 +16,5 @@ // under the License. pub mod consumer; +pub(crate) mod decimal; pub mod producer; diff --git a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs index 75720395aae7c..2dd3355b85ce8 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::logical_plan::decimal::DecimalArithmetic; use crate::logical_plan::producer::{ SubstraitProducer, to_substrait_literal_expr, to_substrait_type, }; @@ -36,6 +37,25 @@ pub fn from_scalar_function( schema: &DFSchemaRef, ) -> datafusion::common::Result { let (_, output_field) = Expr::ScalarFunction(fun.clone()).to_field(schema)?; + if let Some(decimal) = fun.func.inner().downcast_ref::() { + let mut expression = from_function( + producer, + decimal.function_signature(), + &fun.args, + output_field.data_type(), + output_field.is_nullable(), + schema, + )?; + if let Some(RexType::ScalarFunction(function)) = &mut expression.rex_type { + // Preserve the chosen overflow behavior even if the output type + // happens to match a native operator on the next import. + function.options = vec![substrait::proto::FunctionOption { + name: "overflow".into(), + preference: vec!["ERROR".into()], + }]; + } + return Ok(expression); + } from_function( producer, fun.name(), diff --git a/datafusion/substrait/tests/cases/consumer_integration.rs b/datafusion/substrait/tests/cases/consumer_integration.rs index 41ced2a740bc2..dbfaed49ef733 100644 --- a/datafusion/substrait/tests/cases/consumer_integration.rs +++ b/datafusion/substrait/tests/cases/consumer_integration.rs @@ -79,12 +79,18 @@ mod tests { #[tokio::test] async fn tpch_test_01() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(1).await.unwrap_err(); + let plan_str = tpch_plan_to_string(1).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 2): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" - ); + plan_str, + @r#" + Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, sum(LINEITEM.L_QUANTITY) AS SUM_QTY, sum(LINEITEM.L_EXTENDEDPRICE) AS SUM_BASE_PRICE, sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) AS SUM_DISC_PRICE, sum(multiply(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT),Int32(1) + LINEITEM.L_TAX)) AS SUM_CHARGE, avg(LINEITEM.L_QUANTITY) AS AVG_QTY, avg(LINEITEM.L_EXTENDEDPRICE) AS AVG_PRICE, avg(LINEITEM.L_DISCOUNT) AS AVG_DISC, count(Int64(1)) AS COUNT_ORDER + Sort: LINEITEM.L_RETURNFLAG ASC NULLS LAST, LINEITEM.L_LINESTATUS ASC NULLS LAST + Aggregate: groupBy=[[LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS]], aggr=[[sum(LINEITEM.L_QUANTITY), sum(LINEITEM.L_EXTENDEDPRICE), sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)), sum(multiply(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT),Int32(1) + LINEITEM.L_TAX)), avg(LINEITEM.L_QUANTITY), avg(LINEITEM.L_EXTENDEDPRICE), avg(LINEITEM.L_DISCOUNT), count(Int64(1))]] + Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, LINEITEM.L_QUANTITY, LINEITEM.L_EXTENDEDPRICE, multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT), multiply(multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT), CAST(Int32(1) AS Decimal128(15, 2)) + LINEITEM.L_TAX), LINEITEM.L_DISCOUNT + Filter: LINEITEM.L_SHIPDATE <= Date32("1998-12-01") - IntervalDayTime("IntervalDayTime { days: 0, milliseconds: 10368000 }") + TableScan: LINEITEM + "# + ); Ok(()) } @@ -125,12 +131,24 @@ mod tests { #[tokio::test] async fn tpch_test_03() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(3).await.unwrap_err(); + let plan_str = tpch_plan_to_string(3).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" - ); + plan_str, + @r#" + Projection: LINEITEM.L_ORDERKEY, sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) AS REVENUE, ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY + Limit: skip=0, fetch=10 + Sort: sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) DESC NULLS FIRST, ORDERS.O_ORDERDATE ASC NULLS LAST + Projection: LINEITEM.L_ORDERKEY, sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)), ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY + Aggregate: groupBy=[[LINEITEM.L_ORDERKEY, ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY]], aggr=[[sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT))]] + Projection: LINEITEM.L_ORDERKEY, ORDERS.O_ORDERDATE, ORDERS.O_SHIPPRIORITY, multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) + Filter: CUSTOMER.C_MKTSEGMENT = Utf8("BUILDING") AND CUSTOMER.C_CUSTKEY = ORDERS.O_CUSTKEY AND LINEITEM.L_ORDERKEY = ORDERS.O_ORDERKEY AND ORDERS.O_ORDERDATE < CAST(Utf8("1995-03-15") AS Date32) AND LINEITEM.L_SHIPDATE > CAST(Utf8("1995-03-15") AS Date32) + Cross Join: + Cross Join: + TableScan: LINEITEM + TableScan: CUSTOMER + TableScan: ORDERS + "# + ); Ok(()) } @@ -156,23 +174,43 @@ mod tests { #[tokio::test] async fn tpch_test_05() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(5).await.unwrap_err(); + let plan_str = tpch_plan_to_string(5).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" - ); + plan_str, + @r#" + Projection: NATION.N_NAME, sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) AS REVENUE + Sort: sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) DESC NULLS FIRST + Aggregate: groupBy=[[NATION.N_NAME]], aggr=[[sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT))]] + Projection: NATION.N_NAME, multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) + Filter: CUSTOMER.C_CUSTKEY = ORDERS.O_CUSTKEY AND LINEITEM.L_ORDERKEY = ORDERS.O_ORDERKEY AND LINEITEM.L_SUPPKEY = SUPPLIER.S_SUPPKEY AND CUSTOMER.C_NATIONKEY = SUPPLIER.S_NATIONKEY AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_REGIONKEY = REGION.R_REGIONKEY AND REGION.R_NAME = Utf8("ASIA") AND ORDERS.O_ORDERDATE >= CAST(Utf8("1994-01-01") AS Date32) AND ORDERS.O_ORDERDATE < CAST(Utf8("1995-01-01") AS Date32) + Cross Join: + Cross Join: + Cross Join: + Cross Join: + Cross Join: + TableScan: CUSTOMER + TableScan: ORDERS + TableScan: LINEITEM + TableScan: SUPPLIER + TableScan: NATION + TableScan: REGION + "# + ); Ok(()) } #[tokio::test] async fn tpch_test_06() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(6).await.unwrap_err(); + let plan_str = tpch_plan_to_string(6).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 6): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * LINEITEM.L_DISCOUNT derives Decimal128(31, 4); this conversion is unsupported" - ); + plan_str, + @r#" + Aggregate: groupBy=[[]], aggr=[[sum(multiply(LINEITEM.L_EXTENDEDPRICE,LINEITEM.L_DISCOUNT)) AS REVENUE]] + Projection: multiply(LINEITEM.L_EXTENDEDPRICE, LINEITEM.L_DISCOUNT) + Filter: LINEITEM.L_SHIPDATE >= CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-01-01") AS Date32) AND LINEITEM.L_DISCOUNT >= Decimal128(0.05,3,2) AND LINEITEM.L_DISCOUNT <= Decimal128(0.07,3,2) AND LINEITEM.L_QUANTITY < CAST(Int32(24) AS Decimal128(15, 2)) + TableScan: LINEITEM + "# + ); Ok(()) } @@ -202,23 +240,58 @@ mod tests { #[tokio::test] async fn tpch_test_10() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(10).await.unwrap_err(); + let plan_str = tpch_plan_to_string(10).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" - ); + plan_str, + @r#" + Projection: CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) AS REVENUE, CUSTOMER.C_ACCTBAL, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_PHONE, CUSTOMER.C_COMMENT + Limit: skip=0, fetch=20 + Sort: sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) DESC NULLS FIRST + Projection: CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)), CUSTOMER.C_ACCTBAL, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_PHONE, CUSTOMER.C_COMMENT + Aggregate: groupBy=[[CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, CUSTOMER.C_ACCTBAL, CUSTOMER.C_PHONE, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_COMMENT]], aggr=[[sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT))]] + Projection: CUSTOMER.C_CUSTKEY, CUSTOMER.C_NAME, CUSTOMER.C_ACCTBAL, CUSTOMER.C_PHONE, NATION.N_NAME, CUSTOMER.C_ADDRESS, CUSTOMER.C_COMMENT, multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) + Filter: CUSTOMER.C_CUSTKEY = ORDERS.O_CUSTKEY AND LINEITEM.L_ORDERKEY = ORDERS.O_ORDERKEY AND ORDERS.O_ORDERDATE >= CAST(Utf8("1993-10-01") AS Date32) AND ORDERS.O_ORDERDATE < CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_RETURNFLAG = Utf8("R") AND CUSTOMER.C_NATIONKEY = NATION.N_NATIONKEY + Cross Join: + Cross Join: + Cross Join: + TableScan: CUSTOMER + TableScan: ORDERS + TableScan: LINEITEM + TableScan: NATION + "# + ); Ok(()) } #[tokio::test] async fn tpch_test_11() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(11).await.unwrap_err(); + let plan_str = tpch_plan_to_string(11).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 2): declared Decimal128(19, 2), but native expression PARTSUPP.PS_SUPPLYCOST * CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0)) derives Decimal128(35, 2); this conversion is unsupported" - ); + plan_str, + @r#" + Projection: PARTSUPP.PS_PARTKEY, sum(multiply(PARTSUPP.PS_SUPPLYCOST,PARTSUPP.PS_AVAILQTY)) AS value + Sort: sum(multiply(PARTSUPP.PS_SUPPLYCOST,PARTSUPP.PS_AVAILQTY)) DESC NULLS FIRST + Filter: sum(multiply(PARTSUPP.PS_SUPPLYCOST,PARTSUPP.PS_AVAILQTY)) > () + Subquery: + Projection: multiply(sum(multiply(PARTSUPP.PS_SUPPLYCOST,PARTSUPP.PS_AVAILQTY)), Decimal128(0.0001000000,11,10)) + Aggregate: groupBy=[[]], aggr=[[sum(multiply(PARTSUPP.PS_SUPPLYCOST,PARTSUPP.PS_AVAILQTY))]] + Projection: multiply(PARTSUPP.PS_SUPPLYCOST, CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0))) + Filter: PARTSUPP.PS_SUPPKEY = SUPPLIER.S_SUPPKEY AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8("JAPAN") + Cross Join: + Cross Join: + TableScan: PARTSUPP + TableScan: SUPPLIER + TableScan: NATION + Aggregate: groupBy=[[PARTSUPP.PS_PARTKEY]], aggr=[[sum(multiply(PARTSUPP.PS_SUPPLYCOST,PARTSUPP.PS_AVAILQTY))]] + Projection: PARTSUPP.PS_PARTKEY, multiply(PARTSUPP.PS_SUPPLYCOST, CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0))) + Filter: PARTSUPP.PS_SUPPKEY = SUPPLIER.S_SUPPKEY AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8("JAPAN") + Cross Join: + Cross Join: + TableScan: PARTSUPP + TableScan: SUPPLIER + TableScan: NATION + "# + ); Ok(()) } @@ -263,12 +336,19 @@ mod tests { #[tokio::test] async fn tpch_test_14() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(14).await.unwrap_err(); + let plan_str = tpch_plan_to_string(14).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 5): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" - ); + plan_str, + @r#" + Projection: divide(multiply(Decimal128(100.00,5,2), sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT) ELSE Decimal128(0.0000,19,4) END)), sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT))) AS PROMO_REVENUE + Aggregate: groupBy=[[]], aggr=[[sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT) ELSE Decimal128(0.0000,19,4) END), sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT))]] + Projection: CASE WHEN PART.P_TYPE LIKE CAST(Utf8("PROMO%") AS Utf8) THEN multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) ELSE Decimal128(0.0000,19,4) END, multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) + Filter: LINEITEM.L_PARTKEY = PART.P_PARTKEY AND LINEITEM.L_SHIPDATE >= Date32("1995-09-01") AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-10-01") AS Date32) + Cross Join: + TableScan: LINEITEM + TableScan: PART + "# + ); Ok(()) } @@ -305,12 +385,25 @@ mod tests { #[tokio::test] async fn tpch_test_17() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(17).await.unwrap_err(); + let plan_str = tpch_plan_to_string(17).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 4): declared Decimal128(17, 3), but native expression Decimal128(0.2,2,1) * avg(LINEITEM.L_QUANTITY) derives Decimal128(22, 7); this conversion is unsupported" - ); + plan_str, + @r#" + Projection: divide(sum(LINEITEM.L_EXTENDEDPRICE), Decimal128(7.0,2,1)) AS AVG_YEARLY + Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE)]] + Projection: LINEITEM.L_EXTENDEDPRICE + Filter: PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#23") AND PART.P_CONTAINER = Utf8("MED BOX") AND LINEITEM.L_QUANTITY < () + Subquery: + Projection: multiply(Decimal128(0.2,2,1), avg(LINEITEM.L_QUANTITY)) + Aggregate: groupBy=[[]], aggr=[[avg(LINEITEM.L_QUANTITY)]] + Projection: LINEITEM.L_QUANTITY + Filter: LINEITEM.L_PARTKEY = outer_ref(PART.P_PARTKEY) + TableScan: LINEITEM + Cross Join: + TableScan: LINEITEM + TableScan: PART + "# + ); Ok(()) } @@ -343,23 +436,49 @@ mod tests { } #[tokio::test] async fn tpch_test_19() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(19).await.unwrap_err(); + let plan_str = tpch_plan_to_string(19).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 6): declared Decimal128(19, 4), but native expression LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) derives Decimal128(32, 4); this conversion is unsupported" - ); + plan_str, + @r#" + Aggregate: groupBy=[[]], aggr=[[sum(multiply(LINEITEM.L_EXTENDEDPRICE,Int32(1) - LINEITEM.L_DISCOUNT)) AS REVENUE]] + Projection: multiply(LINEITEM.L_EXTENDEDPRICE, CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) + Filter: PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#12") AND (PART.P_CONTAINER = CAST(Utf8("SM CASE") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("SM BOX") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("SM PACK") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("SM PKG") AS Utf8)) AND LINEITEM.L_QUANTITY >= CAST(Int32(1) AS Decimal128(15, 2)) AND LINEITEM.L_QUANTITY <= CAST(Int32(1) + Int32(10) AS Decimal128(15, 2)) AND PART.P_SIZE >= Int32(1) AND PART.P_SIZE <= Int32(5) AND (LINEITEM.L_SHIPMODE = CAST(Utf8("AIR") AS Utf8) OR LINEITEM.L_SHIPMODE = CAST(Utf8("AIR REG") AS Utf8)) AND LINEITEM.L_SHIPINSTRUCT = Utf8("DELIVER IN PERSON") OR PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#23") AND (PART.P_CONTAINER = CAST(Utf8("MED BAG") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("MED BOX") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("MED PKG") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("MED PACK") AS Utf8)) AND LINEITEM.L_QUANTITY >= CAST(Int32(10) AS Decimal128(15, 2)) AND LINEITEM.L_QUANTITY <= CAST(Int32(10) + Int32(10) AS Decimal128(15, 2)) AND PART.P_SIZE >= Int32(1) AND PART.P_SIZE <= Int32(10) AND (LINEITEM.L_SHIPMODE = CAST(Utf8("AIR") AS Utf8) OR LINEITEM.L_SHIPMODE = CAST(Utf8("AIR REG") AS Utf8)) AND LINEITEM.L_SHIPINSTRUCT = Utf8("DELIVER IN PERSON") OR PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#34") AND (PART.P_CONTAINER = CAST(Utf8("LG CASE") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("LG BOX") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("LG PACK") AS Utf8) OR PART.P_CONTAINER = CAST(Utf8("LG PKG") AS Utf8)) AND LINEITEM.L_QUANTITY >= CAST(Int32(20) AS Decimal128(15, 2)) AND LINEITEM.L_QUANTITY <= CAST(Int32(20) + Int32(10) AS Decimal128(15, 2)) AND PART.P_SIZE >= Int32(1) AND PART.P_SIZE <= Int32(15) AND (LINEITEM.L_SHIPMODE = CAST(Utf8("AIR") AS Utf8) OR LINEITEM.L_SHIPMODE = CAST(Utf8("AIR REG") AS Utf8)) AND LINEITEM.L_SHIPINSTRUCT = Utf8("DELIVER IN PERSON") + Cross Join: + TableScan: LINEITEM + TableScan: PART + "# + ); Ok(()) } #[tokio::test] async fn tpch_test_20() -> Result<()> { - // This fixture declares a decimal result type the native operator cannot provide. - let err = tpch_plan_to_string(20).await.unwrap_err(); + let plan_str = tpch_plan_to_string(20).await?; assert_snapshot!( - err, - @"Substrait error: Decimal return type mismatch for multiply:dec_dec (function reference 7): declared Decimal128(17, 3), but native expression Decimal128(0.5,2,1) * sum(LINEITEM.L_QUANTITY) derives Decimal128(28, 3); this conversion is unsupported" - ); + plan_str, + @r#" + Sort: SUPPLIER.S_NAME ASC NULLS LAST + Projection: SUPPLIER.S_NAME, SUPPLIER.S_ADDRESS + Filter: SUPPLIER.S_SUPPKEY IN () AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8("CANADA") + Subquery: + Projection: PARTSUPP.PS_SUPPKEY + Filter: PARTSUPP.PS_PARTKEY IN () AND CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0)) > () + Subquery: + Projection: PART.P_PARTKEY + Filter: PART.P_NAME LIKE CAST(Utf8("forest%") AS Utf8) + TableScan: PART + Subquery: + Projection: multiply(Decimal128(0.5,2,1), sum(LINEITEM.L_QUANTITY)) + Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_QUANTITY)]] + Projection: LINEITEM.L_QUANTITY + Filter: LINEITEM.L_PARTKEY = outer_ref(PARTSUPP.PS_PARTKEY) AND LINEITEM.L_SUPPKEY = outer_ref(PARTSUPP.PS_SUPPKEY) AND LINEITEM.L_SHIPDATE >= CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-01-01") AS Date32) + TableScan: LINEITEM + TableScan: PARTSUPP + Cross Join: + TableScan: SUPPLIER + TableScan: NATION + "# + ); Ok(()) } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 5ee5c90091fc5..bdb24cd6c5f57 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -539,12 +539,28 @@ async fn decimal_arithmetic_output_type_contract() -> Result<()> { }; decimal.precision = declared.0; decimal.scale = declared.1; - let result = from_substrait_plan(&ctx.state(), &proto).await; - if native == (declared.0 as u8, declared.1 as i8) { - assert_eq!(result?.schema().field(0).data_type(), &expected_type); - } else { - let err = result.unwrap_err().to_string(); - assert!(err.contains("Decimal return type mismatch"), "{err}"); + let result = from_substrait_plan(&ctx.state(), &proto).await?; + let declared = (declared.0 as u8, declared.1 as i8); + assert_eq!( + result.schema().field(0).data_type(), + &DataType::Decimal128(declared.0, declared.1) + ); + // Serialize and import again, preserving the declared type and the + // decimal function's chosen overflow behavior. + let exported = to_substrait_plan(&result, &ctx.state())?; + let reimported = from_substrait_plan(&ctx.state(), &exported).await?; + for plan in [result, reimported] { + let batches = ctx.execute_logical_plan(plan).await?.collect().await?; + let value = match op { + "+" => 4 * 10_i128.pow(declared.1 as u32), + "*" => 3 * 10_i128.pow(declared.1 as u32), + "/" => 33_333_333, + _ => unreachable!(), + }; + assert_eq!( + ScalarValue::try_from_array(batches[0].column(0), 0)?, + ScalarValue::Decimal128(Some(value), declared.0, declared.1) + ); } } Ok(())