Skip to content

refactor: use arity helper for Int to Decimal128 reinterpretation #5193

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

Open
0lai0 wants to merge 1 commit into
base: main
Choose a base branch
from
+69 −24
Open

refactor: use arity helper for Int to Decimal128 reinterpretation #5193

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
26 changes: 9 additions & 17 deletions native/core/src/execution/columnar_to_row.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,10 +37,11 @@

use crate::errors::{CometError, CometResult};
use arrow::array::types::{
ArrowDictionaryKeyType, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type,
UInt64Type, UInt8Type,
ArrowDictionaryKeyType, Decimal128Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type,
UInt32Type, UInt64Type, UInt8Type,
};
use arrow::array::*;
use arrow::compute::kernels::arity::unary;
use arrow::compute::{cast_with_options, CastOptions};
use arrow::datatypes::{ArrowNativeType, DataType, TimeUnit};
use std::sync::Arc;
Expand Down Expand Up @@ -1055,14 +1056,10 @@ impl ColumnarToRowContext {
// Parquet stores small-precision decimals as Int32 for efficiency, and the
// reader may surface them as the physical Int32 type. The value is already
// scaled (e.g., -1 means -0.01 for scale 2). Reinterpret (not cast) to
// Decimal128 preserving the value.
let int_array = array.as_any().downcast_ref::<Int32Array>().ok_or_else(|| {
CometError::Internal("Failed to downcast to Int32Array".to_string())
})?;
let decimal_array: Decimal128Array = int_array
.iter()
.map(|v| v.map(|x| x as i128))
.collect::<Decimal128Array>()
// Decimal128 preserving the value. `arity::unary` widens the value and reuses
// the input null buffer zero-copy (an Arrow cast would rescale the value).
let int_array = array.as_primitive::<Int32Type>();
let decimal_array = unary::<_, _, Decimal128Type>(int_array, |x| x as i128)

@andygrove andygrove Aug 1, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This reads well, and the note about an Arrow cast rescaling the value is worth having in the code. I confirmed that dropping the downcast_ref + ok_or_else is safe, since the match is on array.data_type(), so nothing new can panic here.

The substance of the change is null-buffer handling, and test_convert_int32_to_decimal128 and test_convert_int64_to_decimal128 both use all-valid arrays, so that is the one thing the tests do not exercise. Could you add a null to those arrays? A sliced input would be good too, following test_map_data_conversion_sliced_maparray further down the file. I checked both cases by hand and they behave correctly, so this is just about locking the behaviour in.

.with_precision_and_scale(*precision, *scale)
.map_err(|e| {
CometError::Internal(format!("Invalid decimal precision/scale: {}", e))
Expand All @@ -1071,13 +1068,8 @@ impl ColumnarToRowContext {
}
(DataType::Int64, DataType::Decimal128(precision, scale)) => {
// Same as Int32 but for medium-precision decimals stored as Int64.
let int_array = array.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
CometError::Internal("Failed to downcast to Int64Array".to_string())
})?;
let decimal_array: Decimal128Array = int_array
.iter()
.map(|v| v.map(|x| x as i128))
.collect::<Decimal128Array>()
let int_array = array.as_primitive::<Int64Type>();
let decimal_array = unary::<_, _, Decimal128Type>(int_array, |x| x as i128)
.with_precision_and_scale(*precision, *scale)
.map_err(|e| {
CometError::Internal(format!("Invalid decimal precision/scale: {}", e))
Expand Down
67 changes: 60 additions & 7 deletions native/spark-expr/src/math_funcs/internal/make_decimal.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,12 @@

use crate::error::decimal_overflow_error;
use crate::math_funcs::utils::get_precision_scale;
use arrow::compute::kernels::arity::try_unary;
use arrow::datatypes::DataType;
use arrow::error::ArrowError;
use arrow::{
array::{AsArray, Decimal128Builder},
datatypes::{validate_decimal_precision, Int64Type},
array::{AsArray, Decimal128Array},
datatypes::{validate_decimal_precision, Decimal128Type, Int64Type},
};
use datafusion::common::{internal_err, DataFusionError, Result as DataFusionResult, ScalarValue};
use datafusion::physical_plan::ColumnarValue;
Expand All @@ -45,14 +47,41 @@ pub fn spark_make_decimal(
ColumnarValue::Array(a) => match a.data_type() {
DataType::Int64 => {
let arr = a.as_primitive::<Int64Type>();
let mut result = Decimal128Builder::new();
for v in arr.into_iter() {
result.append_option(long_to_decimal(v, precision, scale, fail_on_error)?)
}
let result_type = DataType::Decimal128(precision, scale);
// The Int64 is already the unscaled Decimal128 value, so we only reinterpret
// the bits (an Arrow Int64->Decimal cast would rescale the value). Both arity
// helpers reuse the input null buffer and only invoke the closure on valid rows.
let result: Decimal128Array = if fail_on_error {

@andygrove andygrove Aug 1, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A question on the helper choice here. Arrow's own docs on try_unary note it "is often significantly slower than unary" because LLVM cannot vectorize fallible closures, and unary_opt has the same per-row shape. Meanwhile checkoverflow.rs, in this same directory and solving the same overflow problem, already uses a different idiom: widen unconditionally, then a short-circuiting is_valid_decimal_precision scan, paying for null-masking or error construction only when an overflow actually exists.

I benchmarked the three shapes on 8192-row Int64 batches at Decimal128(18, 2), the widest precision DecimalAggregates produces for MakeDecimal (times are non-ANSI / ANSI):

shape no nulls sparse nulls dense nulls
builder loop (today) 19.0 / 19.1 µs 22.7 / 22.7 µs 19.1 / 19.2 µs
this PR 14.3 / 14.3 µs 15.8 / 14.9 µs 10.7 / 9.8 µs
unary + scan 5.58 / 8.50 µs 5.44 / 7.92 µs 5.48 / 9.04 µs

Your version is a genuine 1.3x to 1.9x win over what is there now. The checkoverflow.rs idiom looks like another 1.7x to 2.6x on top of that. MakeDecimal sits on the decimal sum and avg path, which is why its sibling got optimized in the first place, so it is worth getting right here.

Would you be up for trying that shape and seeing whether you reproduce the gain? The scan still finds the first offending value, so your new ordering test should keep passing. If you would rather stay with unary_opt/try_unary, could you note why in the PR description so the choice is on the record?

Either way, would you mind adding a make_decimal bench? benches/unscaled_value.rs and benches/check_overflow.rs cover the neighbouring expressions and would be easy to mirror. That would make the numbers checkable rather than something a reviewer has to measure by hand.

// ANSI mode: overflow is a hard error. `try_unary` surfaces the closure's
// ArrowError; unwrap the ExternalError back to DataFusionError::External so
// the ANSI error variant (not a generic ArrowError) is preserved.
try_unary::<Int64Type, _, Decimal128Type>(arr, |v| {

@andygrove andygrove Aug 1, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PrimitiveArray::try_unary is generic over the error type. It is only the free function arity::try_unary that pins it to ArrowError. If you call the method form, as you already do for unary_opt just below, the closure can return DataFusionError directly:

arr.try_unary::<_, Decimal128Type, DataFusionError>(|v| {
    let v = v as i128;
    validate_decimal_precision(v, precision, scale)
        .map(|()| v)
        .map_err(|_| {
            DataFusionError::External(Box::new(decimal_overflow_error(v, precision, scale)))
        })
})?

I compiled and ran this and it produces the identical error, same variant and same SparkError payload. It drops the arrow::error::ArrowError import, the map_err match, and the other => arm that cannot currently be reached. It also removes the chance that a later edit loses the unwrap and turns the ANSI failure into a generic ArrowError, which would break the NUMERIC_VALUE_OUT_OF_RANGE mapping. This may become moot if you take the scan approach in my other comment.

let v = v as i128;
validate_decimal_precision(v, precision, scale)
.map(|()| v)
.map_err(|_| {
ArrowError::ExternalError(Box::new(decimal_overflow_error(
v, precision, scale,
)))
})
})
.map_err(|e| match e {
ArrowError::ExternalError(inner) => DataFusionError::External(inner),
other => DataFusionError::from(other),
})?
} else {
// Non-ANSI: overflow becomes null. `unary_opt` applies the closure only to
// valid rows and marks a row null wherever the closure returns None.
arr.unary_opt::<_, Decimal128Type>(|v| {

@andygrove andygrove Aug 1, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

long_to_decimal is still here for the scalar path, and the array path now inlines the same validate-and-convert rule twice more. That is three copies with nothing keeping them in sync. Could one helper cover all three? Something like:

#[inline]
fn to_unscaled(v: i64, precision: u8, scale: i8) -> Option<i128> {
    let v = v as i128;
    validate_decimal_precision(v, precision, scale).ok().map(|()| v)
}

Then the non-ANSI array path becomes arr.unary_opt::<_, Decimal128Type>(|v| to_unscaled(v, precision, scale)), the ANSI path becomes to_unscaled(...).ok_or_else(|| overflow_err(v, precision, scale)), and long_to_decimal builds on the same two pieces.

let v = v as i128;
validate_decimal_precision(v, precision, scale)
.ok()
.map(|()| v)
})
};

Ok(ColumnarValue::Array(Arc::new(
result.finish().with_data_type(result_type),
result.with_data_type(result_type),
)))
}
av => internal_err!("Expected Int64 but found {av:?}"),
Expand Down Expand Up @@ -106,6 +135,30 @@ mod tests {
);
}

#[test]
fn test_array_overflow_reports_first_offending_value() {
// Two distinct overflowing values with a valid value in front. Both the original
// per-row `?` loop and `try_unary` walk rows in index order, so the reported error
// must reference the FIRST overflow (111111), not the later one (222222). This locks
// the ordering semantics rather than merely "some error occurred".
let args = [ColumnarValue::Array(Arc::new(Int64Array::from(vec![
Some(99),
Some(111111),
Some(222222),
])))];
let err = spark_make_decimal(&args, &DataType::Decimal128(3, 0), true)
.expect_err("overflow should error when fail_on_error is set");
let msg = err.to_string();
assert!(
msg.contains("111111"),
"should report first overflow: {msg}"
);
assert!(
!msg.contains("222222"),
"should not report later overflow: {msg}"
);
}

#[test]
fn test_array_overflow_nulls_when_not_fail_on_error() {
let result = spark_make_decimal(&overflow_args(), &DataType::Decimal128(3, 0), false)
Expand Down
Loading