refactor: replace remaining hand-rolled loops with Arrow kernels - #5367
refactor: replace remaining hand-rolled loops with Arrow kernels#53670lai0 wants to merge 1 commit into
Conversation
andygrove
left a comment
There was a problem hiding this comment.
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_booleanis built fromdecimal_array.precision()/scale(), soneqnever 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.sqlalready 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_valueandpos_valueare bothinto_array(batch.num_rows()), so the lengths line up.is_not_nullreturns a mask with no null buffer, same asBooleanArray::from(Vec<bool>)did, andandof two null-free masks stays null-free, soevaluate_selectionbehaves the same. apply's scalar/scalar arm round-trips throughScalarValue::try_from_array, sospark_powstill returns aScalarfor two scalars. The explicit null-scalar short-circuit is genuinely needed, sincebinarywould 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 warningscargo fmt --all -- --check
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?
conversion_funcs/numeric.rs::spark_cast_decimal_to_booleanBooleanBuilderloop ofvalue.is_zero()neqagainst aScalardecimal zero at the source(precision, scale)array_funcs/array_insert.rs::ArrayInsert::evaluate(0..num_rows).map(is_valid).collect()loops materialisingBooleanArraysis_not_null(src)andand(evaluate_pos, is_not_null(pos))math_funcs/pow.rs::spark_powmatchon array/scalar shapes with per-row iteratorsapplyfromdatafusion::physical_expr_common::datum, dispatching tobinaryandunary. Null-scalar short-circuit kept explicit becauseunaryonly preserves the input array's null buffer.spark_cast_decimal_to_booleanalso plumbs the source(precision, scale)through the zero scalar soDecimal128(38, 0)compares against a matching-scale zero.Also included
Three criterion benches under
native/spark-expr/benches/wired inCargo.toml.Attempted and reverted
temporal.rs::days_to_date→Date32Type::to_naive_date_opt.date_truncregressed +15–29% on three of four shapes becausechrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap()is notconst, so the epoch is reconstructed per row. The oldconst i32offset let LLVM fold it. Re-runnable once upstream makes the epochconst.covariance.rs::{update_batch,retract_batch}→and(is_not_null(a), is_not_null(b))+filter. Would have matchedCorrelationAccumulator, butfilterallocates two newFloat64Arrays 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_booleanextended withDecimal128(38, 0), pinning the zero-scalar precision/scale wiring.array_insert.rs:test_array_insert_evaluate_cross_null_patternsdrivesArrayInsert::evaluatewith four rows (src NULL, pos NULL, item NULL, all-non-null), pinning the Spark evaluation-order contract.pow.rs:test_spark_pow_null_scalarcovers the null-scalar short-circuit.SQL-level coverage already exists in
pow.sql,cast_decimal_to_primitive.sql, andarray_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 mainre-run on the same machine. 8192 rows per shape. Allp < 0.05.spark_cast_decimal_to_booleanArrayInsert::evaluateGain comes from dropping the two
BooleanArray::from(Vec<bool>)allocations per batch.spark_powDense-null shapes gain the most because
unary/binaryskip null slots that the olditer().zip().map().collect()still visited viaOptionmatching.