From 67dc9770e25393dbd1e3f21d4dd03090c89ea55c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 10 Jul 2026 08:07:44 -0600 Subject: [PATCH] perf: optimize spark_base64 in spark-expr --- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/base64.rs | 84 ++++++++++++++++++++ native/spark-expr/src/string_funcs/base64.rs | 54 +++++++++++-- 3 files changed, 134 insertions(+), 8 deletions(-) create mode 100644 native/spark-expr/benches/base64.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 800fe3ecb1..b53d83208e 100644 --- a/native/spark-expr/Cargo.toml +++ b/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 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..279c950b7f 100644 --- a/native/spark-expr/src/string_funcs/base64.rs +++ b/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 usize { + n.div_ceil(3) * 4 +} + fn encode_array(array: &GenericBinaryArray, 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::::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)]