diff --git a/native/Cargo.lock b/native/Cargo.lock index ed3642a484..a307afa70f 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2057,7 +2057,6 @@ dependencies = [ "datafusion-comet-common", "datafusion-comet-jni-bridge", "futures", - "hex", "jni 0.22.4", "num", "rand 0.10.2", diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index b0c45353cd..be5a8335a2 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -42,7 +42,6 @@ jni = "0.22.4" futures = { workspace = true } twox-hash = "2.1.2" rand = { workspace = true } -hex = "0.4.3" base64 = "0.22.1" [dev-dependencies] @@ -136,6 +135,10 @@ harness = false name = "parse_url" harness = false +[[bench]] +name = "cast_binary_to_string" +harness = false + [[bench]] name = "arrays_overlap" harness = false diff --git a/native/spark-expr/benches/cast_binary_to_string.rs b/native/spark-expr/benches/cast_binary_to_string.rs new file mode 100644 index 0000000000..8ae9b61761 --- /dev/null +++ b/native/spark-expr/benches/cast_binary_to_string.rs @@ -0,0 +1,115 @@ +// 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 arrow::datatypes::DataType; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::{spark_cast, BinaryOutputStyle, EvalMode, SparkCastOptions}; +use std::hint::black_box; +use std::sync::Arc; + +/// Build a Binary array of `size` rows, each row holding `width` pseudo-random bytes, with a +/// null every `null_every` rows (0 = no nulls). +fn build(size: usize, width: usize, null_every: usize) -> ArrayRef { + let mut builder = BinaryBuilder::with_capacity(size, size * width); + let mut state: u32 = 0x9e37_79b9; + for i in 0..size { + if null_every != 0 && i % null_every == 0 { + builder.append_null(); + } else { + let value: Vec = (0..width) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + (state >> 24) as u8 + }) + .collect(); + builder.append_value(&value); + } + } + Arc::new(builder.finish()) +} + +fn options(style: Option) -> SparkCastOptions { + let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + options.binary_output_style = style; + options +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 8192; + + let short = build(size, 8, 0); + let long = build(size, 64, 0); + let with_nulls = build(size, 16, 17); + + let mut bench = |name: &str, arr: &ArrayRef, style: Option| { + let cast_options = options(style); + c.bench_function(name, |b| { + b.iter(|| { + black_box( + spark_cast( + ColumnarValue::Array(Arc::clone(arr)), + black_box(&DataType::Utf8), + black_box(&cast_options), + ) + .unwrap(), + ) + }) + }); + }; + + bench( + "cast_binary_to_string: hex discrete, short values", + &short, + Some(BinaryOutputStyle::HexDiscrete), + ); + bench( + "cast_binary_to_string: hex discrete, long values", + &long, + Some(BinaryOutputStyle::HexDiscrete), + ); + bench( + "cast_binary_to_string: hex discrete, with nulls", + &with_nulls, + Some(BinaryOutputStyle::HexDiscrete), + ); + bench( + "cast_binary_to_string: hex, long values", + &long, + Some(BinaryOutputStyle::Hex), + ); + bench( + "cast_binary_to_string: basic, long values", + &long, + Some(BinaryOutputStyle::Basic), + ); + bench( + "cast_binary_to_string: base64, long values", + &long, + Some(BinaryOutputStyle::Base64), + ); + bench( + "cast_binary_to_string: default style, long values", + &long, + None, + ); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 8c2e95ab17..02ea142254 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -63,9 +63,8 @@ use datafusion::common::{internal_err, DataFusionError, Result as DataFusionResu use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ColumnarValue; use datafusion_comet_common::decode_utf8_spark_lossy; -use std::borrow::Cow; use std::{ - fmt::{Debug, Display, Formatter}, + fmt::{Debug, Display, Formatter, Write}, hash::Hash, sync::Arc, }; @@ -794,6 +793,60 @@ impl PhysicalExpr for Cast { } } +const UPPER_HEX_DIGITS: [u8; 16] = *b"0123456789ABCDEF"; + +/// Writes `byte` as two uppercase hex digits. +#[inline] +fn write_upper_hex(out: &mut W, byte: u8) -> std::fmt::Result { + out.write_char(UPPER_HEX_DIGITS[(byte >> 4) as usize] as char)?; + out.write_char(UPPER_HEX_DIGITS[(byte & 0x0f) as usize] as char) +} + +/// Writes `byte` reinterpreted as a signed decimal, as Spark does when printing a byte array. +#[inline] +fn write_i8(out: &mut W, byte: u8) -> std::fmt::Result { + // Widened to `i16` so that negating `i8::MIN` does not overflow. + let value = i8::from_ne_bytes([byte]) as i16; + let magnitude = if value < 0 { + out.write_char('-')?; + -value + } else { + value + }; + if magnitude >= 100 { + out.write_char((b'0' + (magnitude / 100) as u8) as char)?; + } + if magnitude >= 10 { + out.write_char((b'0' + (magnitude / 10 % 10) as u8) as char)?; + } + out.write_char((b'0' + (magnitude % 10) as u8) as char) +} + +/// Writes the bytes of `value` between square brackets, encoded with `encode` and joined by +/// `separator`. +fn write_bracketed( + out: &mut W, + value: &[u8], + separator: &str, + encode: fn(&mut W, u8) -> std::fmt::Result, +) -> std::fmt::Result { + out.write_char('[')?; + for (i, byte) in value.iter().enumerate() { + if i > 0 { + out.write_str(separator)?; + } + encode(out, *byte)?; + } + out.write_char(']') +} + +/// Casts a binary array to a string array. Without a binary output style the bytes are +/// reinterpreted as a string as-is, which is what Spark's `Cast` does. +/// +/// The other styles mimic the [BinaryFormatter]: https://github.com/apache/spark/blob/v4.0.0/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ToStringBase.scala#L449-L468 +/// used by SparkSQL's ToPrettyString expression. +/// The BinaryFormatter was [introduced]: https://issues.apache.org/jira/browse/SPARK-47911 in Spark 4.0.0 +/// Before Spark 4.0.0, the default is SPACE_DELIMITED_UPPERCASE_HEX fn cast_binary_to_string( array: &dyn Array, spark_cast_options: &SparkCastOptions, @@ -803,77 +856,59 @@ fn cast_binary_to_string( .downcast_ref::>>() .unwrap(); - // Build with a GenericStringBuilder and append &str straight from the decoder's Cow so the - // common valid-UTF-8 cast path copies the bytes once (into the Arrow value buffer) instead of - // twice (an owned String, then a copy into the buffer). - let mut builder = GenericStringBuilder::::new(); - for value in input.iter() { - match value { - Some(value) => match spark_cast_options.binary_output_style { - // ToPrettyString styles (Spark 4.0+) build owned strings, which is unavoidable. - Some(s) => builder.append_value(spark_binary_formatter(value, s)), - // Default CAST(binary AS string): borrows in the valid path, appends once. - None => builder.append_value(cast_binary_formatter(value)), - }, - None => builder.append_null(), - } - } - Ok(Arc::new(builder.finish())) -} + let num_rows = input.len(); + let offsets = input.value_offsets(); + let value_bytes = offsets[num_rows].as_usize() - offsets[0].as_usize(); + // Upper bound on the encoded length so the value buffer is allocated once. Base64 rounds up + // to a multiple of four per row, hence the extra `num_rows` slack. The default and UTF8 paths + // decode JVM-compatibly-lossily, which can expand each byte to a 3-byte U+FFFD replacement. + let capacity = match spark_cast_options.binary_output_style { + None | Some(BinaryOutputStyle::Utf8) => 3 * value_bytes, + Some(BinaryOutputStyle::Basic) => 6 * value_bytes + 2 * num_rows, + Some(BinaryOutputStyle::Base64) => 4 * value_bytes.div_ceil(3) + num_rows, + Some(BinaryOutputStyle::Hex) => 2 * value_bytes, + Some(BinaryOutputStyle::HexDiscrete) => 3 * value_bytes + 2 * num_rows, + }; -/// This function mimics the [BinaryFormatter]: https://github.com/apache/spark/blob/v4.0.0/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ToStringBase.scala#L449-L468 -/// used by SparkSQL's ToPrettyString expression. -/// The BinaryFormatter was [introduced]: https://issues.apache.org/jira/browse/SPARK-47911 in Spark 4.0.0 -/// Before Spark 4.0.0, the default is SPACE_DELIMITED_UPPERCASE_HEX -fn spark_binary_formatter(value: &[u8], binary_output_style: BinaryOutputStyle) -> String { - match binary_output_style { - // Spark's UTF8 BinaryFormatter renders via `UTF8String.fromBytes(bytes).toString`, i.e. - // `new String(bytes, UTF_8)`. Route through the shared JVM-compatible decoder so invalid - // bytes become U+FFFD (matching Spark) instead of panicking on non-UTF-8 input (#4488). - BinaryOutputStyle::Utf8 => decode_utf8_spark_lossy(value).into_owned(), - BinaryOutputStyle::Basic => { - format!( - "{:?}", - value - .iter() - .map(|v| i8::from_ne_bytes([*v])) - .collect::>() - ) - } - BinaryOutputStyle::Base64 => BASE64_STANDARD_NO_PAD.encode(value), - BinaryOutputStyle::Hex => value - .iter() - .map(|v| hex::encode_upper([*v])) - .collect::(), - BinaryOutputStyle::HexDiscrete => { + let mut builder = GenericStringBuilder::::with_capacity(num_rows, capacity); + // Base64 is the only style that cannot encode straight into the builder. + let mut base64_buffer = String::new(); + for value in input.iter() { + let Some(value) = value else { + builder.append_null(); + continue; + }; + // Encode directly into the builder's value buffer; `append_value("")` then terminates + // the row. Writing to the builder is infallible. + let written = match spark_cast_options.binary_output_style { + // Default CAST(binary AS string) and the UTF8 ToPrettyString style (Spark 4.0+) both + // render via `new String(bytes, UTF_8)`. Route through the shared JVM-compatible lossy + // decoder so ill-formed bytes become U+FFFD (matching Spark) instead of being + // reinterpreted unchecked (UB) or panicking on non-UTF-8 input (#4488, #4763). The + // valid-UTF-8 path borrows, so it copies once (into the Arrow value buffer). Divergence + // for byte-level round-trips such as CAST(CAST(x AS string) AS binary) and value + // identity is documented in the compatibility guide and tracked by #4764. + None | Some(BinaryOutputStyle::Utf8) => { + builder.write_str(&decode_utf8_spark_lossy(value)) + } + Some(BinaryOutputStyle::Basic) => write_bracketed(&mut builder, value, ", ", write_i8), + Some(BinaryOutputStyle::Base64) => { + base64_buffer.clear(); + BASE64_STANDARD_NO_PAD.encode_string(value, &mut base64_buffer); + builder.write_str(&base64_buffer) + } + Some(BinaryOutputStyle::Hex) => value + .iter() + .try_for_each(|byte| write_upper_hex(&mut builder, *byte)), // Spark's default SPACE_DELIMITED_UPPERCASE_HEX - format!( - "[{}]", - value - .iter() - .map(|v| hex::encode_upper([*v])) - .collect::>() - .join(" ") - ) - } + Some(BinaryOutputStyle::HexDiscrete) => { + write_bracketed(&mut builder, value, " ", write_upper_hex) + } + }; + written.expect("writing to a string builder cannot fail"); + builder.append_value(""); } -} - -fn cast_binary_formatter(value: &[u8]) -> Cow<'_, str> { - // CAST(binary AS string) reinterprets the bytes as UTF-8, like Spark's UTF8String.fromBytes. - // Spark keeps the raw bytes, but Arrow's Utf8 type requires valid UTF-8, and building a String - // from non-UTF-8 bytes is undefined behaviour (#4488). Decode JVM-compatibly-lossily instead: - // `decode_utf8_spark_lossy` replaces ill-formed sequences with U+FFFD exactly as Spark's - // `new String(bytes, UTF_8)` does (the same decoder Comet's native shuffle uses, #4521). The - // result is memory-safe valid UTF-8, never feeds invalid bytes into downstream native string - // kernels, and matches Spark's rendered output byte-for-byte. It diverges from Spark for - // operations that read the underlying bytes rather than the rendered text: byte-level round-trips - // such as CAST(CAST(x AS string) AS binary) (Spark still has the original bytes), and value - // identity, since every ill-formed sequence decodes to the same U+FFFD and so two strings that - // differ in Spark can compare equal in Comet. Both are documented in the compatibility guide and - // tracked by #4764. Returning `Cow` lets the valid path borrow so the caller appends without an - // intermediate allocation. - decode_utf8_spark_lossy(value) + Ok(Arc::new(builder.finish())) } #[cfg(test)] @@ -934,6 +969,94 @@ mod tests { assert_eq!(strings.value(3), "\u{FFFD}"); } + #[test] + fn test_cast_binary_to_string_styles() { + let input: ArrayRef = Arc::new(BinaryArray::from_opt_vec(vec![ + Some(b"\x00\x01\xfe".as_slice()), + Some(b"".as_slice()), + None, + Some(b"hi".as_slice()), + ])); + let cast = |style: Option| { + let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + options.binary_output_style = style; + let result = spark_cast( + ColumnarValue::Array(Arc::clone(&input)), + &DataType::Utf8, + &options, + ) + .unwrap() + .into_array(input.len()) + .unwrap(); + let result = result.as_string::(); + (0..result.len()) + .map(|i| (!result.is_null(i)).then(|| result.value(i).to_string())) + .collect::>() + }; + + assert_eq!( + cast(Some(BinaryOutputStyle::HexDiscrete)), + vec![ + Some("[00 01 FE]".to_string()), + Some("[]".to_string()), + None, + Some("[68 69]".to_string()), + ] + ); + assert_eq!( + cast(Some(BinaryOutputStyle::Hex)), + vec![ + Some("0001FE".to_string()), + Some("".to_string()), + None, + Some("6869".to_string()), + ] + ); + assert_eq!( + cast(Some(BinaryOutputStyle::Basic)), + vec![ + Some("[0, 1, -2]".to_string()), + Some("[]".to_string()), + None, + Some("[104, 105]".to_string()), + ] + ); + assert_eq!( + cast(Some(BinaryOutputStyle::Base64)), + vec![ + Some("AAH+".to_string()), + Some("".to_string()), + None, + Some("aGk".to_string()), + ] + ); + } + + #[test] + fn test_cast_binary_to_string_default_style_valid_utf8_through_spark_cast() { + // Exercises the default cast path through the public `spark_cast` entry point. Invalid bytes + // decode JVM-compatibly-lossily to U+FFFD (#4763), while valid multi-byte UTF-8 ("héllo") is + // preserved exactly. + let input: ArrayRef = Arc::new(BinaryArray::from_opt_vec(vec![ + Some(b"\xff\xfe".as_slice()), + None, + Some("héllo".as_bytes()), + ])); + let options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let result = spark_cast( + ColumnarValue::Array(Arc::clone(&input)), + &DataType::Utf8, + &options, + ) + .unwrap() + .into_array(input.len()) + .unwrap(); + let result = result.as_string::(); + assert_eq!(result.value(0), "\u{FFFD}\u{FFFD}"); + assert!(result.is_null(1)); + assert_eq!(result.value(2), "héllo"); + } + #[test] fn test_cast_unsupported_timestamp_to_date() { // Since datafusion uses chrono::Datetime internally not all dates representable by TimestampMicrosecondType are supported