perf: optimize spark_base64 in spark-expr

native/spark-expr/Cargo.toml

@@ -84,6 +84,10 @@ harness = false
 name = "padding"
 harness = false
 
+[[bench]]
+name = "base64"
+harness = false
+
 [[bench]]
 name = "date_trunc"
 harness = false

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);

native/spark-expr/src/string_funcs/base64.rs

@@ -17,7 +17,9 @@
 
 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,11 +64,41 @@ pub fn spark_base64(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionE
     }
 }
 
+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_array<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>, chunk: bool) -> StringArray {
-    array
-        .iter()
-        .map(|value| value.map(|bytes| encode(bytes, chunk)))
-        .collect()
+    // Encode into a builder, reusing scratch buffers across rows. This avoids a
+    // fresh per-row heap allocation for each element's base64 string (and, when
+    // chunking, its line-wrapped copy) that the previous per-element
+    // `encode()` + `collect()` incurred, and sizes the value buffer up front.
+    let data_capacity: usize = (0..array.len())
+        .filter(|&i| array.is_valid(i))
+        .map(|i| base64_encoded_len(array.value(i).len()))
+        .sum();
+    let mut builder = GenericStringBuilder::<i32>::with_capacity(array.len(), data_capacity);
+    let mut encoded = String::new();
+    let mut wrapped = String::new();
+    for i in 0..array.len() {
+        if array.is_null(i) {
+            builder.append_null();
+            continue;
+        }
+        encoded.clear();
+        BASE64_STANDARD.encode_string(array.value(i), &mut encoded);
+        if chunk && encoded.len() > LINE_LEN {
+            wrapped.clear();
+            chunk_into_lines_buf(&encoded, &mut wrapped);
+            builder.append_value(&wrapped);
+        } else {
+            builder.append_value(&encoded);
+        }
+    }
+    builder.finish()
 }
 
 fn encode(bytes: &[u8], chunk: bool) -> String {
@@ -83,12 +115,19 @@ fn encode(bytes: &[u8], chunk: bool) -> String {
 /// 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;
     }
+    let mut out = String::new();
+    chunk_into_lines_buf(&encoded, &mut out);
+    out
+}
+
+/// Append the CRLF-wrapped form of `encoded` (lines of at most 76 characters) into `out`.
+/// The caller is expected to have cleared `out`; only used when `encoded.len() > LINE_LEN`.
+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);
     let mut offset = 0;
     while offset < encoded.len() {
         if offset > 0 {
@@ -98,7 +137,6 @@ fn chunk_into_lines(encoded: String) -> String {
         out.push_str(&encoded[offset..end]);
         offset = end;
     }
-    out
 }
 
 #[cfg(test)]