Skip to content

fix: skip arithmetic and concat operands that Spark skips under ANSI mode - #7456

Open
tlfua wants to merge 4 commits into
apache:mainfrom
tlfua:tlfua/fix/7246-ansi-short-circuit
Open

tlfua wants to merge 4 commits into
apache:mainfrom
tlfua:tlfua/fix/7246-ansi-short-circuit

Conversation

@tlfua

@tlfua tlfua commented Nov 3, 2026 •

Copy link
Copy Markdown

Which issue does this PR close?

Part of #7246. It fixes the arithmetic and concat rows of that issue. OR, AND and IN were fixed by #7326, and the later sort keys of a top-K are left for a separate PR, so the issue stays open.

Rationale for this change

With ANSI mode on, native Comet evaluates an operand over the whole batch even on rows where Spark never evaluates it, so a query that Spark answers fails in Comet with CAST_INVALID_INPUT:

  • +, - and * evaluate the right operand only where the left one is not NULL (BinaryExpression.nullSafeCodeGen).
  • /, DIV and % evaluate the divisor first, and the dividend only where the divisor is not NULL (DivModLike.doGenCode, the same in Spark 3.4 and 4.1).
  • concat evaluates each argument only while no earlier one is NULL (Concat.doGenCode, Spark's default generated code).

For example SELECT n + CAST(txt AS INT) FROM t returns NULL in Spark wherever n is NULL, whatever txt holds, but fails in Comet when such a row has a malformed txt.

What changes are included in this PR?

The fix reuses NullShortCircuit, which already evaluates later function arguments only where no earlier one is NULL (#6716).

  • NullShortCircuit::wrap_in_order evaluates the arguments in a given order, so that the divisor of a division comes first. wrap keeps the order of the children.
  • MathExpr gains null_short_circuit. The serde sets it when the operand that Spark may skip can be skipped (canSkipEvaluation), which is the right operand of +, - and * and the left one of /, DIV and %. The planner then wraps the arithmetic, including the checked ANSI functions, decimal arithmetic and the Decimal256 remainder (create_modulo_expr gains a null_short_circuit flag), keeping CheckedBinaryExpr outermost so errors still carry the query context. skippable is a required argument of createMathExpression, so a new MathExpr serde cannot leave it out.
  • The rebalanced LEGACY integer add and multiply chains set the flag on each node from the operands of its right half, so each operand is evaluated on the same rows as in Spark's left-deep chain.
  • CometConcat wraps the native concat with withNullShortCircuit when it has more than one argument.
  • A masked argument that is deterministic may first be evaluated over the whole batch, selecting the rows only if that fails, as LazyLogical does since fix: evaluate the right operand of AND, OR and IN only for rows the left side has not decided #7326. Without speculation, a * (b - c) under ANSI, as in TPC-H Q1, was up to 2x slower on batches with NULLs, because selecting the rows costs more than the arithmetic. Whether to speculate depends on the cost of the argument:
    • A cheap argument is always speculated. is_cheap is a list of what has been measured: columns, literals, integer and floating point +, - and * (including the ANSI checked_add, checked_sub and checked_mul), and widening casts, over cheap children only. It can grow as more shapes are measured.
    • Any other argument, such as a cast that parses a string, decimal arithmetic or a string function, is speculated only when at least half the rows need it. Otherwise the rows are selected directly, as before this change.
  • Masked, shared by NullShortCircuit and LazyLogical, holds a masked argument: its projection, built lazily on the first evaluate so that planning nested wrappers stays linear, and whether speculation still holds. Both speculate a costly argument when at least half the rows need it; LazyLogical used to require more than half.
  • Cast::cannot_fail also accepts casts that cannot fail for any value in any eval mode: from an integer to a decimal with enough integral digits, and from an integer, floating point number, boolean or decimal to a string. Dates and timestamps are left out because Arrow's formatter cannot format a value outside chrono's range, which Spark can hold, and so are binary and nested types. Such arguments are not wrapped at all. They are added to cannot_fail only, not to is_infallible, which also decides which CASE WHEN branches are evaluated over the whole batch up front.
  • is_deterministic moves from private to pub(super), accepts CheckOverflow, and judges a nested NullShortCircuit by its inner expression alone, so decimal and nested operands speculate too, and the check stays linear in the size of the tree.
  • unwrap_checked in the planner also looks through NullShortCircuit, so the redundant CheckOverflow over a wide decimal add, subtract or multiply is still dropped.
  • The NullShortCircuit and withNullShortCircuit docs state that the function must return NULL wherever an earlier argument is NULL without checking the later ones, since with speculation it can see any value there.
  • create_binary_expr lost its last caller and is removed. Its tests call create_binary_expr_with_options.

An operand that Spark does evaluate still fails as before. A wrapped operand is NULL only on rows whose result is NULL anyway.

The same bug is on branch-1.2, which has NullShortCircuit but not #7326, which this builds on. The backport is tracked in #7542.

How are these changes tested?

  • spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql and spark/src/test/resources/sql-tests/expressions/string/concat_ansi_skipped_arguments.sql run every operator over a batch where the failing operand is malformed only on rows that Spark skips, over one where every row is skipped, and over mostly_null, where three rows in four are skipped, so a cheap argument (n * (big - 1)) is speculated, overflows on a skipped row and falls back to selecting rows, and a costly one (n + CAST(txt AS INT)) selects the rows directly. They cover tinyint, smallint, int, bigint, double, decimal and the Decimal256 paths, try_add and try_divide (including try_divide(CAST(txt AS INT), n) with a zero divisor, integer and decimal), and a zero divisor, where a NULL dividend gives NULL and a malformed one raises CAST_INVALID_INPUT before the zero check, as in Spark. expect_error_class(CAST_INVALID_INPUT) queries check that an operand Spark evaluates still fails, including the reverse order (CAST(txt AS INT) + n, n DIV CAST(txt AS INT)). Before the fix, the first query in each file fails with CAST_INVALID_INPUT.
  • Rust unit tests in null_short_circuit.rs, lazy_logical.rs and modulo_expr.rs: the divisor is evaluated first, it still fails where the dividend is NULL, order must be a permutation, the order survives with_new_children, speculative evaluation falls back to selected rows, a mostly-NULL batch selects the rows of a costly argument without speculating, a cheap one (checked_mul(a, checked_sub(b, c))) is speculated even with three rows in four skipped and still gives the right result when the skipped rows overflow, decimal and nested operands speculate, is_deterministic checks 64 nested wrappers in linear time, 64 nested wrappers that each read a different column are planned in linear time in both null_short_circuit.rs and lazy_logical.rs, and an argument needed by exactly half the rows is speculated. test_check_overflow_skipped_over_short_circuited_wide_decimal in planner.rs checks that the redundant CheckOverflow is still dropped over a wrapped wide decimal add.
  • spark/src/test/resources/sql-tests/expressions/math/arithmetic_legacy_skipped_operands.sql runs LEGACY integer + and * chains of 2 to 5 operands, including one that starts with a BIGINT cast, whose last operand element_at(arr, i) fails on the rows Spark skips, so Comet must return NULL there. It still raises INVALID_INDEX_OF_ZERO where Spark evaluates that operand, and where it comes first. concat_ansi_skipped_arguments.sql also has a query whose arguments cast an int, decimal, boolean and double to a string.
  • cast.rs: test_cannot_fail_casts_succeed_on_extreme_values_in_every_mode casts the extreme values of each type that cannot_fail accepts in LEGACY, TRY and ANSI mode, including decimals of scale 38 and -2, and test_casts_that_can_fail_are_not_claimed checks that one integral digit fewer, a negative scale, dates, timestamps, binary, Decimal256 and lists are not accepted.
  • native/spark-expr/benches/null_short_circuit.rs compares the eager form with the wrapped one at 0, 10, 50, 70 and 90% NULL, for a + CAST(s AS BIGINT), a * (b - c), the decimal p * CheckOverflow(1.00 - d) in ANSI and LEGACY mode, a WideDecimalBinaryExpr multiply, a multiply by a rescaling decimal cast, concat(q, upper(s)), and concat(q, CAST(b AS STRING)) over a long b, which is no longer wrapped. Anyone can rerun it with cd native && cargo bench -p datafusion-comet-spark-expr --bench null_short_circuit.

µs per 8192-row batch, best of two rounds on one GitHub runner (run, from the measurement-only branch https://github.com/tlfua/datafusion-comet/tree/bench/7246-r4, which runs the benchmark on this head and on the source of the previous head c818be3). "eager" is the unwrapped expression, as on main, which gives wrong errors.

Full table
shape NULLs eager previous head (c818be3) this head (b948e2a)
a + CAST(s AS BIGINT) 0% 113.3 114.8 113.3
10% 114.4 115.8 114.8
50% 110.4 95.1 95.0
70% 108.5 57.0 57.0
90% 107.2 23.1 22.9
a * (b - c) over longs 0% 17.4 17.9 17.6
10% 18.1 18.6 18.4
50% 14.7 15.3 15.2
70% 13.0 13.6 13.5
90% 11.5 12.1 12.0
decimal p * (1.00 - d), ANSI 0% 51.4 52.5 51.8
10% 59.0 60.6 59.5
50% 42.6 48.5 48.5
70% 33.8 30.8 30.5
90% 27.2 16.5 16.2
decimal p * (1.00 - d), LEGACY 0% 51.3 52.5 51.8
10% 59.0 60.6 59.5
50% 42.4 48.6 48.5
70% 33.7 30.8 30.5
90% 27.2 16.5 16.3
wide decimal p * (d * d) 0% 317.9 317.9 317.8
10% 293.2 292.6 293.5
50% 234.6 169.6 170.3
70% 204.1 101.1 101.7
90% 177.1 40.4 40.1
p * CAST(d AS DECIMAL(12,4)) 0% 110.9 111.6 111.6
10% 118.7 119.8 119.4
50% 102.6 78.2 78.3
70% 93.7 47.8 47.6
90% 87.2 22.2 21.9
concat(q, upper(s)) 0% 109.0 109.8 110.3
10% 131.7 135.2 135.7
50% 163.2 210.9 209.5
70% 133.3 154.1 155.3
90% 110.7 105.3 105.9
concat(q, CAST(b AS STRING)) 0% 230.1 225.4 230.7
10% 255.1 255.1 253.3
50% 280.4 246.3 280.9
70% 256.6 174.3 256.3
90% 231.4 111.7 230.9

What the numbers show:

  • Apart from concat(q, CAST(b AS STRING)), this head is within 2% of the previous head for every shape and NULL ratio.
  • concat(q, CAST(b AS STRING)) is no longer wrapped, so it now costs the same as eager at every NULL ratio. Wrapped, as on the previous head, it was faster at 50% NULL and above (1.14x at 50%, 1.47x at 70%, 2.07x at 90%), because concat then skips formatting the rows whose result is NULL anyway.
  • concat(q, upper(s)) is still up to 1.3x slower than eager at 50% NULL. Adding string functions to is_cheap is tracked in Speculate string functions such as upper in NullShortCircuit by adding them to is_cheap #7587.
  • An earlier run that forced always or never speculating (run, branch https://github.com/tlfua/datafusion-comet/tree/bench/7246) showed that always speculating a * (b - c) is 1.2x to 2.2x faster than selecting rows at 10 to 70% NULL, while for the other shapes always speculating is 1.2x to 5x slower than selecting rows at 70 to 90% NULL, which is why only cheap arguments are always speculated.

Commands run locally (macOS, default Spark profile), on b948e2a:

  • cargo test -p datafusion-comet-spark-expr: 1356 passed.
  • cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings: passed.
  • make format: no changes beyond this PR's files.
  • python3 dev/ci/check-suites.py and the other dev/ci checks: passed.
  • cargo bench -p datafusion-comet-spark-expr --bench null_short_circuit: run on a GitHub runner, numbers above.

I did not run the Maven suites locally this round. Fork CI (Comet CI (fork run), default Spark profile) passed on b948e2a, including the new SQL tests, the Spark 4.1 Comet suites and the TPC-H and TPC-DS result checks.

I did not run the Spark SQL suite or the other Spark profiles locally. This changes the serde and native expressions, so a maintainer may want to apply run-spark-4.1-tests, and run-benchmark-check for the benchmark changes.

…mode

Spark's generated code evaluates the second operand of + - * only where the
first is not NULL, the dividend of / div % only where the divisor is not NULL,
and each concat argument only where every earlier one is not NULL. Comet
evaluated every operand on every row, so an ANSI cast of a malformed string on
a row that Spark returns NULL for failed the query.

Wrap these expressions in NullShortCircuit, extended with an evaluation order
so the divisor is evaluated first. A deterministic masked argument is first
evaluated over the whole batch and only falls back to the selected rows if
that fails, which keeps the common case as fast as eager evaluation.

Part of apache#7246

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36
@github-actions github-actions Bot added bug Something isn't working area:expressions Expression evaluation labels Nov 3, 2026
@tlfua
tlfua marked this pull request as ready for review November 3, 2026 09:28
@tlfua

tlfua commented Nov 3, 2026

Copy link
Copy Markdown
Author

This changes the serde and native arithmetic and concat expressions, so could a maintainer please apply run-spark-4.1-tests to run the Spark SQL suite? The fork CI run passed: https://github.com/tlfua/datafusion-comet/actions/runs/41600696730

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for picking this up. Reusing NullShortCircuit with an evaluation order keeps the change small. The order matches Spark 3.4.3, 3.5.8, 4.0.1 and 4.1.1: BinaryArithmetic.doGenCode uses nullSafeCodeGen for every numeric type including decimals, DivModLike evaluates the divisor first in both eval and doGenCode, and Concat.doGenCode stops at the first NULL.

Comparisons have the same shape. EqualTo and the ordering comparisons use defineCodeGen, so Spark evaluates the right side only where the left is not NULL. I expect WHERE n = CAST(txt AS INT) under ANSI to fail in Comet the way n + CAST(txt AS INT) did, but I have not run it. Does that belong under #7246, or in its own issue?

Comment on lines +322 to +327
if nulls.is_some() && self.speculate.load(Ordering::Relaxed) {
if let Ok(value) = evaluate_projected(projection, projected, batch, None) {
return Ok(value);
}
self.speculate.store(false, Ordering::Relaxed);
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LazyLogical only speculates when more than half of the rows need the right operand (pending_rows * 2 > num_rows in lazy_logical.rs). Here any NULL is enough, so when nearly every row is skipped the masked argument runs over the whole batch. That also applies to the existing users from #6716, such as array_contains, slice, make_date and GetMapValue, which on main evaluate an argument like CAST(s AS INT) over only the rows that need it. Could this use the same rule, for example nulls.null_count() * 2 < batch.num_rows(), and could the benchmark add a point at 90% NULL? I have not measured it.

The struct doc also still says the function sees NULL for each skipped argument. With speculation it sees the computed value there instead, so the function must also not fail on a row where an earlier argument is NULL, whatever the later ones hold. SparkArraySlice and SparkMakeDate check for NULL first, so I do not see a current bug, but could the doc here and the withNullShortCircuit Scaladoc say so, so that a new kernel does not validate first?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Thanks, measured it before picking a cutoff. The benchmark now also runs at 70% and 90% NULL, has a decimal shape, and I ran it on a GitHub runner always speculating and never speculating (run, from the measurement-only branch https://github.com/tlfua/datafusion-comet/tree/bench/7246, which is not part of this PR). The table is in the PR description.

Where selecting the rows starts to win depends on how much the argument costs. For a + CAST(s AS BIGINT) it is already faster at 50% NULL (65 vs 77 µs). For a * (b - c) it only wins at 90% (7.9 vs 8.6 µs), and at 50% it is 1.8x slower (20.2 vs 11.2 µs). So null_count * 2 < num_rows would cost the cheap arithmetic 1.8x at exactly half NULL.

I went with null_count * 2 <= num_rows, the same split as LazyLogical. For the three shapes it is at most 1.2x off the better choice up to half NULL and 1.4x above it, and at 90% NULL it is 4.5x faster than always speculating for the cast. A batch above the cutoff selects the rows without turning speculation off, so later batches still speculate. That also keeps the existing NullShortCircuit users from #6716 on the old path for mostly-NULL batches.

The doc now says the function must return NULL wherever an earlier argument is NULL without checking the later ones, both on NullShortCircuit and on withNullShortCircuit.

Comment on lines +1541 to +1551
let expr = if options.null_short_circuit {
// Spark evaluates the divisor of a division first
let order: &[usize] = if op == DataFusionOperator::Divide {
&[1, 0]
} else {
&[0, 1]
};
NullShortCircuit::wrap_in_order(expr, order, &input_schema)?
} else {
expr
};

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I think two things change for decimal operands once this wrapper is in the tree. I have not run it, this is from reading the code.

  1. DecimalPrecision.promote puts every decimal +, -, *, / and % under CheckOverflow. The CheckOverflow arm skips the check for a WideDecimalBinaryExpr (perf: Optimize some decimal expressions #3619), but unwrap_checked only looks through CheckedBinaryExpr. For CheckedBinaryExpr(NullShortCircuit(WideDecimalBinaryExpr)), is_wide_decimal is now false and the extra CheckOverflow pass comes back. CAST(n AS DECIMAL(38, 0)) + CAST(txt AS DECIMAL(38, 0)) in the new fixture takes that path.
  2. A decimal operand reaches NullShortCircuit as a CheckOverflow, which is neither in cannot_fail nor in is_deterministic. So price * (1 - discount) is wrapped even in LEGACY mode, where Spark returns NULL on overflow rather than failing, and speculation never applies to it. A batch with NULL prices then takes the selection path that the description measured at up to 2x for BIGINT.

Could unwrap_checked look through NullShortCircuit, and is_deterministic accept CheckOverflow? A decimal shape in the new benchmark would show whether that is enough.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

You're right on both, thanks for reading it that closely.

  1. unwrap_checked now also looks through NullShortCircuit, so the CheckOverflow over a wide decimal add, subtract or multiply is dropped again. test_check_overflow_skipped_over_short_circuited_wide_decimal in planner.rs plans CheckOverflow(a + CAST(s AS DECIMAL(38, 0))) under ANSI and checks that the result is the wrapped WideDecimalBinaryExpr with no CheckOverflow, and that it returns NULL for the row where a is NULL and s is malformed.
  2. is_deterministic now accepts CheckOverflow, and a NullShortCircuit whose inner expression is deterministic, so decimal and nested operands speculate. decimal_and_nested_short_circuited_arguments_are_speculated covers both, and the benchmark has a decimal shape now.

I did not add decimal arithmetic or CheckOverflow to cannot_fail, because I can't show that the decimal kernels never return an error. So a LEGACY decimal operand is still wrapped, but with speculation it costs about the same as unwrapped (decimal rows in the table).

statement
INSERT INTO filled SELECT /*+ REPARTITION(1) */ * FROM VALUES (1, 1, 'bad'), (2, 2, '2') AS t(id, n, txt)

-- Add, subtract and multiply evaluate the right operand only where the left one is not NULL

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Could this file also cover the other paths the planner now wraps?

  • TINYINT and SMALLINT, which use the checked_* kernels on Int8 and Int16
  • DOUBLE for +, - and *, which stays a plain BinaryExpr
  • try_add and try_divide, which Spark plans as EvalMode.TRY arithmetic
  • a zero divisor under ANSI. Spark returns NULL for a NULL dividend over 0, and for CAST(txt AS INT) DIV n with txt = 'bad' and n = 0 it raises CAST_INVALID_INPUT rather than DIVIDE_BY_ZERO, because DivModLike evaluates the dividend before its zero check.

null_short_circuit_skips_the_dividend_where_the_divisor_is_null in modulo_expr.rs repeats the DECIMAL(38, 10) % DECIMAL(38, 0) query in this file, so it could go.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Added all of them:

  • TINYINT + and %, SMALLINT * and DIV
  • DOUBLE +, - and *
  • try_add and try_divide over skipped rows, and try_add with a malformed operand that Spark does evaluate, which still raises CAST_INVALID_INPUT
  • a zero divisor: DIV, % and / return NULL for a NULL dividend over 0, and CAST(txt AS INT) DIV n with txt = 'bad', n = 0 raises CAST_INVALID_INPUT

Comet matches Spark on all of them. I also removed the duplicate test in modulo_expr.rs.


/// Like [`create_modulo_expr`], but with `null_short_circuit` evaluates the dividend only for the
/// rows where the divisor, which Spark evaluates first, is not NULL. See [`NullShortCircuit`].
pub fn create_modulo_expr_with_null_short_circuit(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

After this change RemainderBuilder calls create_modulo_expr_with_null_short_circuit, so create_modulo_expr has no caller left outside its tests. The PR removes create_binary_expr for the same reason. Could create_modulo_expr take the null_short_circuit flag itself, so the crate keeps one public function instead of two?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Done. create_modulo_expr now takes null_short_circuit: bool, create_modulo_expr_with_null_short_circuit is gone, and the tests pass false.

- Speculate a masked argument only when at least half the rows need it,
  chosen from benchmarks at 0-90% NULL across three operand shapes.
- Look through NullShortCircuit in unwrap_checked so the redundant
  CheckOverflow over wide decimal arithmetic is still dropped, and treat
  CheckOverflow and nested NullShortCircuit as deterministic so decimal
  and nested operands speculate.
- Merge create_modulo_expr_with_null_short_circuit into create_modulo_expr.
- Document that a wrapped function must return NULL where an earlier
  argument is NULL without checking the later ones.
- Cover tinyint, smallint, double, try_add, try_divide and a zero divisor
  in the SQL tests, and drop the duplicate modulo unit test.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36
@tlfua

tlfua commented Nov 3, 2026

Copy link
Copy Markdown
Author

Thanks for the review. I've pushed dc1559c for the comments above. The PR description has the updated benchmark table and test list.

On comparisons, you're right. SELECT id, n < CAST(txt AS INT) FROM cmp fails in Comet under ANSI where n is NULL and txt is malformed, while Spark returns NULL. WHERE n = CAST(txt AS INT) happened not to fail, since the inferred isnotnull(n) and the lazy AND skip the row first. I'd keep it out of this PR since it is already large, and opened #7468 for it.

@tlfua

tlfua commented Nov 3, 2026

Copy link
Copy Markdown
Author

Fork CI is green on dc1559c (run). Since this round also changes native/spark-expr/benches/null_short_circuit.rs, could a maintainer please apply run-benchmark-check along with run-spark-4.1-tests? run-all-spark-profiles would also help, since the new SQL cases (try_divide, the zero divisor errors) run on every Spark version. Thanks!

@coderfender coderfender left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for the careful work here, and especially for the benchmark table. I went through Spark's BinaryExpression.nullSafeCodeGen, DivModLike and Concat.doGenCode in 3.4, 3.5, 4.0 and 4.1, and the order you evaluate operands in matches. Operands Spark does evaluate still fail, in both directions. I also like that the zero divisor case outside ANSI falls out of the existing nullIfWhenPrimitive and null_if_zero_primitive rewrites.

A few things I'd like to settle before merge.

  1. native/spark-expr/src/conditional_funcs/null_short_circuit.rs:333. Compared with the release, which evaluates eagerly, your table shows a * (b - c) under ANSI going from 9.5 to 14.0 µs at 70% NULL, and it stays slower until about 85%. That shape is common under ANSI, and eager evaluation there only goes wrong if the subtraction actually fails on a skipped row. Could the select-rows path be kept for arguments where it pays off, such as ones containing a parsing Cast or a string function? Arguments built only from columns, literals and checked arithmetic would then always speculate first. If you'd rather keep a single threshold, could you say in the PR why the slowdown in that band is acceptable?

  2. native/spark-expr/benches/null_short_circuit.rs:162. The bench already produces nulls_0, but the table leaves it out. Could you add it? No NULLs is the common case. TPC-H Q1's l_extendedprice * (1 - l_discount) now goes through the wrapper even in LEGACY mode, because its right side is a CheckOverflow and cannot_fail doesn't accept it. Could you also add a concat shape such as concat(s, upper(t))? cannot_fail rejects every ScalarFunctionExpr, so after this PR any concat whose later argument is a string function is wrapped in every eval mode.

  3. spark/src/test/resources/sql-tests/expressions/math/arithmetic_ansi_skipped_operands.sql:170. Outside a failing eval mode, DivModLike also skips the dividend when the divisor is zero. Could you add SELECT id, try_divide(CAST(txt AS INT), n) FROM zero_bad, and a decimal variant, to pin that? Spark returns NULL there and the release raises CAST_INVALID_INPUT. A table where more than half the rows are NULL would also run the path that selects rows without speculating end to end. Right now skipped is exactly half NULL, so it always speculates first.

  4. Backport. The description says branch-1.2 has the same bug. Could we either backport a version without speculation (1.2 has NullShortCircuit but not #7326), or file a tracking issue for it as backporting.md asks for a correctness fix?

This changes the serde, the planner and native expressions, so I'll make sure the Spark SQL and benchmark checks run before it goes in the queue.

// Its children are those of the function it wraps
|| expr
.downcast_ref::<NullShortCircuit>()
.is_some_and(|n| is_deterministic(n.inner()));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Nested ANSI arithmetic doubles determinism work at each wrapper: n.inner() and the final all visit the same children. My 18-level helper probe counted 1,310,716 visits. Return the inner result directly.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Good catch, thanks. is_deterministic now handles a NullShortCircuit first and returns is_deterministic(inner), without visiting the children again, so the check is linear in the size of the tree. deeply_nested_short_circuits_are_checked_in_linear_time in lazy_logical.rs nests 64 wrappers, which would take 2^64 steps with the old code. Fixed in 9b6c7be.

Address round-3 review of apache#7456:

- Speculate a masked argument over the whole batch regardless of the
  NULL fraction when it is cheap (columns, literals, integer and float
  +, -, *, ANSI checked_add/sub/mul, widening casts), so `a * (b - c)`
  is no slower than before the fix on mostly NULL batches. Costlier
  arguments keep the half threshold.
- is_deterministic no longer visits the children of a NullShortCircuit
  twice, which was exponential in the nesting depth.
- Benchmark LEGACY TPC-H Q1, wide decimal, rescaling decimal cast and
  concat(q, upper(s)) shapes.
- SQL tests for try_divide with a zero divisor and for mostly NULL
  batches with cheap and costly operands.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36
tlfua pushed a commit to tlfua/datafusion-comet that referenced this pull request Nov 6, 2026
Address round-3 review of apache#7456:

- Speculate a masked argument over the whole batch regardless of the
  NULL fraction when it is cheap (columns, literals, integer and float
  +, -, *, ANSI checked_add/sub/mul, widening casts), so `a * (b - c)`
  is no slower than before the fix on mostly NULL batches. Costlier
  arguments keep the half threshold.
- is_deterministic no longer visits the children of a NullShortCircuit
  twice, which was exponential in the nesting depth.
- Benchmark LEGACY TPC-H Q1, wide decimal, rescaling decimal cast and
  concat(q, upper(s)) shapes.
- SQL tests for try_divide with a zero divisor and for mostly NULL
  batches with cheap and costly operands.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36
@tlfua

tlfua commented Nov 6, 2026

Copy link
Copy Markdown
Author

Thanks @coderfender, all four are addressed in 9b6c7be. The full table is in the updated description (bench run).

  1. Done as you suggested. is_cheap in null_short_circuit.rs lists columns, literals, integer and floating point +, - and * (including checked_add, checked_sub and checked_mul) and widening casts, over cheap children only. Such an argument is always speculated first. Anything else, such as a parsing Cast, decimal arithmetic or a string function, keeps the half-NULL threshold. a * (b - c) under ANSI is now within 5% of eager at every NULL ratio (11.0 vs 10.6 µs at 70% NULL, where it was 14.0). I also benchmarked the borderline candidates. Always speculating a WideDecimalBinaryExpr multiply or a rescaling decimal cast is 2.1x slower than selecting rows at 70% NULL, and the decimal p * CheckOverflow(1.00 - d) 1.2x to 1.9x at 70 to 90%, so they stay off the list. concat(q, upper(s)) is the one shape where this falls short: with upper(s) kept off the list, as you suggested for string functions, it is 1.4x slower than eager at 50% NULL and 1.2x at 70%. Always speculating it would bring that to 1.09x and 1.04x, and cost 5% at 90%. So infallible string functions look worth adding to the list, but each needs measuring. I'd prefer a follow-up, but can add upper here if you'd rather.
  2. The table now has nulls_0, and four new shapes: LEGACY p * CheckOverflow(1.00 - d) (your Q1 case), concat(q, upper(s)), a wide decimal multiply and a rescaling decimal cast. With no NULLs, the wrapper is within 3% of eager for every shape (a * (b - c) is 14.8 vs 14.4 µs, LEGACY Q1 42.1 vs 41.8).
  3. Added try_divide(CAST(txt AS INT), n) FROM zero_bad and a decimal variant, both NULL as in Spark. The new mostly_null table has n NULL on three rows in four. n * (big - 1) is cheap, so it is speculated, overflows on a skipped row and falls back to selecting rows. n + CAST(txt AS INT), CAST(txt AS INT) DIV n and the two concat queries select the rows without speculating.
  4. Filed Backport ANSI skipped-operand fix for arithmetic and concat to branch-1.2 #7542 to track the branch-1.2 backport.

@tlfua
tlfua force-pushed the tlfua/fix/7246-ansi-short-circuit branch from 9b6c7be to c818be3 Compare November 7, 2026 09:14

@rich7420 rich7420 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

@tlfua thanks for the patch!

} else {
&[0, 1]
};
NullShortCircuit::wrap_in_order(expr, order, &input_schema)?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

a DIV b over ints, and d + a with an int a, get wrapped even in LEGACY because Cast::is_infallible excludes exact casts such as CAST(a AS DECIMAL(19, 0)). Could integer-to-decimal casts with enough digits count as infallible?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Done in b948e2a. Cast::cannot_fail now accepts a cast from TINYINT, SMALLINT, INT or BIGINT to a DECIMAL(p, s) with s >= 0 and at least 3, 5, 10 or 19 integral digits (p - s), which holds every value of the integer. So the CAST(a AS DECIMAL(19, 0)) operands of a DIV b and of d + a are no longer wrapped. test_cannot_fail_casts_succeed_on_extreme_values_in_every_mode in cast.rs casts the minimum and maximum of each type in every eval mode, and test_casts_that_can_fail_are_not_claimed checks that one digit fewer is not accepted.

I put it in cannot_fail rather than is_infallible, because is_infallible also decides which CASE WHEN branches are evaluated over the whole batch up front, and these casts cost more per row than the ones it holds.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Cool

expr.evalMode,
(builder, mathExpr) => builder.setAdd(mathExpr))
(builder, mathExpr) => builder.setAdd(mathExpr),
skippable = Some(expr.right))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Every LEGACY integer +/*, even with two operands, takes the balanced branch without the flag, so n + element_at(a, 0) fails where n is NULL. Could each balanced node set it from its right slice, matching Spark's left-deep chain?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Good catch, thanks. Each balanced node now sets the flag from its right slice: createBalancedMathExpr keeps each operand next to its proto and calls canSkipEvaluation on the operands of the right half. A node skips its right half where its left half is NULL, which is where some earlier operand is NULL, so every operand is evaluated on the same rows as in Spark's left-deep chain. The new math/arithmetic_legacy_skipped_operands.sql runs chains of 2 to 5 operands under LEGACY, including one that starts with a BIGINT cast, whose last operand element_at(arr, i) fails on the rows Spark skips, and checks that it still raises INVALID_INDEX_OF_ZERO where Spark evaluates it, and where it comes first.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for the update

@coderfender coderfender left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for working through all of this so quickly. I checked the new head against my four points. is_cheap is what I had in mind, and I like that it only changes when speculation is tried, so a failure on a skipped row still falls back to selecting rows. The numbers in the description match the table, and the nulls_0 column shows no real overhead. The zero_bad and mostly_null queries pin what I asked for, and #7542 tracks the backport. Your reply cites 9b6c7be, whose tree is identical to c818be3, so the fork run covers this head.

There are two things I'd like to settle before merge, plus two simplifications.

  1. native/spark-expr/src/conditional_funcs/null_short_circuit.rs:336. rich7420's fix makes is_deterministic linear, but I think the same doubling is still in project. Arg::try_new projects its masked argument right away, and project rewrites column indices with transform_down, which calls with_new_children on every node above a changed column. For a nested NullShortCircuit that is Self::try_new again, which projects its own argument again. So projecting one level rebuilds the level below it twice, once on the way down and once inside the rebuilt wrapper, which comes to 2^k - 1 calls to project for k levels. It needs a projection that isn't the identity at each level, which is the normal case for c0 * (c1 * (c2 * ... (c19 - CAST(s AS INT)))) under ANSI, and nested decimal arithmetic hits it in any mode since its CheckOverflow operand is masked. I traced this through the code rather than running it, so could you check? deeply_nested_short_circuits_are_checked_in_linear_time wouldn't see it, because with only v and s in the schema every projection is the identity. A version that reads a different column at each level, with the outer ones first, should show it. Building the projection lazily on the first evaluate, for example in a OnceLock, looks like it would fix it and still give each argument one copy of its state across batches.

  2. spark/src/main/scala/org/apache/comet/serde/strings.scala:400. Thanks for measuring the concat shape. It is still 1.4x slower than the release at 50% NULL and 1.2x at 70%, and since upper can't fail, that cost buys nothing in any eval mode. I think it reaches further than upper. Cast::cannot_fail doesn't accept a cast to a string, so concat(name, CAST(id AS STRING)) is wrapped too, and that cast isn't on the is_cheap list either, so at half NULL or more it selects string rows. The new concat fixture uses that shape in every query. Could Cast::cannot_fail accept the casts to STRING that can't fail, such as from integers, booleans and decimals, after checking each against the native cast, so those calls aren't wrapped at all? A concat(q, CAST(a AS STRING)) shape in the bench would show the difference. For the string functions, I'd take you up on adding the ones you measured to is_cheap here, or please file the follow-up as an issue so it doesn't get lost.

  3. null_short_circuit.rs:353. Arg::evaluate and LazyLogical::evaluate now run the same steps. They try the whole batch while speculate holds, drop the error, turn speculate off and then evaluate the selected rows. Each keeps its own projection, projected and speculate with a long doc comment, and the two already differ at exactly half (> in LazyLogical, <= here). Could they share one small type next to project, something like Masked { projection, projected, speculate } with fn evaluate(&self, batch, rows: Option<BooleanArray>, try_whole_batch: bool)? Each caller would keep only its threshold, LazyLogical would drop two fields, and the lazy projection from point 1 would only need to go in one place.

  4. spark/src/main/scala/org/apache/comet/serde/arithmetic.scala:50. Every caller of createMathExpression passes skippable, so could it be a required Expression rather than Option[Expression] = None? Then a new MathExpr serde that forgets it won't compile, instead of quietly evaluating both operands over the whole batch, which is the bug this PR fixes.

@coderfender coderfender added run-all-spark-profiles Run the Comet test suites against every Spark profile on this pull request, ahead of the merge queue run-spark-4.1-tests Run the Spark 4.1 SQL tests on this pull request instead of waiting for the merge queue run-benchmark-check Run the benchmark compile and lint check on this pull request instead of waiting for the merge queue labels Nov 9, 2026
@coderfender

Copy link
Copy Markdown
Contributor

@tlfua I added run-all-spark-profiles, run-spark-4.1-tests and run-benchmark-check, since this changes the serde, the planner and native expressions and adds a bench.

- Share masked evaluation between NullShortCircuit and LazyLogical in a
  `Masked` type that projects its expression lazily on first evaluation,
  so nested short-circuits are planned in linear time. Speculate when at
  least half the rows need the expression in both.
- Accept casts to string from integers, floating point numbers, booleans
  and Decimal128, and integer to Decimal128 with enough integral digits,
  in `Cast::cannot_fail`, so they are no longer masked.
- Set the null short-circuit flag on each node of a rebalanced LEGACY
  integer add or multiply chain from its right slice.
- Make `skippable` a required argument of `createMathExpression`.
- Add a `concat_cast` benchmark shape and SQL tests.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 7fbb6474-7dd6-4c0e-bb27-bd78a480ca36
@tlfua

tlfua commented Nov 10, 2026

Copy link
Copy Markdown
Author

@coderfender thanks for checking it so carefully. b948e2a addresses all four points:

  1. You're right, project was doubling. Masked now builds its projection lazily, in a OnceLock, on the first evaluate, so building a nested wrapper no longer projects the level below it. deeply_nested_short_circuits_reading_other_columns_are_planned_in_linear_time in null_short_circuit.rs and deeply_nested_lazy_operands_reading_other_columns_are_planned_in_linear_time in lazy_logical.rs nest 64 levels that each read a different column, outer ones first, which would take 2^64 calls with the old code.
  2. Cast::cannot_fail now accepts casts to STRING from integers, floating point numbers, booleans and decimals, after checking each against the native cast: test_cannot_fail_casts_succeed_on_extreme_values_in_every_mode casts the extreme values of each type in every eval mode. So concat(name, CAST(id AS STRING)) is no longer wrapped. Dates and timestamps are left out on purpose: Arrow's formatter, which those casts use, can't format a value outside chrono's range, which Spark can hold. Binary and nested types are left out too. Other casts that can't fail, such as widening a decimal or casting a boolean to or from a number, can be added later by the same rule. The benchmark has a concat(q, CAST(b AS STRING)) shape now (run, this head against the previous one on one runner). Unwrapped, it now costs the same as eager at every NULL ratio. One thing I didn't expect: wrapped, as on the previous head, it was faster at 50% NULL and above (1.14x at 50%, 1.47x at 70%, 2.07x at 90%: 111.7 vs 230.9 µs), because concat then skips formatting the rows whose result is NULL anyway. So not wrapping it gives up that speedup on mostly-NULL batches. I kept it unwrapped as you suggested, since it matches main and there is no error to guard against, but I'm happy to take the cast to string back out of cannot_fail if you'd rather keep the speedup. For upper and the other string functions I filed Speculate string functions such as upper in NullShortCircuit by adding them to is_cheap #7587.
  3. NullShortCircuit's argument and LazyLogical's right operand now share one Masked type with the projection, the lazy projected expression and speculate. Each caller only passes whether the operand is cheap, and the threshold is the same in both: speculate when at least half the rows need it. That changes LazyLogical at exactly half, from selecting the rows to speculating, which half_the_rows_needing_an_argument_is_enough_to_speculate covers.
  4. skippable is now a required Expression.

Fork CI passed on b948e2a: https://github.com/tlfua/datafusion-comet/actions/runs/42754860727. The PR description has the updated table and test list.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:expressions Expression evaluation bug Something isn't working run-all-spark-profiles Run the Comet test suites against every Spark profile on this pull request, ahead of the merge queue run-benchmark-check Run the benchmark compile and lint check on this pull request instead of waiting for the merge queue run-spark-4.1-tests Run the Spark 4.1 SQL tests on this pull request instead of waiting for the merge queue

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants