perf: build spark_size LargeList lengths from i64 offsets - #5300
Open
0lai0 wants to merge 2 commits into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #5272
Rationale for this change
Follow-up to #5233. That PR routed all list-like
spark_sizethrough Arrow'slengthkernel, which returnsInt64forLargeListand then needs acast_with_options(..., Int32, safe: false)(and ato_vec()when patching null slots). That left an extra Int64 length array and Int32 cast on the LargeList path, and the bench showed it:LargeList (10% null)was only ~1.3x faster than main after #5233, while the pureListpath was ~12x.Skip the length kernel entirely for
LargeList. Subtract adjacent i64 offsets straight intoInt32lengths, then apply the shared null →-1rewrite. No intermediateInt64Array, no cast allocation.CometSize.convertstill wrapssizeinCASE WHEN isnotnull(child), so the production shape is the no-null path (693 ns here, within ~2x of the pureListno-null path at 392 ns).Overflow on LargeList now surfaces as
size(): list length exceeds i32::MAX(same as the scalar path) instead of the Arrowcast_with_options(safe: false)error. Spark arrays are Int-indexed and cannot exceedi32::MAXelements per row, so the overflow branch is unreachable from a Spark plan today; it is kept as a defensive guard for non-Spark producers.What changes are included in this PR?
LargeListgets its own pathspark_size_large_list_from_offsets, splitting off the sharedspark_size_list_like(which now handles onlyListandFixedSizeList).i32):offsets.windows(2).map(|w| (w[1] - w[0]) as i32).collect(). Sound because Arrow offsets are monotonically non-decreasing, so no per-row length can exceed the full span.spark_size_large_list_lengths_checkedfor the rare case where the offset span exceedsi32::MAX. Usesi32::try_fromper row and errors on overflow (matches the oldsafe: falsecast contract). Skipstry_fromon null rows since the caller overwrites them with-1anyway.-1rewrite intoints_with_nulls_as_neg_onesoList/FixedSizeListandLargeListcannot drift.SIZE_OVERFLOW_MSGconstant used by both the array and scalar paths.create_large_list_arrayonwith_nulls, addspark_size: LargeList of long arraysandspark_size: LargeList, no nullsshapes to match theListcoverage.docs/source/contributor-guide/expression-audits/collection_funcs.md(date, PR, technique, speedup, benchmark file) peroptimizing_expressions.md.Benchmark (
array_size, 8192 rows, on top of #5233)Numbers from criterion comparison against a
pre-5272baseline saved on main (macOS).ListandFixedSizeListbenches are unchanged as expected (their code path did not move).How are these changes tested?
spark_sizeunit tests plus three new ones:test_spark_size_sliced_large_list_array: pins slicing behavior (mirrors theListslice test).test_spark_size_large_list_length_overflow: single-rowi32::MAX + 1length errors, exercising the checked fallback directly via a validOffsetBuffer.test_spark_size_large_list_checked_null_row_skips_overflow: null rows with an overflowing offset delta must not error (they get rewritten to-1by the caller, same as the fast path).cargo test -p datafusion-comet-spark-expr --lib -- spark_size(14 passed).cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings.Are there any user-facing changes?
No. Values remain bit-identical for in-range lengths; null still returns
-1. The only observable change is the overflow error message text (see Rationale), which is not reachable from Spark's Int-boundedSize.