mbutrovich commented on code in PR #4885:
URL: https://github.com/apache/datafusion-comet/pull/4885#discussion_r3631945206
##########
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:
This allocates a fresh `String` on every call whose encoding exceeds 76
chars, which is the entire long/chunked benchmark row. The comment on line 100
("the caller's long-lived scratch avoids the per-row allocation the previous
implementation had") overstates the result: the caller's `buf` is reused, but
this temporary reintroduces the allocation for exactly the case the buffer
reuse was meant to remove.
MIME wrapping is aligned to base64's block structure: 57 input bytes encode
to exactly 76 base64 chars with no padding (57 is divisible by 3, `19 * 4 =
76`). So you can encode the input in 57-byte windows straight into `out`, with
`\r\n` between windows, and drop the `encoded` temporary and the second pass
entirely:
```rust
let mut first = true;
for window in bytes.chunks(57) {
if !first {
out.push_str("\r\n");
}
BASE64_STANDARD.encode_string(window, out);
first = false;
}
```
One pass, zero per-row allocation, same output as the current
`chunk_into_lines` semantics.
##########
native/spark-expr/src/string_funcs/base64.rs:
##########
@@ -98,12 +110,43 @@ fn chunk_into_lines(encoded: String) -> String {
out.push_str(&encoded[offset..end]);
offset = end;
}
+}
+
+fn encode_array<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>, chunk:
bool) -> StringArray {
+ // Right-size the value buffer in O(1) from the input's total byte count.
When chunking,
+ // add the CRLF slack the encoded output will contain so the long-input
path does not grow.
+ let total_bytes = array.value_data().len();
+ let encoded_total = base64_encoded_len(total_bytes);
+ let data_capacity = if chunk {
+ chunked_len(encoded_total)
Review Comment:
`base64_encoded_len` is applied to the summed raw bytes, but each row pads
independently to a multiple of 4. For N rows of 1 byte the estimate is about
`4N/3` against an actual `4N`, so it reserves roughly a third of what is
needed, and the chunked case also misses the per-row CRLFs. This is only a
capacity hint (the builder grows correctly), so no correctness impact, but the
"sizes the value buffer up front" claim only holds for the few-large-row shape.
If you take the 57-byte-window approach above, encode straight into the
builder's value buffer and this estimate can stay as a rough lower bound
without the extra care.
--
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]
