Skip to content

refactor: replace remaining hand-rolled loops with Arrow kernels - #5367

Open
0lai0 wants to merge 1 commit into
apache:mainfrom
0lai0:refactor-5091-arrow-kernels
Open

refactor: replace remaining hand-rolled loops with Arrow kernels#5367
0lai0 wants to merge 1 commit into
apache:mainfrom
0lai0:refactor-5091-arrow-kernels

Conversation

@0lai0

@0lai0 0lai0 commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Part of #5091.

Rationale for this change

Each site had a hand-rolled per-row loop that an existing Arrow kernel already covers. Behaviour is unchanged and the kernel versions are faster on the shapes measured.

What changes are included in this PR?

File Before After
conversion_funcs/numeric.rs::spark_cast_decimal_to_boolean BooleanBuilder loop of value.is_zero() neq against a Scalar decimal zero at the source (precision, scale)
array_funcs/array_insert.rs::ArrayInsert::evaluate Two (0..num_rows).map(is_valid).collect() loops materialising BooleanArrays is_not_null(src) and and(evaluate_pos, is_not_null(pos))
math_funcs/pow.rs::spark_pow 4-way match on array/scalar shapes with per-row iterators apply from datafusion::physical_expr_common::datum, dispatching to binary and unary. Null-scalar short-circuit kept explicit because unary only preserves the input array's null buffer.

spark_cast_decimal_to_boolean also plumbs the source (precision, scale) through the zero scalar so Decimal128(38, 0) compares against a matching-scale zero.

Also included

Three criterion benches under native/spark-expr/benches/ wired in Cargo.toml.

Attempted and reverted

temporal.rs::days_to_dateDate32Type::to_naive_date_opt. date_trunc regressed +15–29% on three of four shapes because chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap() is not const, so the epoch is reconstructed per row. The old const i32 offset let LLVM fold it. Re-runnable once upstream makes the epoch const.

covariance.rs::{update_batch,retract_batch}and(is_not_null(a), is_not_null(b)) + filter. Would have matched CorrelationAccumulator, but filter allocates two new Float64Arrays per batch and regressed sparse-null shapes by +18–33%. The no-null path was −12%; not enough to justify a per-batch heuristic.

How are these changes tested?

Existing tests in each file pass unchanged. Three tests were added for the paths the refactors newly reach:

  • numeric.rs: test_spark_cast_decimal_to_boolean extended with Decimal128(38, 0), pinning the zero-scalar precision/scale wiring.
  • array_insert.rs: test_array_insert_evaluate_cross_null_patterns drives ArrayInsert::evaluate with four rows (src NULL, pos NULL, item NULL, all-non-null), pinning the Spark evaluation-order contract.
  • pow.rs: test_spark_pow_null_scalar covers the null-scalar short-circuit.

SQL-level coverage already exists in pow.sql, cast_decimal_to_primitive.sql, and array_insert*.sql.

Benchmarks

Baseline captured on main's versions of the three files (git stash push -- <files>, cargo bench --save-baseline main), refactors restored, cargo bench --baseline main re-run on the same machine. 8192 rows per shape. All p < 0.05.

spark_cast_decimal_to_boolean

shape before after change
no nulls 16.89 µs 1.95 µs −88.5%
sparse nulls 18.24 µs 1.91 µs −89.5%
dense nulls 18.39 µs 1.92 µs −89.6%

ArrayInsert::evaluate

shape before after change
no nulls 168.3 µs 125.5 µs −25.2%
sparse src nulls 281.5 µs 238.2 µs −14.6%
dense src nulls 216.7 µs 187.5 µs −13.3%
mixed src+pos nulls 293.8 µs 243.6 µs −16.9%

Gain comes from dropping the two BooleanArray::from(Vec<bool>) allocations per batch.

spark_pow

shape before after change
array/array no nulls 39.17 µs 27.62 µs −30.1%
array/array sparse nulls 44.56 µs 26.88 µs −38.7%
array/array dense nulls 32.91 µs 16.64 µs −50.0%
scalar/array no nulls 42.35 µs 34.21 µs −18.5%
scalar/array sparse nulls 43.66 µs 30.96 µs −29.4%
scalar/array dense nulls 29.58 µs 20.26 µs −31.7%
array/scalar no nulls 34.93 µs 25.35 µs −27.3%
array/scalar sparse nulls 39.04 µs 25.83 µs −34.4%
array/scalar dense nulls 28.38 µs 17.02 µs −40.4%
null scalar short-circuit 1.46 µs 1.39 µs −5.0%

Dense-null shapes gain the most because unary/binary skip null slots that the old iter().zip().map().collect() still visited via Option matching.

@andygrove andygrove left a comment

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.

Thanks for this. The three refactors all look equivalent to me, and the write-up is unusually thorough. The "Attempted and reverted" section with benchmark numbers is exactly what I want to see in a refactor PR like this.

I traced each change:

  • The zero scalar in spark_cast_decimal_to_boolean is built from decimal_array.precision()/scale(), so neq never hits a type mismatch, and comparing raw i128 against zero at identical scale is exactly !is_zero(). Slices and offsets are handled by the kernel. cast_decimal_to_primitive.sql already covers (10,2), (5,0), (15,5), (20,0), and (38,18), and the new unit test adds (38,0).
  • In array_insert.rs, src_value and pos_value are both into_array(batch.num_rows()), so the lengths line up. is_not_null returns a mask with no null buffer, same as BooleanArray::from(Vec<bool>) did, and and of two null-free masks stays null-free, so evaluate_selection behaves the same.
  • apply's scalar/scalar arm round-trips through ScalarValue::try_from_array, so spark_pow still returns a Scalar for two scalars. The explicit null-scalar short-circuit is genuinely needed, since binary would reject the length-1 against length-N pair.

I also checked the new array_insert test against Spark. ArrayInsert.eval in collectionOperations.scala evaluates first, then second only when first is non-null, then third, so the comment about the evaluation-order contract is accurate, and the expected [NULL, 6, 7] for the null-item row matches nullSafeEval.

Using Rust unit tests rather than SQL file tests is the right call here. pow(NULL, exp) gets rewritten to a null literal by Spark's NullPropagation because Pow is null-intolerant, so the null-scalar branch is not reachable from SQL at all.

A few things I would like addressed before merge.

One factual error in the PR description. Under the spark_pow benchmark table you write that dense-null shapes gain the most "because unary/binary skip null slots that the old iter().zip().map().collect() still visited via Option matching". Neither kernel skips null slots. unary applies the op to every value in the buffer and copies the null buffer through, and binary unions the two null buffers and then computes over all raw values. The win comes from dropping the per-element Option branch and the validity-building collect, not from doing less arithmetic. Could you reword that line? The description becomes the merge commit message, and it is the kind of thing the next person doing a kernel-dedup refactor will read as guidance. The doc comment on spark_pow_kernel itself is accurate, so it is only the description.

Scope of #5091. The title says "remaining", but make_decimal.rs (item 5 in the issue) is still a Decimal128Builder loop on main and is not touched here. That is fine given the PR says "Part of", just confirming #5091 stays open after this merges.

Could you also copy the "Attempted and reverted" section into a comment on #5091? The days_to_date and covariance findings with the benchmark numbers are the most valuable part of this work, and if they only live in a merged PR description nobody will find them before re-attempting the same change.

On the reverted covariance item. The reason covariance was listed in #5091 was the fragile dual-iterator null re-sync in update_batch, where two flatten() iterators get advanced conditionally on is_valid(i). That readability problem is independent of which kernel you reach for, and filter is not the only way out. A plain values1.iter().zip(values2.iter()) with a (Some(a), Some(b)) match drops the flatten() dance with no extra allocation, so it would not hit the regression you measured. Worth noting on #5091 alongside the filter result so the item does not get written off as not worth doing.

CI has not run on this yet, the check rollup is empty. I will get the workflows approved. In the meantime I ran the following locally on d6549f3 and all three are clean:

  • cargo test -p datafusion-comet-spark-expr --lib (629 passed, 0 failed)
  • cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings
  • cargo fmt --all -- --check

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants