Skip to content

perf: reuse validity buffers and use unary kernels in native expressions - #6953

Open
comphead wants to merge 3 commits into
apache:mainfrom
comphead:validity-mask-reuse
Open

comphead wants to merge 3 commits into
apache:mainfrom
comphead:validity-mask-reuse

Conversation

@comphead

@comphead comphead commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

There is no issue for this change. It continues the native Rust cleanup in #6861 with two patterns: validity (null) masks rebuilt row by row, and per-element iterator chains where an arrow unary kernel fits. The negative-scale decimal guard came out of review.

Rationale for this change

Several native expressions do more work per row than needed:

  • They rebuild a null mask that the input already carries. They append one bit per row to a Vec<bool> or a builder and pack it at the end, although the output is null exactly where an input is null. A Vec<bool> also uses a byte per row. nested_extrema collects a Vec<Option<u32>> only to convert it into an array afterwards.
  • They map a primitive array with iter().map(..).collect(), building an Option or a Result<Option<_>> for every element. PrimitiveArray::unary maps the values buffer in one pass and keeps the null buffer, and try_unary does the same for fallible conversions.

What changes are included in this PR?

The PR has three commits.

Reuse validity buffers

Reuse the input's nulls where the output is null exactly when the input is:

  • from_json.rs: the result struct is null only for a null input row, so StructArray::try_new_with_length takes string_array.nulls().cloned() instead of a Vec<bool> filled in the row loop. try_new_with_length is used because a struct with no fields needs the row count when there is no null buffer. The nested-struct builder has no input to borrow from, so its mask becomes a BooleanBufferBuilder instead of a Vec<bool>.
  • regexp_extract_all.rs: the result is null exactly where the subject is, so it reuses array.nulls().
  • split.rs: split_sql_generic_scalar_array reuses the delimiter array's nulls and split_sql_generic_array uses NullBuffer::union of the string and delimiter nulls, instead of a NullBufferBuilder each. split_generic and split_sql_generic_scalar already worked this way.
  • array_insert.rs: a row is null when its list or its position is null, so NullBuffer::union(list.nulls(), pos.nulls()) replaces the Vec<bool> that the row loop filled.

Build an array without an intermediate Vec:

  • array_extrema.rs: nested_extrema appends to a UInt32Builder with append_option, as float_extrema in the same file does, instead of collecting Vec<Option<u32>> and calling UInt32Array::from.

When the input has no nulls, the output now has no null buffer instead of an all-valid one. The two are logically equal.

arrays_zip.rs has the same Vec<bool> pattern and is left alone on purpose. The function is a copy of DataFusion's arrays_zip_inner (the file says so and has a TODO to reuse it), so changing it here would drift from upstream.

Use unary and try_unary for integral casts and unix_timestamp

  • numeric.rs: the macros cast_float_to_int16_down, cast_float_to_int32_up, cast_decimal_to_int16_down and cast_decimal_to_int32_up built an Option or a Result<Option<_>> per element. The legacy float conversions cannot fail and use unary. The ANSI conversions use try_unary, which skips null slots and returns the first overflow in index order, as the collect::<Result<_>>() did. The legacy decimal conversions use try_unary with a closure that always returns Ok, because the division is most of their cost and unary would divide the slots under nulls too. cast_int_to_int_macro! in the same file already has this shape. The decimal macros also compute the power of ten once per array instead of once per value.
  • unix_timestamp.rs: the three copies of an if null_count() == 0 { values().iter()... } else { iter()... } branch become one unary call each. The closure divides by a constant, so running it under a null cannot panic.

Keep negative-scale decimals out of the native integral casts

  • CometCast.canCastFromDecimal now reports a cast from a decimal with a negative scale to ByteType, ShortType, IntegerType or LongType as Unsupported in every eval mode, so Spark's own cast runs. Negative-scale decimals only exist with spark.sql.legacy.allowNegativeScaleOfDecimal=true. The native kernels divide by 10^scale, which overflows for a negative scale. On main that already panics on a debug build as soon as a batch has a value (the same family as Native panic casting to negative-scale decimal when spark.sql.legacy.allowNegativeScaleOfDecimal=true #5013). Hoisting the divisor made an all-null batch fail too, which main answered with nulls, and in a release build the wrapped divisor is zero. A docs section describes the fallback.

How are these changes tested?

The results do not change, so existing tests cover them.

Validity buffers: test_simple_struct, test_empty_and_null_json (no nulls in the input) and test_nested_struct for from_json, null_subject_returns_null for regexp_extract_all, test_split_sql_scalar_string_array_delimiter for the scalar-string split, test_array_insert, and the SQL files array_insert.sql and array_extrema_floating_point.sql (arrays of arrays and of structs reach nested_extrema).

Casts: test_cast_double_to_int_ansi_boundaries, test_cast_double_to_long_ansi_boundaries, test_cast_float_to_int_ansi_boundaries, test_cast_float_to_long_ansi_boundaries, test_cast_float_to_short_and_byte_ansi_boundaries and test_cast_float_to_integral_legacy_saturates for the float macros, cast_decimal_to_primitive.sql for the decimal macros, and test_unix_timestamp_from_timestamp, test_unix_timestamp_truncates_fractional_seconds_toward_zero, test_unix_timestamp_from_date and test_unix_timestamp_with_nulls for unix_timestamp.

Three tests are new:

  • test_split_sql_string_array_delimiter_array: the array-string by array-delimiter path of split_sql had no test. It uses a NULL string and a NULL delimiter.
  • test_cast_to_integral_ignores_values_under_nulls: a value that would overflow sits under a null slot and must neither fail an ANSI cast nor reach the result. It covers the float and decimal macros in both eval modes.
  • cast DecimalType with negative scale to integral types is not supported natively in CometNativeCastSuite: it asserts the cast is Unsupported for every eval mode and compares with Spark for a values table and an all-null table, for all four integral targets. It excludes ConvertToLocalRelation, because otherwise Spark evaluates the projection over the in-memory rows on the driver and the cast never reaches Comet. With the guard removed it fails with the native panic attempt to multiply with overflow.

Locally, cargo clippy --all-targets --workspace -- -D warnings, cargo fmt --check, the 116 conversion_funcs unit tests, the new Scala test together with the existing negative-scale test, spotless:check and prettier --check pass. The other suites were not run locally.

Benchmarks

Measured with criterion on an Apple M3 Max, with a 1 s warm-up and about 2 s of measurement per case, against upstream/main at 55ba50d451. The bench profile used 16 codegen units and no LTO (not the release profile), the same on both sides. Run-to-run noise on identical code was about 1% with a maximum of 3%, so differences under about 3% are inconclusive. These are single-expression microbenchmarks, so query-level gains will be smaller.

Case (8192 rows unless noted) main this PR
unix_timestamp (NTZ input), no nulls 6.7 µs 2.1 µs
unix_timestamp (NTZ input), 10% nulls 21 µs 2.1 µs
unix_timestamp (NTZ input), 65536 rows, no nulls 91 µs 14 µs
Legacy f64 to i32, no nulls / 10% nulls 25 / 33 µs 1.9 / 1.9 µs
Legacy decimal(18,2) to i32, no nulls / 10% nulls 47.5 / 58.5 µs 16.2 / 18.4 µs
ANSI f64 to i32, no nulls / 10% nulls 32 / 40 µs 27 / 24 µs
ANSI decimal(18,2) to i32, no nulls / 10% nulls 51 / 69 µs 31 / 29 µs

Other expressions, by input null density (no nulls / 10% nulls / all null), over 8192 to 524288 rows:

  • array_insert (10 element types): -4% to -12% / -1% to -10% (several cases within noise) / -18% to -41%.
  • split_sql with array delimiters: -1% to -4% / -3% to -6% / -33% to -36%.
  • regexp_extract_all: -2% to +2% (criterion flags one +2% case as slower, which is inside the noise floor) / within noise / -31% to -32%.
  • from_json: -1% to -4% in most cases, close to the noise floor.
  • array_extrema: inconclusive. The nested cases that this PR touches ranged from -21% to +8%, and the float cases that it does not touch moved by -8% to +11%, so this bench is too noisy to show an effect.

The existing benches did not cover the split_sql array-delimiter variants or float and decimal to integer casts. I measured those with temporary cases that are not part of this PR.

…er row

from_json, regexp_extract_all and the split_sql array variants appended one
validity bit per row to a Vec<bool> or a builder, although the output is null
exactly where the input is. Reuse the input null buffer instead.

array_insert derives the same mask with NullBuffer::union of the list and
position nulls, and nested_extrema in array_extrema builds its indices with a
UInt32Builder instead of collecting a Vec<Option<u32>>.

Adds a test for the array/array split_sql path, which had none.
@github-actions github-actions Bot added enhancement New feature or request performance area:expressions Expression evaluation labels Oct 14, 2026
The float and decimal to integral cast macros built an Option or a
Result<Option> per element to feed FromIterator. The legacy conversion cannot
fail, so it now uses PrimitiveArray::unary, and the ANSI conversion uses
try_unary, which skips null slots and returns the first overflow in index
order, as before. cast_int_to_int_macro already has this shape. The decimal
macros also compute the power of ten once per array instead of once per value.

unix_timestamp had three copies of a null-count branch that unary replaces.

Adds a test that a value which would overflow under a null slot neither fails
an ANSI cast nor reaches the result.
@comphead comphead changed the title perf: reuse input validity buffers instead of rebuilding null masks per row perf: reuse validity buffers and use unary kernels in native expressions Oct 14, 2026

@sunchao sunchao 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.

Reviewed the full two-commit, seven-file diff at a9ce673852538cfd4caca40285368e267e410ea9 against 55ba50d45137781f50ec1b7323da7109e66f36f9. The PR is not a draft. Snapshot and live checks found no existing reviews, conversation comments, inline comments, or review threads. One introduced P2 issue was found.

Routed skills: review-comet-pr, review-comet-expression-pr, and optimize-comet-expression for optimization validation guidance.

Summary

  • Prior state and problem: Several expressions rebuilt existing validity masks or collected intermediate optional values. Numeric casts and unix_timestamp also used per-element iterator plumbing.
  • Design approach: Reuse input null buffers, combine binary-input validity with NullBuffer::union, append extrema indices directly to a builder, and use Arrow unary/try_unary kernels.
  • Correctness: The bitmap substitutions preserve the existing null conditions, and the extrema comparator is unchanged. Float cast boundaries, ANSI null-slot skipping, and timestamp truncation remain equivalent in focused checks. However, legacy decimal casts now evaluate null slots and panic for an all-null negative-scale input. Spark and the base return nulls for this input.
  • Compatibility analysis: Compared relevant Spark implementations across the supported profiles: 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Negative-scale decimals are reachable with spark.sql.legacy.allowNegativeScaleOfDecimal=true, and Comet admits their integral casts. Comparison with branch-1.1 at 8ac6fccde53d96fec22184be287747eeae481cac confirms the reported regression also changes released behavior. Spark-specific extrema ordering already present in the base is outside this PR's introduced changes.
  • Key design decisions: try_unary correctly skips null slots in fallible conversions. StructArray::try_new_with_length preserves row counts for empty schemas. The assumption that decimal division is always safe under unary does not hold for negative scales.
  • Implementation sketch: Changes stay within seven expression files. Nested JSON validity uses BooleanBufferBuilder, nested extrema uses UInt32Builder, and the new tests cover split validity and hidden overflowing values under nulls.
  • Performance: Buffer reuse and direct builders remove identifiable allocation work. An isolated decimal-cast timing sample improved the tested nonnegative-scale shapes. This does not establish gains across all expressions or complete queries, and the PR supplies no benchmark results. No reproducible P1/P2 performance regression was identified.
  • Design: The local substitutions are straightforward and retain existing expression boundaries. No separate P1/P2 design issue was identified. Decimal null short-circuiting needs preservation before applying the optimization.
  • Abstraction & complexity: Existing Arrow primitives replace manual bookkeeping without introducing another abstraction layer. The explicit struct length is justified, and retaining the current cast macros keeps the change localized.
  • Behavioral changes worth calling out: Some outputs can omit redundant all-valid bitmaps. Infallible kernels now process physical values beneath nulls. That is harmless for the changed float casts and timestamp division, but produces the reported negative-scale decimal failure.
  • Suggested improvements: Preserve the null-only fast path before calculating the decimal divisor in both macros, or provide a safe negative-scale path. Add an all-null negative-scale regression case covering all four integral targets.

Exact-head CI: 24 checks succeeded, 18 were skipped, and two failed. Native CI passed 2,427 tests, including tests for all changed areas and both new tests. The Spark 4.1 expression job failed resolving surefire-junit4:3.5.4 before running tests, which also failed Required Checks. Spark SQL and macOS jobs were skipped.

Validation limits: The local Comet test build failed both offline and online because the configured Cargo mirror lacks DataFusion 55.2.0. A disposable release-mode harness compiled verbatim base/head cast macros and dispatch code against lockfile Arrow 59.3.0, with lightweight error stubs. Four focused tests passed, including an assertion that reproduces the new panic, plus boundary, slicing, bitmap, empty-struct, builder, and timestamp checks. No full local Spark/Comet run or comprehensive benchmark matrix was completed. Project code remains unchanged, and nothing was published.

}
Ok(truncated as $rust_dest_type)
})?,
_ => cast_array.unary(|value| (value / divisor) as $rust_dest_type),

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.

[P2] Preserve null short-circuiting for negative-scale decimal inputs. With spark.sql.legacy.allowNegativeScaleOfDecimal=true, an all-null DecimalType(38, -1) column can reach this cast because CometCast.canCastFromDecimal accepts integral targets without a scale guard. Previously the iterator skipped every null and returned nulls, matching Spark. Now 10_i128.pow(scale as u32) wraps to zero in release builds, and unary divides by that zero even beneath null slots, failing the query. The same regression affects the byte/short macro at line 546. Please return the null-only result before calculating the divisor in both macros, or handle negative scales safely, and cover this input in regression tests.

Evidence: Reproduced using verbatim decimal-cast macros from the specified base and head with Arrow 59.3.0 and release overflow checks disabled, matching the repository profile. Input: Decimal128Array::from(vec![None::<i128>; 3]).with_precision_and_scale(38, -1).unwrap(). Calling spark_cast_nonintegral_numeric_to_integral with EvalMode::Legacy and each of Int8, Int16, Int32, and Int64 returns three nulls on the base and panics with attempt to divide by zero on the head. Harness: /tmp/pr6953-kernel-review-7d2rw4x3, test negative_scale_nulls_regression. Spark's UnaryExpression.eval returns null before invoking the cast, and Comet's cast dispatcher has no equivalent all-null guard. Full JVM/native execution was unavailable because of the documented dependency-resolution limitation.

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.

I checked this end to end and it goes a bit further than the all-null case. On main, a negative scale already panics as soon as a batch has a non-null value. col("_1").cast(DecimalType(10, -2)).cast(IntegerType) over a Parquet double column runs natively as a CometProject over CometNativeScan and fails on a debug build with attempt to multiply with overflow, while Spark returns 12300. That is the same family as #5013. What this PR adds is the all-null batch. Main returns nulls there, like Spark, and with the PR it panics too. On a debug build that happens in both eval modes. On a release build the divisor wraps to zero, so LEGACY panics because unary divides the null slots by it.

So a null-only fast path would only hide half of it. Could we keep negative scales out of these macros instead? If canCastFromDecimal in CometCast.scala returned Unsupported for a negative scale with an integral target, the codegen dispatcher would run Spark's own cast, and the hoisted divisor could never wrap. A test in CometNativeCastSuite with a negative-scale column, one table with values and one that is all null, would cover both.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thanks, this is fixed in 0ab5db9. I followed andygrove's suggestion below instead of a null-only fast path. A negative scale already panics on main (debug build) as soon as a batch has a non-null value, so the fast path would only have covered the all-null half.

CometCast.canCastFromDecimal now returns Unsupported for a negative-scale decimal with a ByteType, ShortType, IntegerType or LongType target in every eval mode. Spark's own cast runs through the codegen dispatcher, so the hoisted divisor is never computed for such a column. The cast docs have a short section on it.

The new test in CometNativeCastSuite casts a negative-scale column twice, once with values and once all null, to the four integral targets. It also asserts the Unsupported result for the three eval modes. One thing I ran into while writing it: Spark's ConvertToLocalRelation folds Project(LocalRelation) on the driver, so by default the query never reaches the Comet cast and the answer check cannot fail. The test excludes that rule. With the rule excluded and the guard removed, it fails with the native panic attempt to multiply with overflow.

Locally I ran clippy, the 116 conversion_funcs unit tests, and this test together with the existing negative-scale cast test. I only reproduced the debug build panic. I did not run a release build, so the divide by zero in your report is not re-checked here. The guard keeps the cast out of the native kernels, so it should not depend on the build profile.

@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.

I ran the suites the failed [expressions] job never got to against a merge of this PR onto main (Spark 4.1, JDK 17): CometNativeCastSuite, the SQL files for each touched expression, and CometJsonExpressionSuite. They all pass, and the null-mask changes keep the same null positions everywhere I checked. The blocker is the negative-scale decimal cast on sunchao's thread, and I left a perf suggestion for the legacy decimal arms. The [expressions] job failed on a Maven download from Central before any test ran, so it needs a rerun.

})
.collect::<Result<$dest_array_type, _>>()?,
})?,
_ => cast_array.unary(|value| (value / divisor) as i32 as $rust_dest_type),

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.

The legacy decimal arms divide every slot here, including the ones under a null, and the i128 division is most of the cost. I copied the old and new macro bodies into a small harness (arrow 59.3, release, 8192 rows, DECIMAL(38,2) to INT). This PR runs at about 18 µs a batch at every null density, against 44 to 70 µs before, so it's a clear win already. try_unary with a closure that always returns Ok only visits valid slots, and it came in at 16.4 µs with no nulls, 10.1 µs at 50% nulls and 0.4 µs for an all-null batch. Would it be worth using it for the two legacy decimal arms? The float arms are fine with unary, since the as cast is nearly free and vectorizes (1.9 µs at any density, down from 17 to 31 µs).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Done in 0ab5db9. The two legacy decimal arms now use try_unary with a closure that always returns Ok, and a comment says why. The float arms stay on unary.

My numbers for decimal(18,2) to i32 at 8192 rows (criterion, M3 Max, bench profile): with no nulls it is 47.5 µs on main, 18.6 µs with unary and 16.2 µs with try_unary. At 10% nulls it is 58.5 µs on main and 18.4 µs with either, so I see no difference there. I did not measure 50% or all-null batches, so I cannot confirm your 10.1 µs and 0.4 µs figures.

The native decimal to integral casts divide by 10^scale, which overflows for a
negative scale (only creatable with spark.sql.legacy.allowNegativeScaleOfDecimal=true).
A debug build panics on any batch that has a value, and a release build wraps the
divisor to zero, so the hoisted divisor also failed on an all-null batch that used
to return nulls. CometCast now reports these casts as Unsupported in every eval
mode, so Spark's own cast runs.

The legacy decimal arms also use try_unary with a closure that always returns Ok,
so the division only runs for valid slots.

Adds a CometNativeCastSuite test with a values table and an all-null table for all
four integral targets, and a docs note.

@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 the quick turnaround. The guard keeps negative-scale decimals out of the native integral casts, the new test covers both the values and all-null cases, and the legacy decimal arms now skip null slots with try_unary. Looks good to me.

@andygrove

Copy link
Copy Markdown
Member

@comphead could you fix conflicts?

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

Labels

area:expressions Expression evaluation enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants