You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/07/01 16:39:22 UTC
(datafusion-comet) branch main updated: feat: Add support for `base64` (#4778)
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 07614601e3 feat: Add support for `base64` (#4778)
07614601e3 is described below
commit 07614601e30be31d9889955c3ad257d510ce9bb4
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Jul 1 10:39:18 2026 -0600
feat: Add support for `base64` (#4778)
* feat: support base64 natively
* refactor: address base64 PR feedback
Use the AsArray trait (as_binary::<i32/i64>) instead of manual
downcast_ref/unwrap for the binary array cases, pass the encoded String
into chunk_into_lines by value so the common short-input case returns it
without a second allocation, and add a base64_unchunked.sql file test
(spark.sql.chunkBase64String.enabled=false) to cover the unwrapped path
with a real Spark-vs-Comet comparison on Spark 3.5+.
---
docs/source/user-guide/latest/expressions.md | 2 +-
native/spark-expr/src/comet_scalar_funcs.rs | 4 +
native/spark-expr/src/string_funcs/base64.rs | 140 +++++++++++++++++++++
native/spark-expr/src/string_funcs/mod.rs | 2 +
.../org/apache/comet/serde/QueryPlanSerde.scala | 1 +
.../scala/org/apache/comet/serde/statics.scala | 33 ++++-
.../scala/org/apache/comet/serde/strings.scala | 20 ++-
.../sql-tests/expressions/string/base64.sql | 40 ++++++
.../expressions/string/base64_unchunked.sql | 40 ++++++
9 files changed, 277 insertions(+), 5 deletions(-)
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
index ae9530d8ec..4da175db33 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -535,7 +535,7 @@ expression-level). The `outer` variants are wired but marked `Incompatible`; the
| Function | Status | Notes |
| --- | --- | --- |
| `ascii` | ✅ | |
-| `base64` | 🔜 | Lowers to `StaticInvoke(encode)` (not allowlisted); falls back |
+| `base64` | ✅ | |
| `bit_length` | ✅ | |
| `btrim` | ✅ | |
| `char` | ✅ | |
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs
index b8618a9dd2..42ee72c82a 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -193,6 +193,10 @@ pub fn create_comet_physical_fun_with_eval_mode(
let func = Arc::new(spark_log);
make_comet_scalar_udf!("spark_log", func, without data_type)
}
+ "base64" => {
+ let func = Arc::new(crate::string_funcs::spark_base64);
+ make_comet_scalar_udf!("base64", func, without data_type)
+ }
"split" => {
let func = Arc::new(crate::string_funcs::spark_split);
make_comet_scalar_udf!("split", func, without data_type)
diff --git a/native/spark-expr/src/string_funcs/base64.rs b/native/spark-expr/src/string_funcs/base64.rs
new file mode 100644
index 0000000000..8e901e2a60
--- /dev/null
+++ b/native/spark-expr/src/string_funcs/base64.rs
@@ -0,0 +1,140 @@
+// 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 std::sync::Arc;
+
+use arrow::array::{Array, AsArray, GenericBinaryArray, OffsetSizeTrait, StringArray};
+use arrow::datatypes::DataType;
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use datafusion::common::{exec_err, DataFusionError, ScalarValue};
+use datafusion::physical_plan::ColumnarValue;
+
+/// Spark `base64(bin)`: encodes a binary value as a padded base64 string.
+///
+/// The second argument is a boolean `chunk` flag mirroring Spark's
+/// `spark.sql.chunkBase64String.enabled`. When `chunk` is true (Spark's default, and the only
+/// behavior on Spark 3.4), the output matches `java.util.Base64.getMimeEncoder()`: lines of at most
+/// 76 characters joined by a CRLF (`\r\n`), with no trailing separator. When false, the output is a
+/// single unwrapped line, matching `java.util.Base64.getMimeEncoder(-1, [])`.
+pub fn spark_base64(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
+ if args.len() != 2 {
+ return exec_err!("base64 expects exactly two arguments, got {}", args.len());
+ }
+ let chunk = match &args[1] {
+ ColumnarValue::Scalar(ScalarValue::Boolean(Some(chunk))) => *chunk,
+ other => return exec_err!("base64 expects a boolean chunk flag, got {other:?}"),
+ };
+ match &args[0] {
+ ColumnarValue::Array(array) => match array.data_type() {
+ DataType::Binary => Ok(ColumnarValue::Array(Arc::new(encode_array(
+ array.as_binary::<i32>(),
+ chunk,
+ )))),
+ DataType::LargeBinary => Ok(ColumnarValue::Array(Arc::new(encode_array(
+ array.as_binary::<i64>(),
+ chunk,
+ )))),
+ other => exec_err!("base64 expects a binary argument, got {other}"),
+ },
+ ColumnarValue::Scalar(ScalarValue::Binary(value))
+ | ColumnarValue::Scalar(ScalarValue::LargeBinary(value)) => {
+ let encoded = value.as_ref().map(|bytes| encode(bytes, chunk));
+ Ok(ColumnarValue::Scalar(ScalarValue::Utf8(encoded)))
+ }
+ ColumnarValue::Scalar(other) => {
+ exec_err!("base64 expects a binary argument, got {other}")
+ }
+ }
+}
+
+fn encode_array<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>, chunk: bool) -> StringArray {
+ array
+ .iter()
+ .map(|value| value.map(|bytes| encode(bytes, chunk)))
+ .collect()
+}
+
+fn encode(bytes: &[u8], chunk: bool) -> String {
+ let encoded = BASE64_STANDARD.encode(bytes);
+ if chunk {
+ chunk_into_lines(encoded)
+ } else {
+ encoded
+ }
+}
+
+/// 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;
+ }
+ 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() {
+ if offset > 0 {
+ out.push_str("\r\n");
+ }
+ let end = (offset + LINE_LEN).min(encoded.len());
+ out.push_str(&encoded[offset..end]);
+ offset = end;
+ }
+ out
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn unchunked_is_a_single_line() {
+ assert_eq!(encode(b"abc", false), "YWJj");
+ assert_eq!(
+ encode(&[b'a'; 58], false),
+ BASE64_STANDARD.encode([b'a'; 58])
+ );
+ assert_eq!(encode(b"", false), "");
+ }
+
+ #[test]
+ fn chunked_matches_java_mime_encoder() {
+ // Empty and short inputs are returned without a separator.
+ assert_eq!(encode(b"", true), "");
+ assert_eq!(encode(b"abc", true), "YWJj");
+ // 57 input bytes encode to exactly 76 characters: the line limit, so no separator.
+ let exactly_76 = encode(&[b'a'; 57], true);
+ assert_eq!(exactly_76.len(), 76);
+ assert!(!exactly_76.contains("\r\n"));
+ // 58 input bytes encode to 80 characters, wrapping once after 76.
+ let wrapped_once = encode(&[b'b'; 58], true);
+ assert_eq!(wrapped_once.matches("\r\n").count(), 1);
+ assert_eq!(wrapped_once.split("\r\n").next().unwrap().len(), 76);
+ // 120 input bytes encode to 160 characters, wrapping twice (76 + 76 + 8).
+ let wrapped_twice = encode(&[b'c'; 120], true);
+ assert_eq!(wrapped_twice.matches("\r\n").count(), 2);
+ let lines: Vec<&str> = wrapped_twice.split("\r\n").collect();
+ assert_eq!(
+ lines.iter().map(|l| l.len()).collect::<Vec<_>>(),
+ vec![76, 76, 8]
+ );
+ }
+}
diff --git a/native/spark-expr/src/string_funcs/mod.rs b/native/spark-expr/src/string_funcs/mod.rs
index 44c90ec6c2..6e706cc89b 100644
--- a/native/spark-expr/src/string_funcs/mod.rs
+++ b/native/spark-expr/src/string_funcs/mod.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+mod base64;
mod contains;
mod get_json_object;
mod regexp_extract;
@@ -23,6 +24,7 @@ mod regexp_extract_common;
mod split;
mod substring;
+pub use base64::spark_base64;
pub use contains::SparkContains;
pub use get_json_object::spark_get_json_object;
pub use regexp_extract::spark_regexp_extract;
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 44142e75ed..2f7e5d6c17 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -251,6 +251,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
classOf[Overlay] -> CometOverlay,
classOf[SoundEx] -> CometSoundEx,
classOf[StringLocate] -> CometStringLocate,
+ classOf[Base64] -> CometBase64,
classOf[UnBase64] -> CometUnBase64,
classOf[ToCharacter] -> CometToCharacter,
classOf[ToNumber] -> CometToNumber,
diff --git a/spark/src/main/scala/org/apache/comet/serde/statics.scala b/spark/src/main/scala/org/apache/comet/serde/statics.scala
index 1012b01c10..f415c373a1 100644
--- a/spark/src/main/scala/org/apache/comet/serde/statics.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/statics.scala
@@ -19,12 +19,13 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.{Attribute, ExpressionImplUtils, Literal, StringDecode, TryEval, UrlCodec}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64, ExpressionImplUtils, Literal, StringDecode, TryEval, UrlCodec}
import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
import org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils
+import org.apache.spark.sql.types.StringType
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
-import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithFallbackReason, scalarFunctionExprToProto}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithFallbackReason, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType}
object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] {
@@ -44,7 +45,11 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] {
// Spark 4.0 lowers `decode(bin, charset)` to `StaticInvoke(StringDecode.decode, ...)`
// carrying the `legacyCharsets` / `legacyErrorAction` flags. Routing through the codegen
// dispatcher runs Spark's own decoder so both flags are honored. See #4465.
- ("decode", classOf[StringDecode]) -> CometStaticInvokeCodegenDispatch)
+ ("decode", classOf[StringDecode]) -> CometStaticInvokeCodegenDispatch,
+ // Spark 3.5+ makes `Base64` RuntimeReplaceable, lowering `base64(bin)` to
+ // `StaticInvoke(Base64.encode, Seq(child, chunkBase64), ...)`. On Spark 3.4 the `Base64`
+ // node survives and is handled directly (see CometBase64).
+ ("encode", classOf[Base64]) -> CometBase64StaticInvoke)
override def convert(
expr: StaticInvoke,
@@ -90,6 +95,28 @@ object CometUrlDecodeStaticInvoke extends CometExpressionSerde[StaticInvoke] {
}
}
+/**
+ * Handles `base64(bin)` on Spark 3.5+, where it lowers to `StaticInvoke(Base64.encode, Seq(child,
+ * chunkBase64))`. The `chunkBase64` literal carries `spark.sql.chunkBase64String.enabled`
+ * (default true) and is passed through to the native function, which honors both modes.
+ */
+object CometBase64StaticInvoke extends CometExpressionSerde[StaticInvoke] {
+ override def convert(
+ expr: StaticInvoke,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[ExprOuterClass.Expr] = {
+ val childExpr = exprToProtoInternal(expr.arguments.head, inputs, binding)
+ val chunkExpr = exprToProtoInternal(expr.arguments(1), inputs, binding)
+ val optExpr = scalarFunctionExprToProtoWithReturnType(
+ "base64",
+ StringType,
+ failOnError = false,
+ childExpr,
+ chunkExpr)
+ optExprWithFallbackReason(optExpr, expr, expr.arguments: _*)
+ }
+}
+
/** Routes a [[StaticInvoke]] through the JVM codegen dispatcher; used for AES. */
object CometStaticInvokeCodegenDispatch extends CometCodegenDispatch[StaticInvoke]
diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala
index cd0c5db1b1..3186818c9c 100644
--- a/spark/src/main/scala/org/apache/comet/serde/strings.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala
@@ -19,7 +19,7 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.{Attribute, BitLength, Cast, Concat, ConcatWs, Elt, Expression, FindInSet, FormatNumber, FormatString, GetJsonObject, If, InitCap, IsNull, Left, Length, Levenshtein, Like, Literal, Lower, Mask, OctetLength, Overlay, RegExpExtract, RegExpExtractAll, RegExpInStr, RegExpReplace, Right, RLike, SoundEx, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringSplit, StringTranslate, Substring, SubstringIndex, ToCharacter, ToNumb [...]
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64, BitLength, Cast, Concat, ConcatWs, Elt, Expression, FindInSet, FormatNumber, FormatString, GetJsonObject, If, InitCap, IsNull, Left, Length, Levenshtein, Like, Literal, Lower, Mask, OctetLength, Overlay, RegExpExtract, RegExpExtractAll, RegExpInStr, RegExpReplace, Right, RLike, SoundEx, StringLocate, StringLPad, StringRepeat, StringReplace, StringRPad, StringSplit, StringTranslate, Substring, SubstringIndex, ToCharacter [...]
import org.apache.spark.sql.types.{BinaryType, DataTypes, LongType, StringType}
import org.apache.spark.unsafe.types.UTF8String
@@ -734,6 +734,24 @@ object CometSoundEx extends CometCodegenDispatch[SoundEx]
object CometStringLocate extends CometCodegenDispatch[StringLocate]
+// On Spark 3.4 `Base64` is a plain expression node and always chunks the output (it uses
+// `java.util.Base64.getMimeEncoder()` with no arguments). On Spark 3.5+ it is RuntimeReplaceable
+// and lowers to a `StaticInvoke`, handled by CometBase64StaticInvoke instead.
+object CometBase64 extends CometExpressionSerde[Base64] {
+ override def convert(expr: Base64, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = {
+ val childExpr = exprToProtoInternal(expr.child, inputs, binding)
+ val chunkExpr = exprToProtoInternal(Literal(true), inputs, binding)
+ val optExpr =
+ scalarFunctionExprToProtoWithReturnType(
+ "base64",
+ StringType,
+ failOnError = false,
+ childExpr,
+ chunkExpr)
+ optExprWithFallbackReason(optExpr, expr, expr.child)
+ }
+}
+
object CometUnBase64 extends CometCodegenDispatch[UnBase64]
object CometToCharacter extends CometCodegenDispatch[ToCharacter]
diff --git a/spark/src/test/resources/sql-tests/expressions/string/base64.sql b/spark/src/test/resources/sql-tests/expressions/string/base64.sql
new file mode 100644
index 0000000000..2acd10b0aa
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/string/base64.sql
@@ -0,0 +1,40 @@
+-- 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.
+
+statement
+CREATE TABLE test_base64(s string) USING parquet
+
+-- Inputs chosen to exercise Spark's default chunked encoding (spark.sql.chunkBase64String.enabled):
+-- 'abc'/'hello' stay under one line; the 57-char value encodes to exactly 76 base64 characters (the
+-- line limit, so no separator); the 58-char value wraps once; the 120-char value wraps twice.
+statement
+INSERT INTO test_base64 VALUES
+ ('abc'), ('hello'), (''), (NULL),
+ ('aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'),
+ ('bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb'),
+ ('cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc')
+
+query
+SELECT s, base64(cast(s AS binary)) FROM test_base64
+
+-- literal arguments
+query
+SELECT base64(cast('Spark SQL' AS binary)), base64(cast('' AS binary))
+
+-- round trip through unbase64
+query
+SELECT s, cast(unbase64(base64(cast(s AS binary))) AS string) FROM test_base64
diff --git a/spark/src/test/resources/sql-tests/expressions/string/base64_unchunked.sql b/spark/src/test/resources/sql-tests/expressions/string/base64_unchunked.sql
new file mode 100644
index 0000000000..79a252bb86
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/string/base64_unchunked.sql
@@ -0,0 +1,40 @@
+-- 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.
+
+-- Exercises the unchunked base64 path (chunk = false in the native function).
+-- On Spark 3.5+ this config produces single-line output with no CRLF separators.
+-- On Spark 3.4 the config does not exist and is ignored, so both engines still chunk;
+-- the file is harmless there and Comet still matches Spark.
+-- Config: spark.sql.chunkBase64String.enabled=false
+
+statement
+CREATE TABLE test_base64_unchunked(s string) USING parquet
+
+-- The 58-char and 120-char values encode to more than 76 base64 characters, so they would
+-- wrap under the default chunked encoding but stay on a single line here.
+statement
+INSERT INTO test_base64_unchunked VALUES
+ ('abc'), (''), (NULL),
+ ('bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb'),
+ ('cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc')
+
+query
+SELECT s, base64(cast(s AS binary)) FROM test_base64_unchunked
+
+-- round trip through unbase64
+query
+SELECT s, cast(unbase64(base64(cast(s AS binary))) AS string) FROM test_base64_unchunked
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org