refactor: replace remaining hand-rolled loops with Arrow kernels

[0lai0]

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_date → Date32Type::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 both kernels operate over the raw value buffer and propagate validity in bulk, avoiding the per-element Option branch and the validity-building collect from the old iter().zip().map().collect() path.

[andygrove]

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

[0lai0]

PR description updated

dense-null shapes gain the most because both kernels operate over the raw value buffer and propagate validity in bulk, avoiding the per-element Option branch and the validity-building collect from the old iter().zip().map().collect() path.

[sunchao]

Thanks for the focused refactor and benchmark coverage. The array_insert and pow changes look sound. I found one decimal-cast compatibility regression below. On d6549f3, all 629 expression-crate unit tests, formatting, and 5,575 additional differential cases passed.

[sunchao]

[P2] Preserve all-null DECIMAL(0,0) casts

Could we preserve the all-null case before constructing this scalar? Spark accepts an RDD-backed nullable DecimalType(0,0) column, and SELECT CAST(d AS BOOLEAN) returns null. With spark.comet.sparkToColumnar.enabled=true and spark.comet.sparkToColumnar.supportedOperatorList=RDDScan, Comet's row-to-Arrow and FFI paths can carry that schema for null rows.

I reproduced the Spark behavior on 3.5.2 and the native cast separately on this head. Both an all-null Decimal128(0,0) array and ScalarValue::Decimal128(None, 0, 0) now fail with precision cannot be 0, has to be between [1, 38]. The previous loop returned null without validating the precision. The input-column cast survives Spark optimization, so this is not limited to a folded null literal.

An all-null fast path before with_precision_and_scale, with a regression test, would preserve the previous behavior. This does not assume non-null precision-zero values are supported.

[0lai0]

Thanks for review, I'll add an all-null fast path in spark_cast_decimal_to_boolean that returns BooleanArray::new_null(len) before constructing the zero scalar, plus a regression test on an empty Decimal128(0, 0) array . Thanks for the repro details.

[sunchao]

[P2] Preserve non-null precision-zero results from JVM codegen

Thanks for adding the all-null fast path. There is one non-null route I missed in the earlier report: a Java UDF returning java.math.BigInteger.ZERO with declared DecimalType(0, 0). Spark's converter stores that zero as a compact decimal, and Comet's default JVM-UDF codegen writes it through DecimalVector.setSafe(index, long). That overload does not validate precision, so this array can reach the native cast.

For a batch containing zero and null, the base implementation returns [false, null], but this head's zero-scalar construction throws precision cannot be 0, has to be between [1, 38] in Legacy, ANSI, and TRY modes. The all-null guard does not cover the valid zero slot.

I verified the Spark 4.1.3 converter and Arrow Java 18.3 writer, then imported the same array layout through Arrow 58.4 FFI and compared the exact base and head cast implementations. This is component-level reproduction, not an end-to-end SQL run. BigInteger.ZERO matters here: returning BigDecimal.ZERO does not take the same Spark conversion path.

Could we preserve the raw-value cast path for precision-zero arrays and add a JVM-UDF regression? A useful case is CAST(zero_decimal(id) AS BOOLEAN), where a Java UDF1[Long, BigInteger], explicitly declared as DecimalType(0, 0), returns BigInteger.ZERO for one input row and null for another.

[0lai0]

Thanks @sunchao for follow-up, I added a regression test in CometExpressionSuite.scala with checkSparkAnswerAndOperator so a native-plan fallback fails loudly.
Covered by:

• all-null fast path (skips the zero-scalar construction Arrow rejects for precision == 0)
• precision-zero fast path (v != 0 on the raw i128 payload, so an out-of-contract non-zero valid slot still round-trips)

[0lai0]

The CI test didn’t pass because the network connection is unavailable. Please help rerun it, thanks

[Incubating] Problems report is available at: file:///__w/datafusion-comet/datafusion-comet/apache-iceberg/build/reports/problems/problems-report.html
         project :iceberg-spark:iceberg-spark-4.1_2.13
      > Could not resolve org.scala-lang:scala-compiler:2.13.17.
         > Could not get resource 'https://repo.maven.apache.org/maven2/org/scala-lang/scala-compiler/2.13.17/scala-compiler-2.13.17.pom'.
            > Could not GET 'https://repo.maven.apache.org/maven2/org/scala-lang/scala-compiler/2.13.17/scala-compiler-2.13.17.pom'.
               > Got socket exception during request. It might be caused by SSL misconfiguration
                  > Network is unreachable

[sunchao]

[P2] Keep null-skipping for dense nullable inputs

Could we retain a null-aware path for dense inputs and benchmark a nullable child expression? For pow(a + 2.5D, b), with nullable double a and finite fractional exponents in b, Arrow addition preserves the null bits but changes their underlying values from zero to 2.5. These unary/binary calls then compute powers for null slots that the previous iterator skipped.

I compared the exact base/head pow.rs files in a standalone optimized harness with Arrow 58.4 and DataFusion 54.1, using Comet's release optimization settings (thin LTO and one codegen unit). On arm64 macOS, with 8,192-row batches, medians from nine alternating rounds after warm-up showed the following, including the addition in the timed pipeline:

Null fraction	Base	This head	Slowdown
About 90%	26.0 us	47.9 us	1.84x
About 99%	23.8 us	47.0 us	1.97x

Separate repeats and an independent rerun reproduced the regression, while the no-null pipeline improved. Nullness and valid result bits matched. These are local kernel/pipeline measurements, not full Spark-query or production-workload measurements.

The current benchmarks construct null slots directly from None, leaving zero payloads, and stop at 50% nulls. That misses this composed-expression case. Could we cover it and preserve null-skipping where evaluating those masked values is more expensive than the iterator overhead?

[0lai0]

Thanks for reviews !!
Fixed via an adaptive null-aware dispatcher that only pays the null-skip overhead when it wins. Crossover is null_count > 3 * len / 4 (75%).

Before vs. after @ 8192 rows (composed nulls, real payload in null slots, like pow(a + 2.5D, b)):

Shape	90% nulls	99% nulls
Regression reported before	47.9µs (1.84x slow)	47.0µs (1.97x slow)
After null-aware dispatch, array/array	5.56µs	1.83µs
After, pipeline pow(a+2.5D, 3) incl. add	6.35µs	3.37µs

Threshold behaviour is visible in the sweep: array/array composed 70% = 36.9µs (raw-buffer kernel), 80% = 9.1µs (null-skip kicks in). No-null / sparse-null shapes are unchanged (27.4µs / 26.7µs), so the null-oblivious fast path is preserved.

[andygrove]

Note on this review: this was generated by an LLM (Claude Code) at my request while I worked through a review backlog. I have not verified the individual findings myself. Please treat everything below as suggestions to evaluate rather than as authoritative review feedback, and push back on anything that is wrong or already handled.

This is a model PR. Per-shape benchmark tables with a stated baseline method, a section documenting the two refactors that were tried and reverted with the reason and the numbers, and a test added for each newly reached path. The date_trunc finding about NaiveDate::from_ymd_opt not being const is a genuinely useful thing to have written down.

One thing I want to check, plus two smaller notes.

spark_cast_decimal_to_boolean now fails on a non-all-null precision-0 array

if decimal_array.null_count() == decimal_array.len() {
    return Ok(Arc::new(BooleanArray::new_null(decimal_array.len())));
}
let zero = Scalar::new(
    Decimal128Array::from(vec![0i128])
        .with_precision_and_scale(decimal_array.precision(), decimal_array.scale())?,
);

The comment says Decimal128(0, 0) is reachable through Spark's RDD row-to-Arrow path, and the all-null fast path handles it. But if such an array ever contains a non-null value, with_precision_and_scale(0, 0) returns Err and the whole cast fails, where the old BooleanBuilder loop handled it fine.

Is a non-all-null precision-0 array reachable? Spark's DecimalType requires precision at least 1, so my guess is no and the all-null case is a degenerate placeholder. If that is right, the comment should say so, because as written it explains why the fast path exists without saying why the slow path is safe. If it is not right, the zero scalar needs to be built at a clamped precision instead.

is_not_null changes the result length

The old code built a BooleanArray of exactly batch.num_rows() by indexing src_value.is_valid(row). is_not_null(&src_value) returns an array of src_value.len(). Those are the same whenever the evaluated column is materialized to the batch length, which I expect is always. But if a scalar ever came through without expansion, the old code would have produced a full-length mask and the new one produces a length-1 mask, and the subsequent and would fail with a length mismatch rather than silently misbehaving.

That is arguably an improvement, but it is a behavior change worth a sentence. Is there a guarantee upstream that src_array_expr.evaluate always yields a batch-length array here?

One question on spark_pow

The kernel comment says scalar/array uses unary so the scalar is not broadcast, with an explicit null-scalar short-circuit because unary preserves the input array's null buffer rather than the scalar's. What about array/scalar, with the scalar on the right? The comment only describes the left-scalar case. If both directions are handled the wording could say so, and if test_spark_pow_null_scalar only covers one direction it would be worth covering the other.

[0lai0]

Thanks @andygrove for review, all three addressed:

1. benches/spark_pow.rs
   pipeline now includes both pow(a + 2.5D, 3) (array/scalar dispatch) and pow(a + 2.5D, b) (array/array dispatch) with a non-null array of fractional exponents so the powf work stays measurable. Timing wraps the add + spark_pow composition, not a pre-materialised intermediate.

2. array_insert.rs
   reworded the is_not_null comment to state the forward-looking length invariant (into_array(batch.num_rows()) broadcasts scalars, so the mask matches batch length; without it downstream and/evaluate_selection would fail on length mismatch).

3. pow.rs / CometExpressionSuite.scala
   the null-scalar short-circuit doc and test_spark_pow_null_scalar both cover null on either side; Scala regression test uses checkSparkAnswerAndOperator.

[andygrove]

The description fix and the two follow-ups on #5091 from my first pass are all there. The dense-null explanation under the spark_pow benchmark table is reworded correctly, and the days_to_date and covariance findings are copied over to #5091 with the benchmark numbers attached.

I traced the two substantive changes since then. spark_cast_decimal_to_boolean now reads the raw i128 payload directly when precision() == 0 instead of routing through with_precision_and_scale, so the BigInteger.ZERO via DecimalVector.setSafe case sunchao traced through the JVM UDF codegen path round-trips correctly, and the existing all-null fast path still covers the RDD-scan case. That is backed by unit tests that hand-build the precision-zero array, including the mixed valid-and-null row, plus a CometExpressionSuite regression test that drives the actual Java UDF path through checkSparkAnswerAndOperator so a fallback to Spark would fail the test.

For spark_pow, I checked pow_scalar_array_null_aware, pow_array_scalar_null_aware, and pow_binary_null_aware against NullBuffer::union and valid_indices. The is_dense_null threshold only changes which loop runs. It does not change which slots get read or how nulls propagate, so the 75% crossover is a performance choice, not a correctness one. The length-mismatch and payload-in-null-slots tests exercise the new path directly, and the new composed-null benchmarks back up the numbers in the thread.

The two other things I flagged last review are settled as well. The spark_pow_kernel doc comment now says the null-scalar short-circuit applies on either side, with test_spark_pow_null_scalar covering both directions, and the array_insert.rs comment on is_not_null now states the length invariant explicitly. Approving.

[andygrove]

@0lai0 this is conflicting with main now. Could you rebase? The workflows are also sitting at action_required, so nothing has actually run on this head yet.

[0lai0]

Rebased onto latest main

[sunchao]

Reviewed cafe924fce600858c6f2b3d4549b7fb6ec18b21d against fad6230948032012d676bed23156b7c93dd2fa2e.

The earlier precision-zero decimal findings are addressed. I found no new correctness discrepancy in decimal-to-boolean casting, ArrayInsert, or pow. One performance case remains after the dense-null fix: independently nullable power operands can have mostly-null output while the dispatcher still evaluates every raw value. Details and reproduced measurements are in the inline comment.

CI

At submission, 7 checks have succeeded, 9 are skipped, and 12 are queued or running. No failures are reported, but CI is incomplete. Current CI run.

Local validation

• Power: 114,657 base/head call comparisons covering 6,759,600 output rows passed, including distinct slices, threshold-adjacent null densities, scalar/array combinations, NaNs, infinities, signed zeros, subnormals, and arbitrary floating-point payloads. Another 25 malformed invocations returned errors. All 12 current-head power unit tests passed in the component harness.
• Decimal: 180,810 base/head comparisons and 72 precision-zero Arrow FFI round-trips passed.
• ArrayInsert: 13,824 base/head comparisons passed across List/LargeList, slices, empty batches, scalar/column/null children, and both legacy modes. Guarded child expressions checked which rows were evaluated.
• Formatting and diff whitespace checks passed.

The full native expression test command stopped before compilation because the configured dependency mirror lacks pinned DataFusion 55.1.0. The full native build, Clippy, and JVM regression test were therefore not completed on this head. Component tests used exact current expression sources, pinned Arrow 59.3.0, and available DataFusion 55.0.0. I fetched the official 55.0.0 and 55.1.0 datum.rs, physical_expr.rs, and columnar_value.rs and verified that each pair is byte-identical. This is component validation, not a successful full 55.1.0 build or end-to-end Spark run.

[sunchao]

[P2] Use the combined null count for binary pow dispatch

Could we base this dispatch on the combined null mask and reuse it for evaluation? With independent null patterns, each operand can be below 75% null while over 90% of output rows are null. This still selects binary(), which computes powers for every masked payload.

For pow(a + 2.5D, b + 2.5D) with zero-filled input null slots, the additions leave 2.5 under the preserved null bits. The new path evaluates those expensive powers, while the base implementation skips them.

I compared the exact base/head power sources in optimized x86_64 Linux pipelines with 8,192 rows, both additions inside the timed region, thin LTO, one codegen unit, CPU affinity, and warm-up. Seven randomized base/head/candidate triples per regression case were repeated with another input seed:

Nulls per operand	First run paired slowdown	Repeat
About 70%	42%	42%
About 74%	64%	59%

Every regression pair was slower on the head. Computing and reusing the combined mask removed the slowdown. No-null and aligned-dense controls improved on this PR. Raw zero-payload inputs also improved, and adding only the left operand was approximately unchanged. The current benchmarks use aligned masks or a non-null second operand and miss this independently nullable composition.

These are local expression-pipeline measurements, not whole-query throughput. The harness used pinned Arrow 59.3 and DataFusion 55.0 because the mirror lacks 55.1. The relevant DataFusion dispatch/type files are byte-identical between those releases. Could we add this independent-mask composed-expression benchmark alongside the fix?