perf: Optimize some decimal expressions by andygrove · Pull Request #3619 · apache/datafusion-comet

andygrove · 2026-03-02T17:53:38Z

Summary

For some decimal operations, we cast the inputs to Decimal256, perform the math operation, and then cast the result back to Decimal128:

let left = Arc::new(Cast::new(
    left,
    DataType::Decimal256(p1, s1),
    SparkCastOptions::new_without_timezone(EvalMode::Legacy, false),
));
let right = Arc::new(Cast::new(
    right,
    DataType::Decimal256(p2, s2),
    SparkCastOptions::new_without_timezone(EvalMode::Legacy, false),
));
let child = Arc::new(BinaryExpr::new(left, op, right));
Ok(Arc::new(Cast::new(
    child,
    data_type,
    SparkCastOptions::new_without_timezone(EvalMode::Legacy, false),
)))

This causes intermediate allocations and can be optimized with a custom WideDecimalBinaryExpr that performs i256 register arithmetic directly, reducing per-batch allocation from 4 intermediate arrays (3 Decimal256 @ 32 bytes/elem + 1 Decimal128 @ 16 bytes/elem = 112 bytes/elem) to 1 output array (16 bytes/elem)

TPC-H q1

Before:

        "durations": [
            8.183,
            4.728,
            4.597,
            4.562,
            4.538
        ],

After:

        "durations": [
            7.874,
            4.327,
            4.185,
            4.135,
            4.069
        ],

Criterion benchmark results (8192 element batches)

Case	Old	Fused	Speedup
add (same scale)	171 µs	57 µs	3.0x
add (diff scale)	173 µs	57 µs	3.0x
multiply	361 µs	305 µs	1.2x
subtract	173 µs	58 µs	3.0x

How it works

WideDecimalBinaryExpr evaluates left/right children, performs add/sub/mul using i256 intermediates via arrow::compute::kernels::arity::try_binary, applies scale adjustment with HALF_UP rounding, checks precision bounds, and outputs a single Decimal128 array. Follows the same pattern as decimal_div in div.rs.

Overflow handling matches existing behavior:

Ansi mode: returns ArrowError::ComputeError
Legacy/Try mode: uses i128::MAX sentinel + null_if_overflow_precision

Replace the 4-node expression tree (Cast→BinaryExpr→Cast→Cast) used for Decimal128 arithmetic that may overflow with a single fused expression that performs i256 register arithmetic directly. This reduces per-batch allocation from 4 intermediate arrays (112 bytes/elem) to 1 output array (16 bytes/elem). The new WideDecimalBinaryExpr evaluates children, performs add/sub/mul using i256 intermediates via try_binary, applies scale adjustment with HALF_UP rounding, checks precision bounds, and outputs a single Decimal128 array. Follows the same pattern as decimal_div.

Add benchmark comparing old Cast->BinaryExpr->Cast chain vs fused WideDecimalBinaryExpr for Decimal128 add/sub/mul. Covers four cases: add with same scale, add with different scales, multiply, and subtract.

andygrove · 2026-03-02T17:55:52Z

@sqlbenchmark run tpch --iterations 3

andygrove · 2026-03-02T18:44:22Z

@sqlbenchmark run tpch --iterations 3

Eliminate redundant CheckOverflow when wrapping WideDecimalBinaryExpr (which already handles overflow). Fuse Cast(Decimal128→Decimal128) + CheckOverflow into a single DecimalRescaleCheckOverflow expression that rescales and validates precision in one pass.

andygrove · 2026-03-02T20:49:24Z

@sqlbenchmark run tpch --iterations 3

sqlbenchmark · 2026-03-02T21:47:04Z

Benchmark job comet-pr-3619-c3986837690 failed due to an error.

coderfender · 2026-03-04T15:48:57Z

This is awesome !

parthchandra · 2026-03-05T00:12:05Z

native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs

+                let rescale_divisor = if need_rescale {
+                    i256_pow10((max_scale - s_out) as u32)
+                } else {
+                    i256::ONE


Don't we need to scale up to s_out if s_out > max_scale?

@andygrove did you miss this comment?

It is very possible. I will take another look on Monday. I will move to draft for now to avoid accidental merge.

native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs

native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs

native/core/src/execution/planner.rs

native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs

native/core/src/execution/planner.rs

native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs

native/spark-expr/src/math_funcs/wide_decimal_binary_expr.rs

- Handle scale-up when s_out > max(s1, s2) in add/subtract - Propagate errors in scalar path when fail_on_error=true - Guard against large scale delta (>38) overflow in rescale - Assert precision <= 38 in precision_bound - Assert exp <= 76 in i256_pow10 - Remove unnecessary _ prefix on used variables in planner - Use value.signum() instead of manual sign check - Verify Cast target type matches before fusing with CheckOverflow - Validate children count in with_new_children for both expressions - Add tests for scale-up, scalar error propagation, and large delta

andygrove · 2026-03-05T22:34:07Z

Thanks for the review @martin-g and @parthchandra. I pushed a commit to address this feedback and also added some additional tests.

- Validate WideDecimalBinaryExpr output type matches CheckOverflow data_type before bypassing the overflow check - Handle s_out > natural_scale (scale-up) in multiply path for consistency with add/subtract

martin-g · 2026-03-06T09:50:31Z

Close+reopen to re-trigger CI.

mbutrovich

LGTM, thanks @andygrove! I've got some ideas to optimize the fixed-point decimal stuff but that belongs as arrow-rs kernels. This optimization makes sense for Comet!

andygrove · 2026-03-06T21:36:04Z

Moving to draft until I can address #3619 (comment)

andygrove added 2 commits March 2, 2026 10:55

feat: add criterion benchmark for wide decimal binary expr

d7495bd

Add benchmark comparing old Cast->BinaryExpr->Cast chain vs fused WideDecimalBinaryExpr for Decimal128 add/sub/mul. Covers four cases: add with same scale, add with different scales, multiply, and subtract.

andygrove force-pushed the wide-decimal-binary-expr branch from cb52636 to d7495bd Compare March 2, 2026 17:55

andygrove force-pushed the wide-decimal-binary-expr branch from 5a21500 to 91092a6 Compare March 2, 2026 18:46

andygrove changed the title ~~feat: fused WideDecimalBinaryExpr for Decimal128 add/sub/mul~~ perf: Optimize some decimal expressions Mar 3, 2026

Merge branch 'main' into wide-decimal-binary-expr

486eda2

andygrove marked this pull request as ready for review March 3, 2026 15:52

parthchandra reviewed Mar 5, 2026

View reviewed changes

martin-g reviewed Mar 5, 2026

View reviewed changes

andygrove added 2 commits March 5, 2026 15:26

style: apply cargo fmt

1141b39

fix: add defensive checks for CheckOverflow bypass and multiply scale-up

02d6a3b

- Validate WideDecimalBinaryExpr output type matches CheckOverflow data_type before bypassing the overflow check - Handle s_out > natural_scale (scale-up) in multiply path for consistency with add/subtract

martin-g closed this Mar 6, 2026

martin-g reopened this Mar 6, 2026

martin-g approved these changes Mar 6, 2026

View reviewed changes

mbutrovich approved these changes Mar 6, 2026

View reviewed changes

andygrove marked this pull request as draft March 6, 2026 21:35

Conversation

andygrove commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

TPC-H q1

Criterion benchmark results (8192 element batches)

How it works

Uh oh!

andygrove commented Mar 2, 2026

Uh oh!

andygrove commented Mar 2, 2026

Uh oh!

andygrove commented Mar 2, 2026

Uh oh!

sqlbenchmark commented Mar 2, 2026

Uh oh!

coderfender commented Mar 4, 2026

Uh oh!

parthchandra Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

andygrove commented Mar 5, 2026

Uh oh!

martin-g commented Mar 6, 2026

Uh oh!

mbutrovich left a comment

Choose a reason for hiding this comment

Uh oh!

andygrove commented Mar 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

andygrove commented Mar 2, 2026 •

edited

Loading