perf: presize collections and reuse buffers in native loops

[comphead]

Which issue does this PR close?

There is no issue for this change. It is a sweep of the native Rust code for two allocation patterns.

Rationale for this change

Two patterns allocate more than needed across native/:

• A collection is created empty (Vec::new(), HashMap::new(), String::new(), an Arrow builder) and then filled to a size that is already known, so it reallocates as it grows.
• A loop over rows or batches allocates a fresh Vec, String or writer in every iteration, where one could be reused, or the data borrowed or moved instead.

This PR fixes the instances in production code where the bound is exact or the reuse is local.

What changes are included in this PR?

Allocations removed per row:

• bloom_filter_agg.rs: update_batch built a ScalarValue per row, which allocates a String for Utf8 input. String arrays are now hashed in place.
• split.rs: with limit = 0, each row's parts were collected into a Vec so that trailing empty parts could be dropped. Empty parts are now held back until a non-empty part follows.

Allocations removed per batch:

• spark_unsafe/row.rs: process_sorted_row_partition built a ShuffleBlockWriter, which pre-encodes the IPC schema, for every batch. The writer is now kept until the batch schema changes, which happens when a dictionary-encoded column falls back to a plain array.
• spark_unsafe/row.rs: the field-major struct conversion allocated three vectors for every nested struct field. They are now allocated once per call and cleared for each field.
• explode.rs: flatten_struct_cols and the column merge in build_batch allocated a one-element Vec per column before flattening. Both now append into one vector of known capacity.
• multi_partition.rs (hash-all round robin), schema_align.rs and parquet_writer.rs borrow the batch's columns, or move them out with into_parts(), instead of copying them into a new Vec.
• spark_bloom_filter.rs: serialization cloned the bit array and then copied it twice more. It now writes each word into one presized buffer, and SparkBitArray::data() returns a slice.
• Smaller ones: presized output vectors in expand.rs and SparkPercentileGroupsAccumulator::state, an all-true MERGE mask built without a temporary Vec<bool> in merge_rows.rs, and no class_name clone per batch in JvmScalarUdfExpr.

Cold paths, changed only where the bound is exact: presized maps and sets in parquet schema matching (parquet_support.rs, schema_adapter.rs), iceberg_dictionary.rs, the builders in process_sorted_row_partition, explode's per-batch maps, a loop-invariant Schema::empty() hoisted in planner.rs, SparkPlan::new_with_additional moving its Vec instead of copying it, and a reused line buffer in analyze_trace.

Left alone on purpose:

• The RSS writer's per-frame buffers, which are sized and released inside each frame's memory reservation.
• Reusing arrow Rows across batches in range partitioning and RankLimitStream, which would keep memory that the pools do not track alive between batches.
• The per-call vector of completed batches in BufBatchWriter::write. Removing it changes which batches are dropped when a write fails.

How are these changes tested?

None of the changes alter output, so they are covered by existing tests, for example test_split_limit_zero, the bloom filter round-trip tests, CometColumnarShuffleSuite, CometNativeShuffleSuite and CometGenerateExecSuite. The only test change is an import that multi_partition.rs's test module now needs.

cargo clippy --all-targets --workspace -- -D warnings passes locally. The test suites were not run locally, and there are no benchmark numbers for these changes.

[sunchao]

Reviewed full head d156eff7b60e700dbd6dd680688688bc3f07210b against supplied base c1e00e7b92b0ada119b140b016dbefdd198bbc35. The PR is not a draft. No introduced P1/P2 issues found within this review.

Reviewed all 20 PR-changed files and the three additional direct base/head differences. Commit ancestry confirms those additional differences come from base-only fix #6310, not deletions introduced by this PR. The snapshot and live reviews, conversation comments, inline comments, and review threads contained no existing feedback.

Summary

• Prior state and problem: Native loops created temporary collections, cloned strings and column vectors, and repeatedly constructed shuffle writers even when sizes or schemas were already known.
• Design approach: Presize collections, borrow or move existing data, and reuse scratch buffers and schema-dependent writers within their existing execution scope.
• Correctness: Checked surrounding code and relevant Spark implementations for split, Bloom filters, JSON array length, percentile, and generators. String hashing still skips nulls, Bloom serialization preserves V1/V2 headers and big-endian words, and flattening preserves column order. Focused split and MERGE-mask comparisons passed.
• Compatibility analysis: Compared relevant Spark sources across supported versions 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. No configuration, shim, protobuf, or shuffle-format contract changes were introduced. The native split limit-zero difference is inherited, with Spark codegen still used by default. Comparison with branch-1.1 at e9efd9f764ee0a59b7898ff028d6985d4a7a28e1 identified no new result changes attributable to these allocation edits.
• Key design decisions: The cached shuffle writer is rebuilt when the batch schema changes. Dictionary encoding still creates a fresh stream writer per block. Nested-struct scratch vectors are cleared before each sibling field.
• Implementation sketch: Changes cover expression allocation, expand/explode/MERGE construction, Parquet schema matching and batch assembly, shuffle conversion, plan construction, and trace-line reading.
• Performance: The code removes identifiable allocations and copies, including per-string scalar materialization and repeated schema encoding. No benchmarks were supplied or run, so throughput improvements and workload-wide performance neutrality remain unmeasured. No reproducible P1/P2 performance regression was identified.
• Design: The changes remain local and preserve existing ownership and error-handling paths. Reuse is bounded by a call or partition, without introducing process-wide caches.
• Abstraction & complexity: Existing slices, Vec, into_parts(), and Arrow builders are sufficient. The schema/writer pair adds limited state with an explicit invalidation condition. No abstraction issue meeting the reporting bar was identified.
• Behavioral changes worth calling out: The native edits preserve intended results and serialization. Restoring the base's memory-manager tests produced three failures at the pinned head and none at the base, but ancestry establishes that this is the already-fixed race from base-only #6310. It is not a PR regression.
• Suggested improvements: No introduced P1/P2 issue requiring a change was identified.

Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-memory-pr, review-comet-shuffle-pr, review-comet-ffi-pr, and review-comet-iceberg-write-pr.

Exact-head CI: 26 successful checks and 15 skipped checks, with no failures or pending checks. Required Checks, all four Linux Spark 4.1 test groups, native tests, and TPC-H/TPC-DS verification passed. The Rust log reports 2,224 passing tests plus a separate five-test invocation passing.

Local validation: Extracted base/head split functions matched across 122,850 UTF-8/LargeUTF-8 cases. MERGE Boolean masks matched for 1,025 lengths. The isolated memory-manager suite passed 6/9 tests at head and 9/9 at base, explained by the base-only fix above.

Validation limits: No full local Comet build, end-to-end Spark/Iceberg run, or benchmark was performed. Upstream Spark SQL, Iceberg, PyArrow UDF, and macOS CI suites were skipped. No project source or GitHub state was changed.

[andygrove]

LGTM. I read all 20 files and compared the new split, bloom filter, writer cache and nested struct code against the old code with differential and mutation tests. I found no behavior change.

Nit, optional: the spark-expr kernels here (split, bloom_filter_agg, percentile) don't have before and after numbers, and optimizing_expressions.md asks for criterion output for a kernel change. The existing benches/split.rs never passes a limit and benches/bloom_filter_agg.rs only feeds Int64, so neither reaches the new code. If you want to add a limit = 0 case with empty parts to the split bench and a Utf8 column to the bloom filter bench, this is what I saw in my own interleaved A/B on copies of the old and new code. limit = 0 split was 11% to 46% faster across seven shapes. Utf8 hashing was 23% to 49% faster against a model of the old per-row String. I expect your numbers to agree.

The inline comments are nits as well and all optional to address.

[andygrove]

Nit, optional: nothing in the repo would notice if this branch were wrong. test_split_limit_zero goes through split_string, which calls split_with_regex, and this PR doesn't touch that helper. The only other limit = 0 test I found splits a,b,c,d,e, which has no empty parts. I broke the branch three ways (kept the trailing empties, dropped the interior ones, removed the [""] fallback) and all 16 split tests still passed. The logic itself looks right to me. I compared it with the old branch on every string up to length 8 over a three-character alphabet with a handful of patterns, about 88k cases, and they matched. Could we add a test through spark_split with limit 0 on a,b,c,,, a,,b, ,a, a,,b,, and the empty string? A split_rust.sql query with the same inputs would cover it end to end. I'd leave ,, out of the SQL file since Java returns an empty array for it and the native path returns [""].

[andygrove]

Nit, optional: the cache invalidation here has no test. If I replace the condition with block_writer.is_none(), all 172 unit tests in the shuffle crate still pass. Without the schema check, a column that goes from plain to dictionary fails the write with Ipc error: no dict id for field. One that goes from dictionary to plain writes a block that decodes to the wrong type. I wrote a throwaway test that calls process_sorted_row_partition with batch_size = 4 and a dictionary ratio of 1.5 over x,x,x,x,a,b,c,d,y,y,y,y. The column goes dictionary, plain, dictionary inside one call, and the PR decodes all three blocks correctly. Could we add a test like that? It needs only a few one-column UnsafeRows, laid out the way benches/row_columnar.rs does.

[andygrove]

Nit, optional: same question for the scratch vectors here and at line 981. Reusing them only matters when a struct has two or more nested struct fields, and I didn't see a case with that shape in CometColumnarShuffleSuite. If I remove the three clear() calls in both functions, the shuffle crate's tests still pass. A struct with a: struct<x>, b: struct<y> and m: struct<c: struct<p>, d: struct<q>>, with a null at each level, passes on the PR and fails without the clear() calls. Could the test above also cover that shape?

[andygrove]

Nit, optional: in a local A/B against the old code, this per-word extend_from_slice loop ran at about the same speed. It was a few microseconds either way from 64 KiB to 1 MiB, and the sign flipped between runs at 8 MiB. Peak memory is 3x lower, which is the real win. Filling the presized buffer with resize and chunks_exact_mut(8) ran about 2.5x faster than this loop at 1 MiB, roughly 18 µs against 50 µs, with one allocation and byte-identical output. Could we use that form?