perf: Optimize string split functions and avoid intermediate heap allocations

[kazantsev-maksim]

perf: Optimize string split functions and avoid intermediate heap allocations

Which issue does this PR close?

• N/A

Rationale for this change

Optimize existing expression.

spark_split compiled a regex::Regex even for plain literal delimiters, its limit == 0 path collected a fresh temporary vector per row, and both scalar paths materialized owned intermediate strings before building Arrow arrays. It also trimmed trailing empty parts for limit == 0, which Spark does not do. This PR removes that matching and allocation overhead and fixes the limit == 0 semantics.

What changes are included in this PR?

1. Fast path for literal / single-char delimiters

• is_regex_literal classifies the pattern conservatively: any of . ^ $ * + ? ( ) [ ] { } | \ keeps it on the regex path.
• Non-regex patterns bypass the regex::Regex DFA. Single-character delimiters use the standard library's character-pattern split; multi-character literals use str::match_indices.
• push_split_char advances by the delimiter's UTF-8 byte length, so multi-byte delimiters are handled without re-encoding.
• Regex patterns keep the previous behavior; test_literal_and_regex_helpers_agree cross-checks that the literal and regex helpers produce identical output for a delimiter both can handle.

2. Correct limit == 0 semantics (Spark parity)

limit == 0 is normalized to -1 once, at the top of spark_split, before the split helpers are called, with the Spark source reference in a comment next to it. Spark keeps trailing empty parts for limit <= 0; the helpers no longer trim them.

This is a correctness fix, not a performance change. split('a,b,c,,', ',', 0) now returns ["a","b","c","",""], matching Spark 3.5 and master, and is pinned by test_split_limit_zero_keeps_trailing_empties_like_spark.

3. Removed the scratch buffer

Once limit == 0 is remapped before the helpers run, push_split_parts, push_split_literal and push_split_char never read the scratch: &mut Vec<&str> parameter: both arms of each append straight into the offset and value buffers. The parameter, the three per-batch Vec::new() allocations, and the now-unused lifetime parameter on the two literal helpers are removed. This also clears the clippy unused_variables and ptr_arg findings on those three sites, and Arc::clone is used in place of .clone() on a ref-counted pointer.

4. Pre-allocated Arrow builders

String value and offset buffers are pre-allocated from the input batch size and total byte length (value_data().len()), avoiding repeated reallocations as the buffers grow.

5. Optimized scalar branches

Scalar inputs build the Arrow buffers directly from borrowed string slices instead of going through an owned Vec<String>. The scalar SQL-split branch assembles offsets and values in place, and the three new_unchecked constructions carry SAFETY comments naming the two invariants they depend on: offsets are monotonic and are pushed after the corresponding bytes, and the values buffer is valid UTF-8 because every slice comes from a
&str.

Note that these branches previously emitted one item per character for an empty delimiter, which does not match StringSplitSQL. That was caught in review and is fixed here: an empty delimiter now appends the whole string once, matching Spark and the array paths. This is not a user-facing change relative to main — main was already correct — it is an internal regression that was fixed before merge.

How are these changes tested?

• cargo test -p datafusion-comet-spark-expr string_funcs::split — 21 tests pass, including the value-asserting tests for each helper.
• cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings — clean.
• New SQL-file coverage in split_rust.sql for a \d+ regex pattern and a multi-byte delimiter.
• Existing expression and SQL-file suites run as before.

Benchmarks

native/spark-expr/benches/split.rs already covered spark_split and spark_split_sql, but only on one shape: a scalar "," against a comma-only payload, with no limit argument. That is the shape the fast path is supposed to win on, so it could not separate the fast path from the allocation work. This PR extends the existing bench rather than
adding a second file, with the axes the claim rests on:

• pattern class: literal_char, literal_multi_char, regex;
• limit: default vs 0 vs -1, on one shared input so the rows are comparable;
• null density, via the existing NULL_RATIOS grid from benches/common/mod.rs;
• spark_split_sql: a multi-char literal delimiter and the empty delimiter.

The extended bench was applied identically to base and head so the comparison is like-for-like. Mid estimate of the 95% CI reported; change is against main.

spark_split:

Shape	NULL	8192	65536	524288
literal_char/default_limit	no_nulls	−12.6% (675 µs)	−16.1% (5.25 ms)	−24.9% (45.4 ms)
literal_char/default_limit	sparse	−12.3% (617 µs)	−17.0% (4.73 ms)	−17.2% (41.4 ms)
literal_multi_char/default_limit	no_nulls	−29.5% (488 µs)	−27.7% (3.93 ms)	−33.1% (32.4 ms)
literal_multi_char/default_limit	sparse	−26.1% (447 µs)	−27.8% (3.61 ms)	−34.1% (29.8 ms)
regex/default_limit	no_nulls	+2.7% (1.49 ms)	+2.3% (12.2 ms)	−5.2% (97.4 ms)
regex/default_limit	sparse	+3.3% (1.37 ms)	+3.3% (11.2 ms)	−4.4% (88.5 ms)
literal_char/limit_0	no_nulls	−67.4% (391 µs)	−66.7% (3.09 ms)	−64.7% (26.4 ms)
literal_char/limit_0	sparse	−65.9% (359 µs)	−66.3% (2.82 ms)	−65.5% (24.1 ms)
literal_char/limit_-1	no_nulls	−13.0% (393 µs)	−16.6% (3.08 ms)	−19.1% (26.3 ms)
literal_char/limit_-1	sparse	−14.0% (361 µs)	−17.4% (2.83 ms)	−21.1% (23.9 ms)

spark_split_sql:

Shape	NULL	8192	65536	524288
literal_multi_char	no_nulls	−32.1% (481 µs)	−21.9% (4.13 ms)	−33.7% (33.1 ms)
literal_multi_char	sparse	−28.9% (447 µs)	−23.7% (3.71 ms)	−33.0% (30.8 ms)
empty_delimiter	no_nulls	−5.8% (91.5 µs)	−11.5% (447 µs)	−55.4% (3.86 ms)
empty_delimiter	sparse	−10.7% (80.2 µs)	−8.7% (493 µs)	−48.6% (4.07 ms)

All-null inputs, spark_split:

Shape	8192	65536	524288
literal_char/default_limit	−22.8%	−18.6%	−20.7%
literal_multi_char/default_limit	−23.2%	−15.6%	−18.4%
regex/default_limit	−5.5%	−17.6%	−21.6%
literal_char/limit_0	−22.9%	−21.0%	−20.9%
literal_char/limit_-1	−22.4%	−19.1%	−19.8%

spark_split_sql all-null rows are within noise (between +2.2% and −2.2%), which is
expected: with no non-null input the string buffers stay empty and only the null buffer is
touched.

The 16–23% improvement on all-null inputs is the bulk-NULL change: null rows contribute no
parts, so the input null buffer is reused instead of a per-row validity bit being built.
The literal fast path shows up on non-null rows, and the two effects are independent.

limit == 0 is a correctness change, not a speedup

With limit == 0 remapped to -1 before the helpers run, literal_char/limit_0 and
literal_char/limit_-1 execute the same code and must agree. Measured on one shared input,
they do:

rows / NULL	limit_0	limit_-1	gap
8192 / no_nulls	390.50 µs	393.49 µs	0.8%
8192 / sparse	358.79 µs	360.62 µs	0.5%
8192 / all_null	11.376 µs	11.220 µs	1.4%
65536 / no_nulls	3.0906 ms	3.0833 ms	0.2%
65536 / sparse	2.8245 ms	2.8253 ms	0.03%
65536 / all_null	85.367 µs	85.476 µs	0.1%
524288 / no_nulls	26.372 ms	26.326 ms	0.2%
524288 / sparse	24.081 ms	23.886 ms	0.8%
524288 / all_null	690.18 µs	677.91 µs	1.8%

Seven of nine combinations are within 1%, the largest gap is 1.8%, the sign of the gap is not systematic, and the only two rows with non-overlapping confidence intervals are degenerate all-null inputs where the split helper is never called. The remap covers all input paths exercised here.

The reported −65% on literal_char/limit_0 is therefore not an optimization. On main that same input took about 1.2 ms against roughly 452 µs for limit_-1, because the limit == 0 path collected a Vec<&str> and trimmed trailing empties through it — work the new code does not perform at all, because Spark does not ask for it. The allocation and
fast-path wins are in the default_limit rows above.

An earlier revision of this description listed literal_char_limit_0 at 39.8 µs against literal_char_default_limit at 68.5 µs. Those two rows used different payloads (data_{i},,,, versus a five-field CSV row), so the comparison was never like-for-like, and both predate the remap. Both are corrected above, and the bench file the table came from is superseded by the extension to benches/split.rs described here.

Known regression: regex patterns at 8192 and 65536 rows

regex/default_limit is 2.3–3.4% slower than main at 8192 and 65536 rows, with p < 0.05 and a tight confidence interval, while being 4–5% faster at 524288. The sign is consistent across all four non-degenerate rows at the two smaller sizes, so this is not measurement noise, and an earlier revision of this description described it as "within noise", which was wrong.

Regex patterns do not take the literal fast path, so the candidates are the buffer pre-allocation sized from value_data().len() — for \s+ the split output is smaller than the input, so capacity may be reserved that is never filled — and any difference in the iteration strategy on the regex arm for the default limit. This is being investigated by
splitting the patch into classification, pre-allocation and scratch-removal parts and benching them separately. The regex path should return to noise before merge; if it cannot, the regression will be documented with a follow-up issue rather than presented as noise.

Measurement caveats

The regex rows exceed criterion's default target time and ran with fewer than 100 samples, so they should be re-taken with an increased --measurement-time on a quiet machine. Several rows report 12–21% high-severe outliers, consistent with background load rather than with the change. The literal_char/default_limit row at 524288 has a wide confidence interval (−31.8% to −19.9%) and its point estimate should be treated as indicative only. Base and head were both run with default criterion settings, on the same machine, in one session.

[sunchao]

Summary

Reviewed immutable head 4f382219f7c1dbddeff784501a5b1b518c859f6b against base 2699f59b71788e17a2714910e166a3f83deed937 with five independent review scopes, followed by separate verification. One introduced P2 remains: scalar split_sql changes the meaning of an empty literal delimiter and can return incorrect split_part results. The inline comment includes the supported scalar-subquery case and executable before/after evidence.

Prior state and problem

The previous regex split path compiled a regex even for plain literal delimiters, and its limit-zero path collected a fresh temporary vector for each row. Scalar paths also materialized owned intermediate strings before constructing Arrow arrays. The optimization targets this allocation and matching overhead; preserving existing literal SQL-split semantics is essential because split_sql also implements Spark 4.x split_part.

Design approach

The patch classifies patterns conservatively using regex metacharacters, then selects character, substring, or regex matching. It reuses a borrowed-slice scratch vector across rows and reserves Arrow output buffers using input size estimates. Scalar results are assembled directly into Arrow buffers instead of going through the previous owned-string vectors.

Correctness / compatibility analysis

Exact-base/head native probes, real Spark 4.0.2 reference queries with ANSI on and off, and a native DataFusion/UDF/ListExtract composition independently reproduce the empty-delimiter regression. Scalar subqueries survive normal constant folding, and exact routing source shows their string results reach the changed scalar branch; the unchanged array paths still preserve the whole string. The ordinary StringSplit default JVM-backed route, its opt-in Rust regex differences, and non-default-collation fallback are unchanged and are not additional findings.

The existing split-module tests pass at both pins: six at head and fourteen at base. These are focused module tests, not a full Comet native/JVM or newly built JNI integration run. The requested 72-file two-commit comparison was inspected in full; the three-file merge-base contribution was also checked so unrelated newer-base changes were not misidentified as regressions.

Key design decisions

The literal detector leaves regex metacharacters on the regex path, while the single-character path advances by that character's UTF-8 byte length. Scratch references remain borrowed from the input until their bytes are copied into the output buffers. Array offset widths and validity propagation retain their existing structure; the verified problem is the new scalar empty-delimiter policy, not a demonstrated aliasing or lifetime failure.

Implementation sketch

split.rs adds literal/character split helpers and direct offset/value-buffer assembly, and adapts the regex helper to reuse scratch storage. The scalar SQL-split branch is rewritten separately, which is where it diverges from push_split_sql_parts. The other two contributed files register and implement the new Criterion benchmark.

Behavioral changes worth calling out

The intended changes reduce matching and allocation overhead, but scalar SQL splitting currently changes observable results for an empty delimiter: a multi-character input becomes individual characters while the column-input form stays intact. The benchmark's six array-input scenarios compile and pass smoke execution; the PR's release speedup figures were not independently reproduced. At the reviewed head, CodeQL, Delta Contrib Build Gate, and CI report action_required, each with zero jobs, so there is no successful CI execution to rely on.

Suggested improvements

Preserve the original string as one element in the scalar empty-delimiter branch, including an empty input, matching Spark and the existing array helper. Restore equivalent public-entry-point coverage for empty delimiters across scalar and array inputs, and include a scalar-subquery SQL case so constant folding cannot conceal the regression. This directly addresses the single inline P2 without expanding the scope of the optimization.

[sunchao]

[P2] Preserve the whole string for an empty SQL delimiter

StringSplitSQL treats an empty delimiter as "do not split", but this scalar branch now emits one item per character. For a Parquet table t(s) containing 'abc', SELECT split_part((SELECT max(s) FROM t), '', 1) FROM t retains the scalar subquery through normal Spark optimization, and Comet's 4.x route passes its string result to split_sql as a scalar. Exact-base/head native UDF + ListExtract probes return 'abc' versus 'a' (part 2 changes from '' to 'b'), while Spark 4.0.2 and the unchanged array paths preserve the whole input. This is the default literal SQL-split route, not the incompatible-regex opt-in. Please append the whole string once for an empty delimiter and retain a scalar/array regression test.

[kazantsev-maksim]

Thanks, i tried addressed your comments.

[sunchao]

[P2] Make both literal split loops pass the Clippy gate

This manual counter and the equivalent loop in push_split_char (line 333) trigger clippy::explicit_counter_loop. Fresh Clippy runs on the exact base/head module, with the crate's lint attributes and -D warnings, pass at the base but fail at the head on these two loops. .github/actions/rust-test/action.yaml runs cargo clippy --color=never --all-targets --workspace -- -D warnings, so these loops will fail the Rust-test check once the currently blocked workflows run. Please use .enumerate() in both helpers.

[kazantsev-maksim]

Thanks, fixed.

[sunchao]

[P2] Preserve the non-null item type in scalar SQL splits

StringSplitSQL declares containsNull = false, so this UDF is planned to return List(non-null Utf8). Setting this field to true instead returns List(Utf8) for every non-null scalar/scalar call, including ordinary nonempty delimiters. DataFusion's result-type check in debug/CI builds rejects that value before element_at: the existing split_part.sql:38 native query (split_part('a.b.c', '.', 2), with constant folding disabled by the test harness) now fails in all four Spark 4.x expression jobs (Linux 4.0 example). Fresh exact-base/previous/head typed-UDF probes reproduce pass/pass/type-mismatch. Please keep the item field non-nullable, as in the other branches, and cover the declared return type through the UDF wrapper; the new direct-helper tests bypass that check.

[kazantsev-maksim]

thanks, fixed

[andygrove]

Spark's UTF8String.split remaps limit == 0 to -1 before calling Java's String.split, with an explicit comment saying it does that to avoid Java's "drop trailing empty strings" behavior. So split('a,b,c,,', ',', 0) gives ["a","b","c","",""] in Spark, on 3.5 and on master. push_split_parts and the two new helpers push_split_literal and push_split_char in native/spark-expr/src/string_funcs/split.rs all trim trailing empties when limit == 0, so they give ["a","b","c"]. The divergence predates this PR, but the PR now carries three copies of it, and limit == 0 is exactly what the headline benchmark row measures. If the limit is normalized to -1 up front then the whole scratch-vector branch becomes unreachable and the 165% number goes with it. Would you be willing to either fix that here or open an issue for it? The buffer preallocation win holds either way, I would just rather not shape the optimization around a branch Spark never reaches.

The unit tests that checked actual split output are gone. test_split_regex, test_split_limit_positive, test_split_limit_zero, test_split_limit_negative and test_split_empty_string went away with split_string, and what is left for spark_split is test_split_basic and test_split_with_limit, which only assert matches!(result, ColumnarValue::Array(_)). All eight new tests exercise split_sql. That leaves the rewritten regex path with no coverage at all, because split_rust.sql only uses , and :: and both of those now take the new literal fast path. Could you add cases that pin the values? A regex delimiter and a multi-byte delimiter in split_rust.sql, plus a direct test asserting the literal and regex helpers agree on the same input, would cover the interesting part of the rewrite.

I went through is_regex_literal against the regex crate's default syntax and it looks right to me. Everything that carries meaning outside a class or a group is in the list, and # and whitespace only matter under (?x), which needs a ( to turn on. The three GenericStringArray::new_unchecked calls could use a SAFETY: comment though. Soundness rests on every byte reaching append_str coming from a &str and on OffsetBuffer::new rejecting a non-monotonic offset buffer, and neither of those facts is visible at the unsafe block.

[kazantsev-maksim]

Thanks for the review @andygrove. I tried addressed your comments.

[andygrove]

All three of my points are addressed. The limit == 0 remap is in the right place, at the top of spark_split with the Spark reference alongside it, and test_split_limit_zero_keeps_trailing_empties_like_spark pins split('a,b,c,,', ',', 0) to ["a","b","c","",""], which is what Spark gives on 3.5 and master. The value-asserting tests are back and better than what was removed, test_literal_and_regex_helpers_agree is exactly the cross-check I wanted, split_rust.sql now covers \\d+ and a multi-byte →, and the three new_unchecked sites carry SAFETY comments naming both facts the soundness rests on. All 21 string_funcs::split tests pass locally.

Two things left, and the first will fail CI as it stands.

Clippy is red on this head

The remap made the scratch buffer dead, which is what I expected, but the parameter and the allocation stayed behind. cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings:

error: unused variable: `scratch`
   --> spark-expr/src/string_funcs/split.rs:281:5
error: unused variable: `scratch`
   --> spark-expr/src/string_funcs/split.rs:308:5
error: unused variable: `scratch`
   --> spark-expr/src/string_funcs/split.rs:746:5
error: writing `&mut Vec` instead of `&mut [_]` involves a new object where a slice will do
   --> spark-expr/src/string_funcs/split.rs:281:14, 308:14, 746:14
error: using `.clone()` on a ref-counted pointer
    --> spark-expr/src/string_funcs/split.rs:1186:67

push_split_parts, push_split_literal and push_split_char never read scratch now: both arms of each append straight into the buffers. Dropping the parameter also removes the three let mut scratch = Vec::new() allocations at lines 110, 344 and 555, which are one per batch each and now buy nothing. That should clear five of the seven errors, and the Arc::clone at 1186 is a one-liner.

The benchmark table needs re-running

With limit == 0 remapped to -1 before the helpers see it, literal_char_limit_0 and literal_char_default_limit execute identical code. The table still shows them 1.7x apart at 1024 rows (39.8 µs versus 68.5 µs), so those numbers were measured before the remap and cannot both be current.

Worth being precise about what the remaining number means, too. The +165% on literal_char_limit_0 is a real before-and-after difference, but most of it comes from the correctness fix rather than from the allocation work: main was doing a trailing-empty trim through a scratch vector, and this PR does not do that work at all because Spark never asked for it. Presented under "Eliminated per-row heap allocations on limit = 0" it reads as an optimization win, when the honest framing is that the old path was doing something wrong and the new one skips it. The allocation and fast-path wins are the default_limit rows, and those stand on their own.

There is also no benchmark in the tree, so the table is not reproducible from the branch. Since native/spark-expr/benches already has several criterion benches, would you add this one? It is the only way the fast-path claim stays honest after the next change to this file.

is_regex_literal still looks right to me on a second pass, and I have no further correctness concerns.

[kazantsev-maksim]

I dropped scratch from the three helpers (clippy is clean), extended benches/split.rs with the limit and regex axes, and reworded the description so the −65% on limit_0 reads as the correctness fix it is rather than an allocation win — the one thing I won't call noise is the +2.3–3.4% regression on regex at 8192/65536 rows (p < 0.05), which I'm currently tracing to the buffer pre-allocation sized from value_data().len().