Skip to content

refactor: use Arrow cast for decimal rescale check #5444

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
peterxcli wants to merge 1 commit into from
Closed

refactor: use Arrow cast for decimal rescale check #5444

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs/source/contributor-guide/expression-audits/math_funcs.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,9 +121,9 @@ Internal decimal wrapper emitted around every decimal `+ - * /`, `sum`, and `avg

## DecimalRescaleCheckOverflow (internal)

Internal fused expression that rescales a Decimal128 value (changing scale) and checks output precision in one pass, replacing the `CheckOverflow(Cast(expr, Decimal128(p, s)))` pattern used by decimal-to-decimal casts. Native impl: `math_funcs/internal/decimal_rescale_check.rs`.
Internal fused expression that rescales a Decimal128 value and checks output precision with Arrow's decimal cast kernel, replacing the `CheckOverflow(Cast(expr, Decimal128(p, s)))` pattern used by decimal-to-decimal casts. Native impl: `math_funcs/internal/decimal_rescale_check.rs`.

- Performance (tuned 2026-07-15, PR [#4938](https://github.com/apache/datafusion-comet/pull/4938)): the legacy path ran `null_if_overflow_precision` (a second full pass that allocates a new array) on every batch to turn overflow sentinels into nulls, even when nothing overflowed. Now that pass runs only when a sentinel is present (`contains(&i128::MAX)`, short-circuiting), so the common no-overflow case skips the allocation. 8 to 26% faster on no-overflow shapes; overflow and ANSI shapes unchanged. Benchmark: `benches/decimal_rescale.rs`.
- Implementation (updated 2026-08-04, issue [#5094](https://github.com/apache/datafusion-comet/issues/5094)): decimal casts delegate rescaling, HALF_UP rounding, and precision validation to Arrow's decimal cast kernel. Legacy overflow becomes null and ANSI overflow raises Spark's `NUMERIC_VALUE_OUT_OF_RANGE`; scale changes beyond Arrow's Decimal128 power table preserve Spark's value-sensitive zero, null, and overflow behavior. Benchmark: `benches/decimal_rescale.rs`.

## e

Expand Down
137 changes: 119 additions & 18 deletions native/core/src/execution/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@ use datafusion::physical_expr::expressions::{Literal, StatsType};
use datafusion::physical_expr::window::WindowExpr;
use datafusion::physical_expr::LexOrdering;

use crate::execution::expressions::arithmetic::CheckedBinaryExpr;
use crate::parquet::parquet_exec::init_datasource_exec;
use arrow::array::{
new_empty_array, Array, ArrayRef, BinaryBuilder, BooleanArray, Date32Array, Decimal128Array,
Expand Down Expand Up @@ -602,6 +603,10 @@ impl PhysicalPlanner {
self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?;
let data_type = to_arrow_datatype(expr.datatype.as_ref().unwrap());
let fail_on_error = expr.fail_on_error;
let query_context = spark_expr.expr_id.and_then(|expr_id| {
let registry = &self.query_context_registry;
registry.get(expr_id)
});

// WideDecimalBinaryExpr already handles overflow — skip redundant check
// but only if its output type matches CheckOverflow's declared type
Expand All @@ -622,23 +627,23 @@ impl PhysicalPlanner {
{
let cast_target = cast.data_type(&input_schema)?;
if cast_target == data_type {
return Ok(Arc::new(DecimalRescaleCheckOverflow::new(
Arc::clone(&cast.child),
s_in,
*p_out,
*s_out,
fail_on_error,
)));
let fused: Arc<dyn PhysicalExpr> =
Arc::new(DecimalRescaleCheckOverflow::new(
Arc::clone(&cast.child),
s_in,
*p_out,
*s_out,
fail_on_error,
));
return if query_context.is_some() {
Ok(Arc::new(CheckedBinaryExpr::new(fused, query_context)))
} else {
Ok(fused)
};
}
}
}

// Look up query context from registry if expr_id is present
let query_context = spark_expr.expr_id.and_then(|expr_id| {
let registry = &self.query_context_registry;
registry.get(expr_id)
});

Ok(Arc::new(CheckOverflow::new(
child,
data_type,
Expand Down Expand Up @@ -1121,8 +1126,6 @@ impl PhysicalPlanner {
Arc::new(ConfigOptions::default()),
));

// Wrap with CheckedBinaryExpr to add query_context to errors
use crate::execution::expressions::arithmetic::CheckedBinaryExpr;
Ok(Arc::new(CheckedBinaryExpr::new(scalar_expr, query_context)))
} else {
Ok(Arc::new(BinaryExpr::new(left, op, right)))
Expand Down Expand Up @@ -4598,7 +4601,8 @@ mod tests {
use std::{sync::Arc, task::Poll};

use arrow::array::{
Array, DictionaryArray, Int32Array, Int8Array, ListArray, RecordBatch, StringArray,
Array, Decimal128Array, DictionaryArray, Int32Array, Int8Array, ListArray, RecordBatch,
StringArray,
};
use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema};
use datafusion::catalog::memory::DataSourceExec;
Expand All @@ -4616,7 +4620,9 @@ mod tests {
use tempfile::TempDir;
use tokio::sync::mpsc;

use crate::execution::{operators::InputBatch, planner::PhysicalPlanner};
use crate::execution::{
expressions::arithmetic::CheckedBinaryExpr, operators::InputBatch, planner::PhysicalPlanner,
};

use crate::execution::operators::ExecutionError;
use crate::execution::planner::literal_to_array_ref;
Expand All @@ -4632,7 +4638,7 @@ mod tests {
spark_operator,
spark_operator::{operator::OpStruct, Operator},
};
use datafusion_comet_spark_expr::EvalMode;
use datafusion_comet_spark_expr::{EvalMode, SparkError, SparkErrorWithContext};

#[test]
fn test_unpack_dictionary_primitive() {
Expand Down Expand Up @@ -4807,6 +4813,101 @@ mod tests {
assert_eq!(comet_err.to_string(), "Error from DataFusion: exec error.");
}

#[test]
fn fused_decimal_overflow_keeps_query_context() {
use spark_expression::data_type::{
data_type_info::DatatypeStruct, DataTypeInfo, DecimalInfo,
};

fn decimal_type(precision: i32, scale: i32) -> spark_expression::DataType {
spark_expression::DataType {
type_id: 10,
type_info: Some(Box::new(DataTypeInfo {
datatype_struct: Some(DatatypeStruct::Decimal(DecimalInfo {
precision,
scale,
})),
})),
}
}

let input_type = decimal_type(3, 1);
let output_type = decimal_type(2, 1);
let bound = Expr {
expr_struct: Some(Bound(spark_expression::BoundReference {
index: 0,
datatype: Some(input_type),
})),
query_context: None,
expr_id: None,
};
let cast = Expr {
expr_struct: Some(ExprStruct::Cast(Box::new(spark_expression::Cast {
child: Some(Box::new(bound)),
datatype: Some(output_type.clone()),
timezone: "UTC".to_string(),
eval_mode: spark_expression::EvalMode::Ansi as i32,
allow_incompat: false,
is_spark4_plus: true,
}))),
query_context: None,
expr_id: None,
};
let sql = "SELECT CAST(10.1 AS DECIMAL(2, 1))";
let check_overflow = Expr {
expr_struct: Some(ExprStruct::CheckOverflow(Box::new(
spark_expression::CheckOverflow {
child: Some(Box::new(cast)),
datatype: Some(output_type),
fail_on_error: true,
},
))),
query_context: Some(spark_expression::QueryContext {
sql_text: sql.to_string(),
start_index: 7,
stop_index: 34,
object_type: None,
object_name: None,
line: 1,
start_position: 7,
sql_text_idx: None,
}),
expr_id: Some(1),
};

let schema = Arc::new(Schema::new(vec![Field::new(
"value",
DataType::Decimal128(3, 1),
false,
)]));
let expr = PhysicalPlanner::default()
.create_expr(&check_overflow, Arc::clone(&schema))
.unwrap();
assert!(expr.downcast_ref::<CheckedBinaryExpr>().is_some());

let values = Decimal128Array::from(vec![101])
.with_precision_and_scale(3, 1)
.unwrap();
let batch = RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap();
let DataFusionError::External(error) = expr.evaluate(&batch).unwrap_err() else {
panic!("expected external Spark error")
};
let error = error
.downcast_ref::<SparkErrorWithContext>()
.expect("expected SparkErrorWithContext");
assert!(matches!(
&error.error,
SparkError::NumericValueOutOfRange {
value,
precision: 2,
scale: 1,
} if value == "10.1"
));
let context = error.context.as_ref().expect("expected query context");
assert_eq!(context.sql_text.as_str(), sql);
assert_eq!((context.start_index, context.stop_index), (7, 34));
}

// Creates a filter operator which takes an `Int32Array` and selects rows that are equal to
// `value`.
fn create_filter(child_op: spark_operator::Operator, value: i32) -> spark_operator::Operator {
Expand Down
Loading
Loading