andygrove commented on code in PR #4885:
URL: https://github.com/apache/datafusion-comet/pull/4885#discussion_r3692630530
##########
native/spark-expr/src/string_funcs/base64.rs:
##########
@@ -62,33 +64,43 @@ pub fn spark_base64(args: &[ColumnarValue]) ->
Result<ColumnarValue, DataFusionE
}
}
-fn encode_array<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>, chunk:
bool) -> StringArray {
- array
- .iter()
- .map(|value| value.map(|bytes| encode(bytes, chunk)))
- .collect()
+const LINE_LEN: usize = 76;
+
+/// Length of the padded base64 encoding of `n` input bytes.
+fn base64_encoded_len(n: usize) -> usize {
+ n.div_ceil(3) * 4
}
-fn encode(bytes: &[u8], chunk: bool) -> String {
- let encoded = BASE64_STANDARD.encode(bytes);
- if chunk {
- chunk_into_lines(encoded)
+/// Length after CRLF wrapping if `encoded_len` bytes are chunked at
`LINE_LEN` chars per line.
+fn chunked_len(encoded_len: usize) -> usize {
+ if encoded_len == 0 {
+ 0
} else {
- encoded
+ encoded_len + ((encoded_len - 1) / LINE_LEN) * 2
}
}
-/// Wrap a base64 string into lines of at most 76 characters joined by CRLF,
with no trailing
-/// separator. Matches `java.util.Base64.getMimeEncoder()`. base64 output is
pure ASCII, so byte
-/// offsets and character offsets coincide. Takes the string by value so the
common short-input
-/// case (no wrapping needed) returns it without a second allocation.
-fn chunk_into_lines(encoded: String) -> String {
- const LINE_LEN: usize = 76;
- if encoded.len() <= LINE_LEN {
- return encoded;
+/// Encodes `bytes` into `out`, wrapping at `LINE_LEN` when `chunk` is true.
`out` is reused
+/// across rows to avoid per-row heap allocations; the caller clears it before
each call.
+fn encode_into(bytes: &[u8], chunk: bool, out: &mut String) {
+ if !chunk {
+ BASE64_STANDARD.encode_string(bytes, out);
+ return;
}
- let separators = (encoded.len() - 1) / LINE_LEN;
- let mut out = String::with_capacity(encoded.len() + separators * 2);
+ // Encode into a scratch, then wrap. Two passes are unavoidable because
the base64 crate
+ // does not emit CRLF for us and computing chunk boundaries mid-encode
would require a
+ // custom writer that carries per-row state.
+ let unwrapped_len = base64_encoded_len(bytes.len());
+ if unwrapped_len <= LINE_LEN {
+ BASE64_STANDARD.encode_string(bytes, out);
+ return;
+ }
+ // Reuse `out` for the wrapped result: encode into a temporary owned by
the outer scratch,
+ // then copy CRLF-wrapped chunks in. The temporary is short-lived per row,
but the caller's
+ // long-lived scratch avoids the per-row allocation the previous
implementation had.
+ let mut encoded = String::with_capacity(unwrapped_len);
+ BASE64_STANDARD.encode_string(bytes, &mut encoded);
Review Comment:
You are right that the temporary reintroduces the allocation for exactly the
case the reuse was meant to remove, and the doc comment overstated the result.
Both fixed in b1acaefbc — but not via 57-byte windows, because I benchmarked
that and it is slower.
I implemented your suggestion verbatim first. It is correct (57 is divisible
by 3, `19 * 4 == 76`, and only the final window can carry padding, so the
concatenation is byte-identical to encode-then-split — I added an exhaustive
test for that below). But it regresses the case it targets:
| Benchmark | 57-byte windows |
| --- | --- |
| short, unchunked | -2.5% |
| short, chunked | **+3.7%** |
| long, unchunked | -6.1% |
| long, chunked | **+10.5%** |
The per-call overhead of many small `encode_string` calls (each recomputes
the encoded length and reserves before delegating to `encode_slice`) costs more
than the single bulk encode plus the copy. A 200-byte value is 4 windows, and
base64 encoding is fast enough in bulk that call overhead dominates at that
size. The allocation you flagged is real but is not the dominant cost — a
~270-byte `String` is a cheap thread-cached malloc.
So I kept the bulk encode and removed the allocation a different way:
`encode_into` now writes to a `fmt::Write` sink and takes a caller-owned
`scratch` for the unwrapped encoding. `GenericStringBuilder` implements
`fmt::Write`, so the array path passes **the builder itself** as the sink and
the CRLF-wrapped result lands directly in its value buffer. That removes the
per-row temporary *and* the copy that previously staged each row in `buf`
before `append_value`:
| Benchmark | `fmt::Write` into builder |
| --- | --- |
| short, unchunked | ~0% |
| short, chunked | -1.7% |
| long, unchunked | -7.2% |
| long, chunked | **-20.8%** |
Same baseline, same machine. So the batch now performs no per-row
allocation, which was your point, and long chunked input got 21% faster rather
than 10% slower.
I kept your window insight documented on `encode_into`, including why it is
not used, so nobody re-derives it and re-regresses the benchmark. The
exhaustive equivalence test
(`chunked_matches_encode_then_split_at_every_length`, sweeping every length
across three line boundaries) also stays, since it pins the line structure
regardless of which strategy is used.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
