This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 38de369fb7 perf: optimize spark_base64 in spark-expr (#4885)
38de369fb7 is described below
commit 38de369fb7ab564c65d459a924633592bb01e3b1
Author: Andy Grove <[email protected]>
AuthorDate: Mon Aug 3 17:57:17 2026 -0600
perf: optimize spark_base64 in spark-expr (#4885)
* perf: optimize spark_base64 in spark-expr
* refactor: unify base64 encoding, O(1) capacity, add array-path tests
Replaces the O(n) pre-pass over the array with an O(1) capacity estimate
from array.value_data().len(), extended with the CRLF slack when chunk is true
so the wrapped-output path no longer under-reserves. The wrap-decision is
single-sourced in a shared encode_into helper, so the array and scalar callers
cannot drift. Removes the near-duplicate chunk_into_lines/chunk_into_lines_buf
split. Adds encode_array_binary_chunked (BinaryArray with mixed nulls, empty,
exactly-76, wrap-once, wrap [...]
* review: remove per-row allocation from chunked base64, fix capacity bound
Encode straight into the output builder's value buffer. `encode_into` now
writes to a `fmt::Write` sink and takes a caller-owned scratch for the
unwrapped encoding, so the array path passes the builder itself: the
temporary
`String` per wrapping row is gone, and so is the copy that staged each row
in a
second buffer. Encoding a batch now performs no per-row heap allocation.
Measured with the committed `base64` benchmark (criterion, 8192 rows, vs the
previous commit on this branch):
| Benchmark | Change |
| --- | --- |
| short, unchunked | ~0% |
| short, chunked | -1.7% |
| long, unchunked | -7.2% |
| long, chunked | -20.8% |
Also make the value-buffer capacity a true upper bound. Rows pad
independently,
so the encoded total is `sum ceil(len_i / 3) * 4`, which cannot be derived
from
the input's total byte count alone: N one-byte rows encode to `4N`, not to
`base64_encoded_len(N)`, so the old estimate under-reserved by ~3x for
arrays of
tiny values. Adding `4 * (N - 1)` makes it a genuine bound and stays O(1),
with
no pass over the offsets.
Tests:
- `chunked_matches_encode_then_split_at_every_length` sweeps every input
length
across three line boundaries, asserting the output equals encoding the
whole
input and splitting every 76 chars, that every non-final line is full, and
that no line exceeds the limit.
- `encoded_capacity_is_an_upper_bound` checks the estimate against the
actual
encoded total for eight array shapes, including the all-one-byte worst
case,
under both chunk settings.
- `encode_array_many_tiny_rows` covers the shape the old estimate
under-reserved
for.
---
native/spark-expr/Cargo.toml | 4 +
native/spark-expr/benches/base64.rs | 84 ++++++++++
native/spark-expr/src/string_funcs/base64.rs | 221 ++++++++++++++++++++++++---
3 files changed, 284 insertions(+), 25 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 775311935c..5d518387f5 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -95,6 +95,10 @@ harness = false
name = "padding"
harness = false
+[[bench]]
+name = "base64"
+harness = false
+
[[bench]]
name = "date_trunc"
harness = false
diff --git a/native/spark-expr/benches/base64.rs
b/native/spark-expr/benches/base64.rs
new file mode 100644
index 0000000000..d205e75ab2
--- /dev/null
+++ b/native/spark-expr/benches/base64.rs
@@ -0,0 +1,84 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::builder::BinaryBuilder;
+use arrow::array::ArrayRef;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::common::ScalarValue;
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_comet_spark_expr::spark_base64;
+use std::hint::black_box;
+use std::sync::Arc;
+
+/// Build a binary array with a mix of short values and longer values that
+/// exercise the CRLF line-wrapping path when chunking is enabled.
+fn create_binary_array(size: usize, value_len: usize) -> ArrayRef {
+ let mut builder = BinaryBuilder::new();
+ for i in 0..size {
+ if i % 10 == 0 {
+ builder.append_null();
+ } else {
+ let byte = (i % 251) as u8;
+ let value = vec![byte; value_len];
+ builder.append_value(&value);
+ }
+ }
+ Arc::new(builder.finish())
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let size = 8192;
+ // Short values (encode to <= 76 chars, no wrapping).
+ let short = create_binary_array(size, 16);
+ // Long values (encode to several wrapped lines when chunking).
+ let long = create_binary_array(size, 200);
+
+ c.bench_function("spark_base64: short, unchunked", |b| {
+ let args = vec![
+ ColumnarValue::Array(Arc::clone(&short)),
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))),
+ ];
+ b.iter(|| black_box(spark_base64(black_box(&args)).unwrap()))
+ });
+
+ c.bench_function("spark_base64: short, chunked", |b| {
+ let args = vec![
+ ColumnarValue::Array(Arc::clone(&short)),
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))),
+ ];
+ b.iter(|| black_box(spark_base64(black_box(&args)).unwrap()))
+ });
+
+ c.bench_function("spark_base64: long, unchunked", |b| {
+ let args = vec![
+ ColumnarValue::Array(Arc::clone(&long)),
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))),
+ ];
+ b.iter(|| black_box(spark_base64(black_box(&args)).unwrap()))
+ });
+
+ c.bench_function("spark_base64: long, chunked", |b| {
+ let args = vec![
+ ColumnarValue::Array(Arc::clone(&long)),
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))),
+ ];
+ b.iter(|| black_box(spark_base64(black_box(&args)).unwrap()))
+ });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/src/string_funcs/base64.rs
b/native/spark-expr/src/string_funcs/base64.rs
index 8e901e2a60..d394cfdccf 100644
--- a/native/spark-expr/src/string_funcs/base64.rs
+++ b/native/spark-expr/src/string_funcs/base64.rs
@@ -15,9 +15,12 @@
// specific language governing permissions and limitations
// under the License.
+use std::fmt::{self, Write};
use std::sync::Arc;
-use arrow::array::{Array, AsArray, GenericBinaryArray, OffsetSizeTrait,
StringArray};
+use arrow::array::{
+ Array, AsArray, GenericBinaryArray, GenericStringBuilder, OffsetSizeTrait,
StringArray,
+};
use arrow::datatypes::DataType;
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
@@ -62,48 +65,108 @@ 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;
+
+/// Panic message for the `fmt::Write` sinks used here; neither can actually
fail.
+const INFALLIBLE_SINK: &str = "writing base64 to an in-memory buffer cannot
fail";
+
+/// 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` as base64 into `out`, wrapping at `LINE_LEN` when `chunk`
is true.
+///
+/// `scratch` receives the unwrapped encoding from a single bulk
`encode_string` call. Both
+/// `scratch` and `out` are owned by the caller and reused across rows, so
encoding a batch performs
+/// no per-row heap allocation. Because `out` is a `fmt::Write` sink, the
array path can pass the
+/// output builder directly and have the wrapped result land in its value
buffer, rather than
+/// staging each row in a second buffer and copying it in.
+///
+/// An alternative is to skip `scratch` entirely and encode line-sized windows
straight into `out`:
+/// MIME wrapping is aligned to base64's block structure, so 57 input bytes
encode to exactly
+/// `LINE_LEN` chars with no padding, and only the final window can carry
padding. That is correct
+/// and allocation-free, but measurably slower — the per-call overhead of many
small `encode_string`
+/// calls exceeds the cost of one bulk encode plus the copy out. See the
benchmark discussion on
+/// <https://github.com/apache/datafusion-comet/pull/4885>.
+fn encode_into<W: Write>(
+ bytes: &[u8],
+ chunk: bool,
+ scratch: &mut String,
+ out: &mut W,
+) -> fmt::Result {
+ scratch.clear();
+ BASE64_STANDARD.encode_string(bytes, scratch);
+ if !chunk || scratch.len() <= LINE_LEN {
+ return out.write_str(scratch);
}
- let separators = (encoded.len() - 1) / LINE_LEN;
- let mut out = String::with_capacity(encoded.len() + separators * 2);
let mut offset = 0;
- while offset < encoded.len() {
+ while offset < scratch.len() {
if offset > 0 {
- out.push_str("\r\n");
+ out.write_str("\r\n")?;
}
- let end = (offset + LINE_LEN).min(encoded.len());
- out.push_str(&encoded[offset..end]);
+ let end = (offset + LINE_LEN).min(scratch.len());
+ out.write_str(&scratch[offset..end])?;
offset = end;
}
+ Ok(())
+}
+
+/// O(1) upper bound on the total encoded length of `array`, for sizing the
output value buffer.
+///
+/// Each row pads independently, so the encoded total is `sum ceil(len_i / 3)
* 4`, which cannot be
+/// derived from the input's total byte count alone: N rows of one byte each
encode to `4N`, not to
+/// `base64_encoded_len(N)`. Since `sum ceil(x_i) <= ceil(sum x_i) + (N - 1)`,
adding `4 * (N - 1)`
+/// to the whole-input encoding turns the estimate into a true upper bound
while staying O(1) — no
+/// pass over the offsets. It over-reserves by at most 4 bytes per row, which
matters only for
+/// arrays of very short values, where the buffer is small in absolute terms
anyway.
+fn encoded_capacity<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>, chunk:
bool) -> usize {
+ let encoded_total =
+ base64_encoded_len(array.value_data().len()) + 4 *
array.len().saturating_sub(1);
+ if chunk {
+ chunked_len(encoded_total)
+ } else {
+ encoded_total
+ }
+}
+
+fn encode_array<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>, chunk:
bool) -> StringArray {
+ let mut builder =
+ GenericStringBuilder::<i32>::with_capacity(array.len(),
encoded_capacity(array, chunk));
+ // Reused across rows, so a batch pays no per-row allocation.
+ let mut scratch = String::new();
+ for i in 0..array.len() {
+ if array.is_null(i) {
+ builder.append_null();
+ continue;
+ }
+ encode_into(array.value(i), chunk, &mut scratch, &mut
builder).expect(INFALLIBLE_SINK);
+ // Finalizes the value written through `fmt::Write` above.
+ builder.append_value("");
+ }
+ builder.finish()
+}
+
+fn encode(bytes: &[u8], chunk: bool) -> String {
+ let mut scratch = String::new();
+ let mut out = String::new();
+ encode_into(bytes, chunk, &mut scratch, &mut out).expect(INFALLIBLE_SINK);
out
}
#[cfg(test)]
mod tests {
use super::*;
+ use arrow::array::{BinaryArray, LargeBinaryArray};
#[test]
fn unchunked_is_a_single_line() {
@@ -137,4 +200,112 @@ mod tests {
vec![76, 76, 8]
);
}
+
+ #[test]
+ fn chunked_matches_encode_then_split_at_every_length() {
+ // The window encoding must be byte-for-byte what encoding the whole
input and splitting
+ // every LINE_LEN chars produces. Sweep every length across three
window boundaries so a
+ // padding or off-by-one error in the window size cannot hide.
+ for n in 0..=(57 * 3 + 4) {
+ let input = vec![b'z'; n];
+ let unwrapped = BASE64_STANDARD.encode(&input);
+ let expected = unwrapped
+ .as_bytes()
+ .chunks(LINE_LEN)
+ .map(|line| std::str::from_utf8(line).unwrap())
+ .collect::<Vec<_>>()
+ .join("\r\n");
+ assert_eq!(encode(&input, true), expected, "length {n}");
+ // Every line is full except the last, and no line exceeds the
limit.
+ let actual = encode(&input, true);
+ let lines: Vec<&str> = actual.split("\r\n").collect();
+ for (i, line) in lines.iter().enumerate() {
+ if i + 1 < lines.len() {
+ assert_eq!(line.len(), LINE_LEN, "length {n}, line {i}");
+ } else {
+ assert!(line.len() <= LINE_LEN, "length {n}, last line");
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn encoded_capacity_is_an_upper_bound() {
+ // Regression test for the capacity estimate: rows pad independently,
so an estimate
+ // derived from the summed input bytes alone under-reserves by ~3x for
many tiny rows.
+ let shapes: Vec<Vec<usize>> = vec![
+ vec![1; 64], // worst case: every row pads to 4
chars
+ vec![2; 64], //
+ vec![3; 64], // no padding at all
+ vec![0; 16], // all empty
+ vec![57; 8], // exactly one line each
+ vec![58; 8], // wraps once each
+ vec![200, 1, 0, 57, 58, 3], // mixed
+ vec![4096; 4], // few large rows
+ ];
+ for shape in shapes {
+ let values: Vec<Option<Vec<u8>>> = shape.iter().map(|&n|
Some(vec![b'q'; n])).collect();
+ let input = BinaryArray::from_iter(values);
+ for chunk in [false, true] {
+ let actual: usize = (0..input.len())
+ .map(|i| encode(input.value(i), chunk).len())
+ .sum();
+ let estimated = encoded_capacity(&input, chunk);
+ assert!(
+ estimated >= actual,
+ "shape {shape:?}, chunk={chunk}: estimate {estimated} <
actual {actual}"
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn encode_array_many_tiny_rows() {
+ // The shape that the old capacity estimate under-reserved for.
Verifies output
+ // correctness independently of the reservation.
+ let values: Vec<Option<&[u8]>> = (0..100).map(|_|
Some(&b"a"[..])).collect();
+ let input = BinaryArray::from(values);
+ for chunk in [false, true] {
+ let out = encode_array(&input, chunk);
+ assert_eq!(out.len(), 100);
+ for i in 0..out.len() {
+ assert_eq!(out.value(i), "YQ==");
+ }
+ }
+ }
+
+ #[test]
+ fn encode_array_binary_chunked() {
+ // Nulls, empty values, values that wrap, and one on the boundary.
+ let a57 = [b'a'; 57];
+ let b58 = [b'b'; 58];
+ let c120 = [b'c'; 120];
+ let input = BinaryArray::from(vec![
+ Some(&b""[..]),
+ None,
+ Some(&a57[..]), // exactly 76 encoded chars: no wrap
+ Some(&b58[..]), // wraps once
+ None,
+ Some(&c120[..]), // wraps twice
+ ]);
+ let out = encode_array(&input, true);
+ assert_eq!(out.len(), 6);
+ assert_eq!(out.value(0), "");
+ assert!(out.is_null(1));
+ assert_eq!(out.value(2), encode(&a57, true));
+ assert_eq!(out.value(3), encode(&b58, true));
+ assert!(out.is_null(4));
+ assert_eq!(out.value(5), encode(&c120, true));
+ }
+
+ #[test]
+ fn encode_array_large_binary_unchunked() {
+ // Exercises the LargeBinaryArray (i64 offsets) instantiation of the
generic.
+ let input = LargeBinaryArray::from(vec![None, Some(&b"abc"[..]),
Some(&b"hi!"[..])]);
+ let out = encode_array(&input, false);
+ assert_eq!(out.len(), 3);
+ assert!(out.is_null(0));
+ assert_eq!(out.value(1), "YWJj");
+ assert_eq!(out.value(2), BASE64_STANDARD.encode(b"hi!"));
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
