perf: optimize spark_base64 in spark-expr#4885
Conversation
mbutrovich
left a comment
There was a problem hiding this comment.
The scratch-buffer reuse is the real win and it stands on its own. One concern with the capacity precompute in encode_array:
let data_capacity: usize = (0..array.len())
.filter(|&i| array.is_valid(i))
.map(|i| base64_encoded_len(array.value(i).len()))
.sum();Two issues:
-
This adds a full O(n) pass over the array before the encode loop makes its own pass. For a change whose goal is removing per-row allocations, the extra walk is a cost that partly offsets the win.
-
base64_encoded_lenreturns the unwrapped length. Whenchunkis true (the Spark default), the output includes CRLF separators (about 2 bytes per 76 chars), so the builder under-reserves and can still reallocate on long chunked input. The PR description says it "sizes the value buffer up front", which only holds for the unchunked path.
I would drop the precompute entirely. The builder's amortized doubling is fine, and removing it saves the extra pass. If you want to keep it, make the estimate chunk-aware so the reservation is actually tight for the default path.
Separately, the Evidence block in this PR is clean (all four reported benchmarks exist in the added base64.rs), which is worth noting since the sibling perf PRs had reported numbers that do not map to committed benchmarks.
| fn chunk_into_lines_buf(encoded: &str, out: &mut String) { | ||
| let separators = (encoded.len() - 1) / LINE_LEN; | ||
| let mut out = String::with_capacity(encoded.len() + separators * 2); | ||
| out.reserve(encoded.len() + separators * 2); |
There was a problem hiding this comment.
should we reserve every row? perhaps we can do a reusable buffer?
What changed
Optimizes
spark_base64in the nativespark-exprcrate.Reuse scratch String buffers for base64 encode + line-wrap and build directly into a preallocated StringBuilder, eliminating two per-row heap allocations in the array path.
Evidence
main).