From 269cedaa61543a1a3c9153ee1a6b73bc0d672e4e Mon Sep 17 00:00:00 2001 From: Manu Zhang Date: Wed, 28 Oct 2026 15:06:50 +0000 Subject: [PATCH] fix: decode invalid UTF-8 in nested Parquet BINARY values read as strings (#7106) * fix: decode invalid UTF-8 in nested Parquet BINARY values read as strings When a Parquet BINARY column without the UTF8 annotation is read as a string, through a read schema or spark.sql.parquet.binaryAsString, the native scan converts nested values (struct fields, list elements, map keys and values) in parquet_convert_array. That fell through to Arrow's safe cast, which turns invalid UTF-8 into NULL, so those values silently became NULL where Spark keeps the bytes. Convert Binary to Utf8 with Spark's cast instead. It replaces each ill-formed sequence with U+FFFD, the decode CAST(binary AS string) and native shuffle already share (#4764). Top-level columns are unchanged: DataFusion's file-schema coercion still decodes them straight into Utf8 (#7032). To keep the zero-copy conversion Arrow's cast gave valid nested values, cast_binary_to_string checks the bytes the rows span once for the default and UTF8 styles. It doesn't check the whole value buffer, which a slice shares with its parent. When the bytes are valid, it returns the input's value buffer narrowed to them, with a slice's offsets rebased. Otherwise rows inside the valid prefix are copied without being checked again, and only the rest go through the U+FFFD decoder. This also speeds up plain CAST(binary AS string). Split out of #6246. Co-Authored-By: Claude Opus 5.5 * fix: skip the UTF-8 fast path when nulls hide most of the payload cast_binary_to_string validated every byte the rows span, including the payload of null rows, which the row-by-row decoder skipped. After nullif or a CASE with column branches, a mostly-null batch can keep most of its bytes under nulls, so the cast scanned all of them and the reused buffer kept them alive downstream. Count the bytes of the non-null rows first. When null rows span more than half of the bytes, skip the check and decode the live rows one at a time as before, into a value buffer sized by the live bytes instead of the whole span. Co-Authored-By: Claude Opus 5.5 * fix: decode invalid UTF-8 in nested LargeBinary and BinaryView values The Parquet reader keeps the types a file's ARROW:schema stores, and DataFusion's file-schema coercion only changes top-level columns. So when an Arrow writer stores a list element, struct field or map value as LargeBinary or BinaryView, the native scan hands it to parquet_convert_array as that type. Only the Binary arm used Spark's cast, so these values still went through Arrow's safe cast and a malformed value became NULL. Convert LargeBinary and BinaryView to Binary first and use the same Spark cast. For LargeBinary that rewrites the offsets and keeps the value buffer. Co-Authored-By: Claude Opus 5.5 * fix: copy a valid slice that fills little of its value buffer A slice shares its parent's value buffer, and Buffer::slice_with_length narrows the view but keeps the whole allocation. When cast_binary_to_string reused a valid slice's bytes, every output batch cut from one large array, such as aggregate output, held that array's whole allocation. DataFusion's sort charges each batch the capacity of its buffers, so it counted the allocation once per batch and spilled where a compact copy didn't. Reuse the buffer only when the rows' bytes fill at least half of its allocation. Otherwise copy those already validated bytes into a compact buffer with fresh offsets, which still skips the row-by-row decoder. Co-Authored-By: Claude Opus 5.5 * refactor: share offsets when casting valid binary to string cast_binary_to_string rebuilt the offsets whenever it copied the values of valid input, even for an unsliced array whose offsets needed no rebasing. Decide each buffer on its own: reuse the offsets when they start at zero and fill at least half of their allocation, and the values when the rows' bytes fill at least half of theirs. An unsliced array keeps its offsets when only its value buffer is oversized, while the first slice of a large array still gets offsets that don't hold the parent's allocation. Also update the function's doc comment, which still said the bytes are taken as is, and pass parquet_convert_array_impl's Binary to Utf8 cast only the options it reads. Add rows to the nested ParquetReadSuite fixture: a null list element, struct field and map value ahead of the invalid row, and a row whose list, struct and map are null. Co-Authored-By: Claude Opus 5.5 * fix: decode invalid UTF-8 in nested dictionary-encoded BINARY values A file's ARROW:schema can dictionary-encode a nested BINARY leaf, which the reader then returns as Dictionary(_, Binary | LargeBinary). Those fell through to Arrow's safe cast and turned malformed values into NULL. Decode the dictionary values through the Binary arm and expand the keys. Co-Authored-By: Claude Opus 5.5 * fix: copy the validity bitmap a small slice shares with its parent The fast path for valid input compacted the offsets and values of a slice that fills little of their allocations, but cloned the input's validity bitmap, which still held the parent's whole allocation. A sort that charges every batch for what it retains then spilled where the base did not, for example on nullable binary group keys emitted in small slices. Apply the same rule to the bitmap: reuse it when the slice's bits fill at least half of its allocation, and copy them otherwise. Co-Authored-By: Claude Opus 5.5 * fix: keep the bytes of a map key read from unannotated BINARY Decoding turned the keys FF and EF BF BD into the same U+FFFD, so a lookup of one returned the other's value and map_from_entries failed with DUPLICATED_MAP_KEY where Spark, which tells keys apart by their bytes, succeeds. A key now keeps its bytes at any depth under a map's key field, unchecked, as a top-level value does. Map values still decode. Decoding a dictionary now touches only the entries the batch references, so a wide malformed entry that one batch uses is not decoded again for every other batch of the row group. A batch whose decoded values would not fit a string column is reported with the byte count instead of an offset overflow. In the binary-to-string fallback, count the live bytes over runs of non-null rows instead of row by row, so sparse nulls don't cost a pass over the batch. In the row that holds the first ill-formed sequence, copy its validated prefix and decode only the rest. Size the builder for the validated prefix plus three bytes per remaining byte, and copy the result compact when a replacement still left it filling under half of its allocation. Co-Authored-By: Claude Fable 5.1 * fix: copy the bytes of a valid binary cast outside a scan A plain CAST(binary AS string) can run on a batch imported over the Arrow C stream, where every buffer keeps the whole batch alive, wide sibling columns included, for as long as a sort or a join holds the result. Share the input's buffers only when a scan adapts its schema, where the reader owns them, and copy the validated bytes once otherwise, as the cast did before. Also cover a Dictionary(Int32, Binary) map key, link the tracking issue for the per-batch dictionary overflow from the code, and list that limit under the known issues with the batch size and the growth of ill-formed bytes. Co-Authored-By: Claude Fable 5.1 --------- Co-authored-by: Claude Opus 5.5 (cherry picked from commit 621b9ccdb4658ad2671e4e88e5c51da67e3c8fbc) --- .../user-guide/latest/compatibility/index.md | 11 +- native/core/src/parquet/parquet_exec.rs | 82 +++ native/core/src/parquet/parquet_support.rs | 305 +++++++++- .../spark-expr/src/conversion_funcs/cast.rs | 539 +++++++++++++++++- .../comet/parquet/ParquetReadSuite.scala | 90 +++ 5 files changed, 1002 insertions(+), 25 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 36cd8cf2b86..87aea310a18 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -77,7 +77,8 @@ codegen dispatcher covers only scalar expressions; see the Spark's `StringType` can hold arbitrary bytes, including sequences that are not valid UTF-8 (for example `CAST(X'FF' AS STRING)`). Arrow's string type requires valid UTF-8, so Comet cannot store the raw bytes natively. When Comet produces a string from arbitrary bytes (such as -`CAST(binary AS string)` or a columnar shuffle), it decodes them the same way the JVM does +`CAST(binary AS string)`, a columnar shuffle, or a nested Parquet BINARY value read as a string), +it decodes them the same way the JVM does (`new String(bytes, UTF_8)`), replacing each ill-formed sequence with the Unicode replacement character `U+FFFD`. Spark itself applies the identical replacement whenever such a string is materialized (collected, printed, or passed to most string functions), so the rendered result @@ -96,8 +97,12 @@ underlying bytes rather than on the rendered text: such as `contains` can therefore disagree with Spark when non-UTF-8 bytes are involved. Both differences require string data that is not valid UTF-8, which does not occur for text read from -Parquet or produced by string expressions. Consistent handling of invalid UTF-8 across all native -string paths is tracked by +UTF8-annotated Parquet columns or produced by string expressions. A BINARY column without the UTF8 +annotation can hold such bytes when it is read as `string`, through the read schema or +`spark.sql.parquet.binaryAsString=true`. Comet decodes such a value in a struct field, list element +or map value. Top-level values and map keys are not checked +([#7032](https://github.com/apache/datafusion-comet/issues/7032)). Consistent handling of invalid +UTF-8 across all native string paths is tracked by [#4764](https://github.com/apache/datafusion-comet/issues/4764). Separately, Comet's native Parquet scan currently rejects string columns whose stored bytes are not diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 1ff3b3f7289..be5a01db194 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -535,6 +535,88 @@ mod tests { } } + /// The reader keeps the nested types of a file's ARROW:schema, so an Arrow writer's + /// `LargeBinary`, `BinaryView` or dictionary-encoded list elements reach the conversion to + /// Spark's string type as those types, not as `Binary`. A malformed value must still decode + /// to U+FFFD. + #[tokio::test] + async fn nested_arrow_binary_types_read_as_string_decode_invalid_utf8() { + use arrow::array::{Array, AsArray, BinaryArray, ListArray}; + use arrow::buffer::OffsetBuffer; + + for binary_type in [ + DataType::LargeBinary, + DataType::BinaryView, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)), + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::LargeBinary)), + ] { + let element = Arc::new(Field::new("element", binary_type.clone(), true)); + let values = arrow::compute::cast( + &BinaryArray::from(vec![b"one".as_slice(), &[0xff]]), + &binary_type, + ) + .unwrap(); + let list = ListArray::new(element, OffsetBuffer::new(vec![0, 2].into()), values, None); + let file_schema = Arc::new(Schema::new(vec![Field::new( + "s", + list.data_type().clone(), + true, + )])); + let batch = + RecordBatch::try_new(Arc::clone(&file_schema), vec![Arc::new(list)]).unwrap(); + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, file_schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let string_list = DataType::List(Arc::new(Field::new("element", DataType::Utf8, true))); + let required_schema = Arc::new(Schema::new(vec![Field::new("s", string_list, true)])); + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + required_schema, + None, + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![PartitionedFile::from_path(filename).unwrap()]], + None, + None, + None, + "UTC", + true, + false, + false, + false, + &session_ctx, + false, + false, + false, + ) + .unwrap(); + + let mut strings = Vec::new(); + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + let values = batch.column(0).as_list::().values(); + let values = values.as_string::(); + strings.extend(values.iter().map(|v| v.map(str::to_string))); + } + assert_eq!( + strings, + vec![Some("one".to_string()), Some("\u{fffd}".to_string())], + "{binary_type}" + ); + } + } + // Regression test for #4990: a fresh `TableParquetOptions::new()` ignored session-level // `datafusion.execution.parquet.*` settings entirely, so `spark.comet.datafusion. // execution.parquet.*` (behind `respectDataFusionConfigs`) and `spark.comet.parquet. diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index e8987a79da6..d284a788b5d 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -18,10 +18,11 @@ use crate::execution::operators::ExecutionError; use crate::parquet::name_fold::fold_names; use arrow::array::{ - make_array, FixedSizeBinaryArray, GenericListViewArray, MapArray, OffsetSizeTrait, StringArray, + make_array, BooleanBufferBuilder, FixedSizeBinaryArray, GenericListViewArray, MapArray, + OffsetSizeTrait, StringArray, }; use arrow::buffer::NullBuffer; -use arrow::compute::can_cast_types; +use arrow::compute::{can_cast_types, take}; use arrow::datatypes::{FieldRef, Fields}; use arrow::{ array::{ @@ -40,7 +41,7 @@ use datafusion::execution::object_store::{ use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::ColumnarValue; use datafusion_comet_common::SparkError; -use datafusion_comet_spark_expr::EvalMode; +use datafusion_comet_spark_expr::{spark_cast, EvalMode, SparkCastOptions}; use log::debug; use object_store::path::Path; use object_store::{parse_url, ObjectStore, ObjectStoreScheme}; @@ -187,14 +188,17 @@ fn parquet_convert_array( to_type: &DataType, parquet_options: &SparkParquetOptions, ) -> DataFusionResult { - parquet_convert_array_impl(array, to_type, parquet_options, None) + parquet_convert_array_impl(array, to_type, parquet_options, None, false) } +// `in_map_key` is set below a map's key field, where a BINARY value read as a string keeps its +// bytes instead of being decoded; see the `Binary` arm. fn parquet_convert_array_impl( array: ArrayRef, to_type: &DataType, parquet_options: &SparkParquetOptions, parent_nulls: Option<&NullBuffer>, + in_map_key: bool, ) -> DataFusionResult { use DataType::*; let from_type = array.data_type(); @@ -217,6 +221,7 @@ fn parquet_convert_array_impl( to_type, parquet_options, visible.as_ref(), + in_map_key, )?), ( List(_) | LargeList(_) | FixedSizeList(_, _) | ListView(_) | LargeListView(_), @@ -234,6 +239,7 @@ fn parquet_convert_array_impl( to_inner_type.data_type(), parquet_options, child_visibility.as_ref(), + in_map_key, )?; // Resolve element fields with Spark's rules before Arrow changes list layout. // Casting the original list directly can match missing struct fields by position. @@ -288,7 +294,7 @@ fn parquet_convert_array_impl( )) } (Map(_, ordered_from), Map(_, ordered_to)) if ordered_from == ordered_to => - parquet_convert_map_to_map(array.as_map(), to_type, parquet_options, *ordered_to, visible.as_ref(), checked_timestamp_overflow) + parquet_convert_map_to_map(array.as_map(), to_type, parquet_options, *ordered_to, visible.as_ref(), checked_timestamp_overflow, in_map_key) , // Iceberg stores UUIDs as 16-byte fixed binary but Spark expects string representation. // Arrow doesn't support casting FixedSizeBinary to Utf8, so we handle it manually. @@ -312,6 +318,60 @@ fn parquet_convert_array_impl( Ok(Arc::new(string_array)) } + // Unannotated BINARY read as a string. Spark's cast replaces malformed UTF-8 with + // U+FFFD, where Arrow's safe cast below would return NULL. Top-level columns don't + // reach this arm (#7032). The reader keeps a file's ARROW:schema types, so an Arrow + // writer's `LargeBinary` or `BinaryView` arrives as is; Spark's cast takes only + // `Binary`, and converting `LargeBinary` reuses its value buffer. + (Binary | LargeBinary | BinaryView, Utf8) => { + let array = match from_type { + Binary => array, + _ => cast_with_options(&array, &Binary, &PARQUET_OPTIONS)?, + }; + if in_map_key { + // Spark tells map keys apart by their bytes. Decoding turns the keys `FF` and + // `EF BF BD` into the same U+FFFD, so a lookup of one returns the other's value + // and `map_from_entries` fails with DUPLICATED_MAP_KEY. A key keeps its bytes + // instead, unchecked, as a top-level value does (#7032). + let (offsets, values, nulls) = array.as_binary::().clone().into_parts(); + let keys = unsafe { StringArray::new_unchecked(offsets, values, nulls) }; + return Ok(Arc::new(keys)); + } + // Binary to Utf8 reads only the binary output style, which a scan leaves unset, and + // whether the reader owns the input's buffers, which lets valid values share them. + let mut cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + cast_options.is_adapting_schema = true; + let len = array.len(); + spark_cast(ColumnarValue::Array(array), to_type, &cast_options)?.into_array(len) + } + // The same schema can dictionary-encode such a value. Decode each dictionary value once + // through the arm above, then expand the keys, as Arrow's cast would. Expanding flattens + // the batch into one string array with 32-bit offsets, so a batch whose rows reference + // wide entries many times can decode past `i32::MAX` bytes and fail where Spark, which + // keeps the dictionary, succeeds (#7254). + (Dictionary(_, value_type), Utf8) if matches!(value_type.as_ref(), Binary | LargeBinary) => { + let dictionary = array.as_any_dictionary(); + let keys = dictionary.normalized_keys(); + let key_nulls = dictionary.keys().nulls(); + let values = parquet_convert_array_impl( + referenced_dictionary_values(dictionary.values(), &keys, key_nulls)?, + to_type, + parquet_options, + None, + in_map_key, + )?; + let decoded_bytes = + decoded_bytes(values.as_string::().value_offsets(), &keys, key_nulls); + if decoded_bytes > i32::MAX as usize { + return Err(DataFusionError::Execution(format!( + "Decoding the dictionary-encoded BINARY values of one batch as strings \ + takes {decoded_bytes} bytes, more than the {} a string column holds. \ + Lower spark.comet.batchSize so that a batch holds fewer rows.", + i32::MAX + ))); + } + Ok(take(values.as_ref(), dictionary.keys(), None)?) + } // If Arrow cast supports the cast, delegate the cast to Arrow _ if can_cast_types(from_type, to_type) => { Ok(cast_with_options(&array, to_type, &PARQUET_OPTIONS)?) @@ -522,6 +582,7 @@ fn parquet_convert_struct_to_struct( to_type: &DataType, parquet_options: &SparkParquetOptions, parent_nulls: Option<&NullBuffer>, + in_map_key: bool, ) -> DataFusionResult { match (from_type, to_type) { (DataType::Struct(from_fields), DataType::Struct(to_fields)) => { @@ -536,6 +597,7 @@ fn parquet_convert_struct_to_struct( to_field.data_type(), parquet_options, parent_nulls, + in_map_key, )?); field_overlap = true; } else { @@ -575,6 +637,7 @@ fn parquet_convert_map_to_map( to_ordered: bool, parent_nulls: Option<&NullBuffer>, checked_timestamp_overflow: bool, + in_map_key: bool, ) -> Result { match to_data_type { DataType::Map(entries_field, _) => { @@ -595,12 +658,14 @@ fn parquet_convert_map_to_map( key_field.data_type(), parquet_options, child_visibility.as_ref(), + true, )?; let value_array = parquet_convert_array_impl( Arc::clone(from.values()), value_field.data_type(), parquet_options, child_visibility.as_ref(), + in_map_key, )?; Ok(Arc::new(MapArray::try_new( @@ -621,6 +686,47 @@ fn parquet_convert_map_to_map( } } +/// Hides the dictionary entries that no key of the batch references behind a null, so that +/// decoding the values touches only the entries the batch uses. A row group's dictionary backs +/// every batch read from the row group, so a wide ill-formed entry that one batch uses would +/// otherwise be decoded again for each of the others. A dictionary of up to 64 KiB is returned as +/// is: decoding it costs less than finding the referenced entries. +fn referenced_dictionary_values( + values: &ArrayRef, + keys: &[usize], + key_nulls: Option<&NullBuffer>, +) -> DataFusionResult { + const SMALL_DICTIONARY_BYTES: usize = 64 * 1024; + if values.get_buffer_memory_size() <= SMALL_DICTIONARY_BYTES { + return Ok(Arc::clone(values)); + } + let mut referenced = BooleanBufferBuilder::new(values.len()); + referenced.append_n(values.len(), false); + for (row, key) in keys.iter().enumerate() { + if key_nulls.is_none_or(|nulls| nulls.is_valid(row)) { + referenced.set_bit(*key, true); + } + } + let referenced = referenced.finish(); + if referenced.count_set_bits() == values.len() { + return Ok(Arc::clone(values)); + } + let nulls = NullBuffer::union(values.nulls(), Some(&NullBuffer::new(referenced))); + Ok(make_array( + values.to_data().into_builder().nulls(nulls).build()?, + )) +} + +/// The bytes that expanding the dictionary `values` (given by their `offsets`) through `keys` +/// writes, counting the rows that `key_nulls` keeps. +fn decoded_bytes(offsets: &[i32], keys: &[usize], key_nulls: Option<&NullBuffer>) -> usize { + keys.iter() + .enumerate() + .filter(|(row, _)| key_nulls.is_none_or(|nulls| nulls.is_valid(*row))) + .map(|(_, key)| (offsets[key + 1] - offsets[*key]) as usize) + .sum() +} + /// Gets the key field from the entries of a map. For all other types returns None. fn key_field(entries_field: &FieldRef) -> Option { if let DataType::Struct(fields) = entries_field.data_type() { @@ -976,6 +1082,195 @@ pub(crate) fn prepare_object_store_with_configs( #[cfg(test)] mod tests { + #[test] + fn nested_binary_to_string_decodes_invalid_utf8_like_spark() { + use super::{parquet_convert_array, SparkParquetOptions}; + use arrow::array::{ArrayRef, AsArray, BinaryArray, ListArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field}; + use datafusion_comet_spark_expr::EvalMode; + use std::sync::Arc; + + let values = BinaryArray::from(vec![Some(b"one".as_slice()), Some(&[0xff])]); + let input: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("element", DataType::Binary, true)), + OffsetBuffer::new(vec![0, 2].into()), + Arc::new(values), + None, + )); + let to_type = DataType::List(Arc::new(Field::new("element", DataType::Utf8, true))); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + + let result = parquet_convert_array(input, &to_type, &options).unwrap(); + let strings = result.as_list::().values().as_string::(); + assert_eq!(strings.value(0), "one"); + assert_eq!(strings.value(1), "\u{fffd}"); + } + + /// A map key keeps its bytes, also inside a struct key, so that keys that differ only in + /// ill-formed bytes stay distinct, as in Spark. The map's values are decoded. + #[test] + fn nested_binary_to_string_keeps_map_key_bytes() { + use super::{parquet_convert_array, SparkParquetOptions}; + use arrow::array::{ + ArrayRef, AsArray, BinaryArray, DictionaryArray, Int32Array, MapArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, Fields, Int32Type}; + use datafusion_comet_spark_expr::EvalMode; + use std::sync::Arc; + + let ff: &[u8] = &[0xFF]; + let replacement = "\u{fffd}".as_bytes(); + let key_bytes = BinaryArray::from(vec![ff, replacement]); + // The same keys dictionary-encoded, in an order `take` has to follow, past an entry no + // key references. + let dictionary_keys = DictionaryArray::::new( + Int32Array::from(vec![2, 1]), + Arc::new(BinaryArray::from(vec![b"zzz".as_slice(), replacement, ff])), + ); + let struct_keys = StructArray::from(vec![( + Arc::new(Field::new("s", DataType::Binary, true)), + Arc::new(key_bytes.clone()) as ArrayRef, + )]); + let values = BinaryArray::from(vec![ff, b"two".as_slice()]); + let map_type = |key_type: DataType, value_type: DataType| { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", key_type, false), + Field::new("value", value_type, true), + ])), + false, + )), + false, + ) + }; + let struct_type = + |inner: DataType| DataType::Struct(Fields::from(vec![Field::new("s", inner, true)])); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + + for (keys, from_key_type, to_key_type) in [ + ( + Arc::new(key_bytes.clone()) as ArrayRef, + DataType::Binary, + DataType::Utf8, + ), + ( + Arc::new(struct_keys) as ArrayRef, + struct_type(DataType::Binary), + struct_type(DataType::Utf8), + ), + ( + Arc::new(dictionary_keys) as ArrayRef, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)), + DataType::Utf8, + ), + ] { + let from_type = map_type(from_key_type.clone(), DataType::Binary); + let DataType::Map(entries, _) = &from_type else { + unreachable!() + }; + let DataType::Struct(entry_fields) = entries.data_type() else { + unreachable!() + }; + let input: ArrayRef = Arc::new(MapArray::new( + Arc::clone(entries), + OffsetBuffer::new(vec![0, 2].into()), + StructArray::new( + entry_fields.clone(), + vec![keys, Arc::new(values.clone())], + None, + ), + None, + false, + )); + let to_type = map_type(to_key_type.clone(), DataType::Utf8); + + let result = parquet_convert_array(input, &to_type, &options).unwrap(); + let map = result.as_map(); + let key_strings = match to_key_type { + DataType::Utf8 => map.keys().as_string::(), + _ => map.keys().as_struct().column(0).as_string::(), + }; + // The keys hold their original bytes, unchecked, so read them as bytes. + assert_eq!(key_strings.value_offsets(), &[0, 1, 4]); + assert_eq!(key_strings.value_data(), [ff, replacement].concat()); + let value_strings = map.values().as_string::(); + assert_eq!(value_strings.value(0), "\u{fffd}"); + assert_eq!(value_strings.value(1), "two"); + } + } + + /// Decoding a dictionary touches only the entries the batch references, and the wide + /// ill-formed entry another batch uses stays hidden behind a null. + #[test] + fn referenced_dictionary_values_hides_unreferenced_entries() { + use super::referenced_dictionary_values; + use arrow::array::{Array, ArrayRef, BinaryArray}; + use arrow::buffer::NullBuffer; + use std::sync::Arc; + + let wide = vec![0xFF; 86_016]; + let values: ArrayRef = Arc::new(BinaryArray::from(vec![ + b"normal".as_slice(), + wide.as_slice(), + b"other".as_slice(), + ])); + let hidden = referenced_dictionary_values(&values, &[0, 2, 0], None).unwrap(); + assert!(hidden.is_valid(0)); + assert!(hidden.is_null(1)); + assert!(hidden.is_valid(2)); + + // A null key's entry doesn't count as referenced. + let nulls = NullBuffer::from(vec![true, false, true]); + let hidden = referenced_dictionary_values(&values, &[0, 1, 0], Some(&nulls)).unwrap(); + assert!(hidden.is_null(1)); + + // Every entry referenced, or a small dictionary: the values are returned as they are. + let all = referenced_dictionary_values(&values, &[0, 1, 2], None).unwrap(); + assert!(Arc::ptr_eq(&all, &values)); + let small: ArrayRef = Arc::new(BinaryArray::from(vec![b"a".as_slice(), b"b"])); + let same = referenced_dictionary_values(&small, &[0], None).unwrap(); + assert!(Arc::ptr_eq(&same, &small)); + } + + /// The bytes a batch's keys expand a decoded dictionary to are counted from the decoded + /// values, so that a batch past what a string column holds is reported before `take` + /// overflows the offsets. The error itself needs a 2 GiB batch and is not exercised here. + #[test] + fn nested_dictionary_binary_to_string_counts_decoded_bytes() { + use super::{decoded_bytes, parquet_convert_array, SparkParquetOptions}; + use arrow::array::{Array, ArrayRef, AsArray, BinaryArray, DictionaryArray, Int32Array}; + use arrow::buffer::NullBuffer; + use arrow::datatypes::{DataType, Int32Type}; + use datafusion_comet_spark_expr::EvalMode; + use std::sync::Arc; + + let offsets = [0, 1 << 30, (1 << 30) + 5]; + assert_eq!(decoded_bytes(&offsets, &[0, 0, 0], None), 3 << 30); + let nulls = NullBuffer::from(vec![true, false, true]); + assert_eq!( + decoded_bytes(&offsets, &[0, 0, 1], Some(&nulls)), + (1 << 30) + 5 + ); + + // The expanded, decoded size is checked, not the input's: 8 keys that reference one + // 1 KiB entry decode to 24 KiB, which fits. + let wide = vec![0xFF; 1024]; + let values = BinaryArray::from(vec![b"one".as_slice(), wide.as_slice()]); + let input: ArrayRef = Arc::new(DictionaryArray::::new( + Int32Array::from(vec![1; 8]), + Arc::new(values), + )); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let result = parquet_convert_array(input, &DataType::Utf8, &options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.len(), 8); + assert_eq!(strings.value(7).len(), 3 * 1024); + } + /// Checks parser-backed I/O labels without constructing stores, including libhdfs overrides /// and rejection of unknown native schemes. Configured S3 aliases follow URL normalization. #[test] diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 0c14afe23af..04b806e3330 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -43,8 +43,10 @@ use crate::{cast_whole_num_to_binary, BinaryOutputStyle}; use crate::{EvalMode, SparkError}; use arrow::array::builder::{GenericStringBuilder, StringBuilder}; use arrow::array::{ - new_null_array, BinaryBuilder, GenericByteArray, ListArray, MapArray, StringArray, StructArray, + new_null_array, BinaryBuilder, GenericByteArray, GenericStringArray, ListArray, MapArray, + StringArray, StructArray, }; +use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, GenericBinaryType, Schema}; use arrow::error::ArrowError; use arrow::{ @@ -140,7 +142,8 @@ pub struct SparkCastOptions { /// Support casting unsigned ints to signed ints (used by Parquet SchemaAdapter) pub allow_cast_unsigned_ints: bool, /// We also use the cast logic for adapting Parquet schemas, so this flag is used - /// for that use case + /// for that use case. The reader owns the arrays then, so a cast may share their buffers; + /// see [cast_binary_to_string]. pub is_adapting_schema: bool, /// String to use to represent null values pub null_string: String, @@ -811,8 +814,13 @@ fn write_bracketed( 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. +/// Casts a binary array to a string array. Without a binary output style, or with the `UTF8` +/// style, the bytes are decoded the way the JVM's `new String(bytes, UTF_8)` does, replacing each +/// ill-formed sequence with U+FFFD. Spark's `Cast` keeps the bytes, so the strings match what +/// Spark returns but byte-level results can differ (#4764). Valid input is copied once, or +/// shares the input's buffers when a scan adapts its schema, where the reader owns them. A plain +/// cast can run on a batch imported over the Arrow C stream, where each buffer keeps the whole +/// batch alive, wide sibling columns included, for as long as a sort or a join holds the result. /// /// 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. @@ -829,12 +837,94 @@ fn cast_binary_to_string( let num_rows = input.len(); let offsets = input.value_offsets(); - let value_bytes = offsets[num_rows].as_usize() - offsets[0].as_usize(); + let start = offsets[0].as_usize(); + let end = offsets[num_rows].as_usize(); + // Bytes of the non-null rows. A null row can still span bytes, for example after `nullif`. + // Each run of non-null rows is one subtraction, so sparse nulls don't cost a pass over the rows. + let live_bytes = match input.nulls().filter(|nulls| nulls.null_count() > 0) { + Some(nulls) => nulls + .valid_slices() + .map(|(from, to)| offsets[to].as_usize() - offsets[from].as_usize()) + .sum(), + None => end - start, + }; + + // Check the bytes the rows span once. If they are valid UTF-8, copy them, or reuse the + // input's buffers when the reader owns them and the rows fill at least half of their + // allocations; otherwise copy the rows in the valid prefix and decode the rest. Skip this + // when null rows span more than half of those bytes: decoding the live rows one at a time + // then reads less and returns a compact buffer instead of keeping the null rows' bytes alive. + let utf8_style = matches!( + spark_cast_options.binary_output_style, + None | Some(BinaryOutputStyle::Utf8) + ); + let mut valid_prefix = ""; + let mut checked = false; + if utf8_style && 2 * live_bytes >= end - start { + checked = true; + let data = input.value_data(); + match std::str::from_utf8(&data[start..end]) { + Ok(valid) => { + // An offset splits a code point only if it points at a continuation byte. + if offsets.iter().all(|offset| { + let offset = offset.as_usize(); + offset == end || !matches!(data[offset], 0x80..=0xBF) + }) { + // Reuse a buffer only when the reader owns it and the rows fill at least + // half of its allocation. A slice of a large array shares the whole + // allocation, and memory accounting such as a sort's charges all of it to + // every batch that holds it, so the offsets, values and validity bitmap are + // copied then, without decoding the values again. An imported buffer reports + // its length as its capacity, so it is told apart by the caller instead. + let share = spark_cast_options.is_adapting_schema; + let offset_buffer = input.offsets().inner().inner(); + let offsets = + if share && start == 0 && fills(offset_buffer, offset_buffer.len()) { + input.offsets().clone() + } else { + // Rebases a slice's offsets onto the narrowed value buffer. + OffsetBuffer::from_lengths(input.offsets().lengths()) + }; + let values = if share && fills(input.values(), end - start) { + input.values().slice_with_length(start, end - start) + } else { + Buffer::from_slice_ref(&data[start..end]) + }; + let nulls = input.nulls().map(|nulls| { + if share && fills(nulls.buffer(), nulls.len().div_ceil(8)) { + nulls.clone() + } else { + // `sliced` shares the bitmap when the slice starts on a byte boundary. + let bits = Buffer::from_slice_ref(&nulls.inner().sliced()); + NullBuffer::new(BooleanBuffer::new(bits, 0, nulls.len())) + } + }); + // SAFETY: the value buffer is valid UTF-8 and every offset falls on a char + // boundary, which is what `GenericStringArray::try_new` checks. + let strings = + unsafe { GenericStringArray::::new_unchecked(offsets, values, nulls) }; + return Ok(Arc::new(strings)); + } + valid_prefix = valid; + } + // SAFETY: `valid_up_to` is the length of the longest prefix that is valid UTF-8. + Err(error) => { + valid_prefix = + unsafe { std::str::from_utf8_unchecked(&data[start..][..error.valid_up_to()]) } + } + } + } + + let value_bytes = live_bytes; // Upper bound on the encoded length so the value buffer is allocated once. For the - // JVM-lossy UTF-8 decoder path, valid UTF-8 sits at 1x; the builder grows on the rare - // invalid byte that expands to U+FFFD. Base64 rounds up to a full 4-char group via - // div_ceil, which already covers the group padding. + // JVM-lossy UTF-8 decoder path, valid UTF-8 sits at 1x and the bytes after the first + // ill-formed sequence at up to 3x, since U+FFFD takes three bytes; a batch that was not + // checked is sized for valid input. Base64 rounds up to a full 4-char group via div_ceil, + // which already covers the group padding. let capacity = match spark_cast_options.binary_output_style { + None | Some(BinaryOutputStyle::Utf8) if checked => { + valid_prefix.len() + 3 * value_bytes.saturating_sub(valid_prefix.len()) + } None | Some(BinaryOutputStyle::Utf8) => value_bytes, Some(BinaryOutputStyle::Basic) => 6 * value_bytes + 2 * num_rows, Some(BinaryOutputStyle::Base64) => 4 * value_bytes.div_ceil(3), @@ -845,7 +935,7 @@ fn cast_binary_to_string( 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() { + for (value, range) in input.iter().zip(offsets.windows(2)) { let Some(value) = value else { // The previous iteration always finalized its row with `append_value("")`, so no // bytes are pending in the builder here; a future edit that adds a `continue` or @@ -857,14 +947,31 @@ fn cast_binary_to_string( // 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. + // render via `new String(bytes, UTF_8)`. Values in the valid prefix are copied as is; + // the rest go 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). 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)) + let row_start = range[0].as_usize() - start; + let row_end = range[1].as_usize() - start; + match valid_prefix.get(row_start..row_end) { + Some(valid) => builder.write_str(valid), + None => { + // The row holds the first ill-formed sequence, follows it, or starts or + // ends inside a code point. Copy the part of it that was checked when + // the row starts on a character boundary, and decode only the rest: the + // decoder restarts at every boundary, so the result is the same as + // decoding the whole row. + let validated = valid_prefix + .get(row_start..row_end.min(valid_prefix.len())) + .unwrap_or(""); + builder.write_str(validated).and_then(|()| { + builder.write_str(&decode_utf8_spark_lossy(&value[validated.len()..])) + }) + } + } } Some(BinaryOutputStyle::Basic) => write_bracketed(&mut builder, value, ", ", write_i8), Some(BinaryOutputStyle::Base64) => { @@ -883,7 +990,25 @@ fn cast_binary_to_string( written.expect("writing to a string builder cannot fail"); builder.append_value(""); } - Ok(Arc::new(builder.finish())) + let strings = builder.finish(); + // A replacement that grows the value buffer past its capacity doubles the allocation, which + // memory accounting then charges in full. Return a compact copy when the rows fill less than + // half of it. + if utf8_style && !fills(strings.values(), strings.values().len()) { + let values = Buffer::from_slice_ref(strings.values().as_slice()); + let (offsets, _, nulls) = strings.into_parts(); + // SAFETY: the builder produced valid UTF-8 and offsets on character boundaries, and + // only the value buffer's allocation changes. + let strings = unsafe { GenericStringArray::::new_unchecked(offsets, values, nulls) }; + return Ok(Arc::new(strings)); + } + Ok(Arc::new(strings)) +} + +/// True when the `used` bytes of `buffer` fill at least half of its allocation, which is when +/// sharing the buffer costs less than a compact copy. +fn fills(buffer: &Buffer, used: usize) -> bool { + 2 * used >= buffer.capacity() } #[cfg(test)] @@ -1045,6 +1170,386 @@ mod tests { assert_eq!(result.value(2), "héllo"); } + /// Cast options of a scan adapting its schema, where the reader owns the input's buffers + /// and a valid input can share them. + fn scan_cast_options() -> SparkCastOptions { + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + cast_options.is_adapting_schema = true; + cast_options + } + + #[test] + fn test_cast_binary_to_string_valid_utf8_reuses_buffers() { + use arrow::buffer::NullBuffer; + + // The rows fill the value buffer's allocation, which a builder would round up. + let input = BinaryArray::new( + OffsetBuffer::new(vec![0, 6, 6, 6].into()), + Buffer::from_vec("héllo".as_bytes().to_vec()), + Some(NullBuffer::from(vec![true, false, true])), + ); + for style in [None, Some(BinaryOutputStyle::Utf8)] { + let mut cast_options = scan_cast_options(); + cast_options.binary_output_style = style; + + let result = cast_binary_to_string::(&input, &cast_options).unwrap(); + + let strings = result.as_string::(); + assert_eq!(strings.values().as_ptr(), input.values().as_ptr()); + assert_eq!(strings.nulls(), input.nulls()); + assert_eq!(strings.value(0), "héllo"); + assert!(strings.is_null(1)); + assert_eq!(strings.value(2), ""); + } + } + + /// Outside a scan, the input may be a batch imported over the Arrow C stream, where each + /// buffer keeps the whole batch alive. A plain cast copies the valid bytes once instead. + #[test] + fn test_cast_binary_to_string_outside_scan_copies_valid_utf8() { + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let payload = [b'a'; 1024]; + let input = BinaryArray::from_opt_vec( + (0..64) + .map(|i| (i % 8 != 3).then_some(payload.as_slice())) + .collect(), + ); + for input in [input.clone(), input.slice(8, 48)] { + let result = cast_binary_to_string::(&input, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_ne!(strings.values().as_ptr(), input.value(0).as_ptr()); + assert_ne!( + strings.value_offsets().as_ptr(), + input.value_offsets().as_ptr() + ); + assert_ne!( + strings.nulls().unwrap().buffer().as_ptr(), + input.nulls().unwrap().buffer().as_ptr() + ); + // The copies are compact and hold the same rows, rebased to the slice's bytes. + let offsets = input.value_offsets(); + let bytes = &input.value_data()[offsets[0] as usize..offsets[input.len()] as usize]; + assert_eq!(strings.value_data(), bytes); + assert!(strings.values().capacity() < bytes.len() + 128); + assert_eq!(strings.value_offsets()[0], 0); + assert!(strings.offsets().lengths().eq(input.offsets().lengths())); + assert_eq!(strings.nulls(), input.nulls()); + assert_eq!(strings.value(0).as_bytes(), payload); + assert!(strings.is_null(3)); + } + + let no_nulls = BinaryArray::from(vec![payload.as_slice(); 8]); + let result = cast_binary_to_string::(&no_nulls, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_ne!(strings.values().as_ptr(), no_nulls.value(0).as_ptr()); + assert!(strings.nulls().is_none()); + assert_eq!(strings.value_offsets(), no_nulls.value_offsets()); + assert_eq!(strings.value(7).as_bytes(), payload); + } + + #[test] + fn test_cast_binary_to_string_invalid_bytes_under_null_fall_back() { + use arrow::buffer::{Buffer, NullBuffer}; + + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + + let under_null = BinaryArray::new( + OffsetBuffer::new(vec![0, 3, 4, 6].into()), + Buffer::from(b"abc\xff\xc3\xa9".as_slice()), + Some(NullBuffer::from(vec![true, false, true])), + ); + let result = cast_binary_to_string::(&under_null, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.len(), 3); + assert_eq!(strings.value(0), "abc"); + assert!(strings.is_null(1)); + assert_eq!(strings.value(2), "é"); + + // An offset that splits a code point turns both halves into U+FFFD. + let split = BinaryArray::from(vec![b"\xc3".as_slice(), b"\xa9".as_slice()]); + let result = cast_binary_to_string::(&split, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.value(0), "\u{FFFD}"); + assert_eq!(strings.value(1), "\u{FFFD}"); + } + + #[test] + fn test_cast_binary_to_string_slice_checks_only_its_rows() { + use arrow::buffer::NullBuffer; + + let cast_options = scan_cast_options(); + // `\xff`, `ab`, null, `cé` and `\xff`, in a value buffer exactly as large as its bytes. + let parent = BinaryArray::new( + OffsetBuffer::new(vec![0, 1, 3, 3, 6, 7].into()), + Buffer::from_vec(b"\xffabc\xc3\xa9\xff".to_vec()), + Some(NullBuffer::from(vec![true, true, false, true, true])), + ); + + let valid = parent.slice(1, 3); + let result = cast_binary_to_string::(&valid, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.values().as_ptr(), valid.value(0).as_ptr()); + assert_eq!(strings.value_data(), "abcé".as_bytes()); + assert_eq!(strings.value_offsets(), [0, 2, 2, 5]); + assert_eq!(strings.nulls(), valid.nulls()); + assert_eq!(strings.value(0), "ab"); + assert!(strings.is_null(1)); + assert_eq!(strings.value(2), "cé"); + + // In a slice with an invalid row, the rows before it are unchanged and it becomes U+FFFD. + let invalid = parent.slice(1, 4); + let result = cast_binary_to_string::(&invalid, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.len(), 4); + assert_eq!(strings.value(0), "ab"); + assert!(strings.is_null(1)); + assert_eq!(strings.value(2), "cé"); + assert_eq!(strings.value(3), "\u{FFFD}"); + + // An offset inside a slice that splits a code point turns both halves into U+FFFD. + let split = BinaryArray::from(vec![ + b"\xff".as_slice(), + b"\xc3".as_slice(), + b"\xa9".as_slice(), + ]) + .slice(1, 2); + let result = cast_binary_to_string::(&split, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.value(0), "\u{FFFD}"); + assert_eq!(strings.value(1), "\u{FFFD}"); + } + + #[test] + fn test_cast_binary_to_string_small_slice_gets_compact_copy() { + use arrow::buffer::MutableBuffer; + + let cast_options = scan_cast_options(); + let payload = [b'a'; 1024]; + let parent = BinaryArray::from(vec![payload.as_slice(); 64]); + + // The slice's bytes fill an eighth of the allocation it shares, so they are copied. + let small = parent.slice(16, 8); + let result = cast_binary_to_string::(&small, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_ne!(strings.values().as_ptr(), small.value(0).as_ptr()); + assert_eq!(strings.value_data().len(), 8 * payload.len()); + assert!(strings.values().capacity() < 2 * 8 * payload.len()); + assert_eq!(strings.value_offsets()[0], 0); + assert!((0..8).all(|i| strings.value(i).as_bytes() == payload)); + + // The first slice's offsets need no rebasing, but they fill little of the parent's + // offsets allocation, so they are rebuilt as well. + let first = parent.slice(0, 8); + let result = cast_binary_to_string::(&first, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_ne!(strings.values().as_ptr(), first.value(0).as_ptr()); + assert_ne!( + strings.value_offsets().as_ptr(), + first.value_offsets().as_ptr() + ); + assert_eq!(strings.value_offsets(), first.value_offsets()); + + // An unsliced array whose value buffer is oversized keeps its offsets. + let mut oversized = MutableBuffer::with_capacity(8 * payload.len()); + oversized.extend_from_slice(&payload); + let unsliced = BinaryArray::new( + OffsetBuffer::new(vec![0, 1024].into()), + oversized.into(), + None, + ); + let result = cast_binary_to_string::(&unsliced, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_ne!(strings.values().as_ptr(), unsliced.value(0).as_ptr()); + assert_eq!( + strings.value_offsets().as_ptr(), + unsliced.value_offsets().as_ptr() + ); + assert_eq!(strings.value(0).as_bytes(), payload); + + // A slice whose bytes fill at least half of it still reuses the buffer. + let large = parent.slice(16, 40); + let result = cast_binary_to_string::(&large, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.values().as_ptr(), large.value(0).as_ptr()); + assert_eq!(strings.value_data().len(), 40 * payload.len()); + } + + #[test] + fn test_cast_binary_to_string_small_slice_gets_compact_nulls() { + let cast_options = scan_cast_options(); + // One null row in every eight, like a nullable grouped key that is emitted in slices. + let parent = BinaryArray::from_opt_vec( + (0..4096) + .map(|i| (i % 8 != 3).then_some(b"ab".as_slice())) + .collect(), + ); + let parent_nulls = parent.nulls().unwrap().buffer(); + + // The slice's bits fill a thirty-second of the parent's bitmap, so they are copied, + // whether or not the slice starts on a byte boundary. + for offset in [1024, 1027] { + let small = parent.slice(offset, 128); + let result = cast_binary_to_string::(&small, &cast_options).unwrap(); + let strings = result.as_string::(); + let nulls = strings.nulls().unwrap(); + assert_ne!(nulls.buffer().as_ptr(), parent_nulls.as_ptr()); + assert!(2 * nulls.buffer().capacity() < parent_nulls.capacity()); + assert_eq!(strings.nulls(), small.nulls()); + assert_eq!(strings.null_count(), 16); + } + + // A slice whose bits fill at least half of the bitmap shares it. + let large = parent.slice(1024, 3072); + let result = cast_binary_to_string::(&large, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!( + strings.nulls().unwrap().buffer().as_ptr(), + large.nulls().unwrap().buffer().as_ptr() + ); + assert_eq!(strings.nulls(), large.nulls()); + } + + #[test] + fn test_cast_binary_to_string_mostly_null_copies_live_rows() { + use arrow::buffer::NullBuffer; + + let cast_options = scan_cast_options(); + let payload = [b'a'; 1024]; + let values = BinaryArray::from(vec![payload.as_slice(); 64]); + let with_nulls = |live: &dyn Fn(usize) -> bool| { + let nulls = NullBuffer::from((0..64).map(live).collect::>()); + BinaryArray::new( + values.offsets().clone(), + values.values().clone(), + Some(nulls), + ) + .slice(16, 32) + }; + + // Nulls hide most of the slice's bytes: only the two live rows are copied. + let mostly_null = with_nulls(&|i| i == 20 || i == 40); + let result = cast_binary_to_string::(&mostly_null, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_ne!(strings.values().as_ptr(), mostly_null.value(0).as_ptr()); + assert_eq!(strings.value_data().len(), 2 * payload.len()); + assert!(strings.values().capacity() < 4 * payload.len()); + assert_eq!(strings.null_count(), 30); + assert_eq!(strings.value(4).as_bytes(), payload); + assert_eq!(strings.value(24).as_bytes(), payload); + + // Nulls hide less than half: the slice's bytes are checked and reused. + let mostly_live = with_nulls(&|i| i % 4 != 0); + let result = cast_binary_to_string::(&mostly_live, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.values().as_ptr(), mostly_live.value(0).as_ptr()); + assert_eq!(strings.nulls(), mostly_live.nulls()); + } + + /// One null per 128 short rows still takes the fast path, which counts the live bytes by + /// run of non-null rows rather than by row. + #[test] + fn test_cast_binary_to_string_sparse_nulls_share_buffer() { + let cast_options = scan_cast_options(); + let rows: Vec> = (0..8192) + .map(|i| (i % 128 != 127).then_some(b"eightbyt".as_slice())) + .collect(); + let input = BinaryArray::from(rows); + let result = cast_binary_to_string::(&input, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.values().as_ptr(), input.value(0).as_ptr()); + assert_eq!(strings.null_count(), 64); + assert_eq!(strings.value(126), "eightbyt"); + assert!(strings.is_null(127)); + + // Null rows that span bytes count for nothing, run by run. + let hidden = BinaryArray::new( + input.offsets().clone(), + input.values().clone(), + Some(NullBuffer::from( + (0..8192).map(|i| i % 2 == 0).collect::>(), + )), + ); + let result = cast_binary_to_string::(&hidden, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.values().as_ptr(), hidden.value(0).as_ptr()); + assert_eq!(strings.null_count(), 4096); + } + + /// The rows from the first ill-formed sequence on decode one at a time. The row that holds + /// it reuses its checked prefix, and a row that starts or ends inside a code point is decoded + /// whole; both match decoding each row on its own. + #[test] + fn test_cast_binary_to_string_row_decoding_matches_per_row_decoding() { + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let wide = "é".repeat(50_000).into_bytes(); + let truncated: Vec = [wide.as_slice(), b"abc", &[0xE9]].concat(); + let rows: Vec> = vec![ + "é".as_bytes().to_vec(), + truncated, + b"ok".to_vec(), + vec![0xC3], + vec![0xA9, 0xFF, b'x'], + vec![0xFF], + "日本".as_bytes().to_vec(), + ]; + // Rows 3 and 4 split a code point between them, so the span of all rows is valid up to + // row 1's truncated byte and row 3 starts on a boundary while row 4 does not. + for (first, last) in [(0, 7), (3, 7), (0, 4), (1, 2)] { + let input = BinaryArray::from_iter_values(rows[first..last].iter()); + let result = cast_binary_to_string::(&input, &cast_options).unwrap(); + let strings = result.as_string::(); + for (i, row) in rows[first..last].iter().enumerate() { + assert_eq!( + strings.value(i), + decode_utf8_spark_lossy(row), + "rows {first}..{last}, row {i}" + ); + } + } + + // A slice whose first row is the one with the truncated byte reuses its checked prefix. + let input = BinaryArray::from_iter_values(rows.iter()).slice(1, 2); + let result = cast_binary_to_string::(&input, &cast_options).unwrap(); + let strings = result.as_string::(); + assert_eq!(strings.value(0), decode_utf8_spark_lossy(&rows[1])); + assert_eq!(strings.value(1), "ok"); + } + + /// A replacement that grows the builder's value buffer past its capacity doubles the + /// allocation. The result holds a compact copy instead. + #[test] + fn test_cast_binary_to_string_fallback_buffer_is_compact() { + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let payload = [b'a'; 64]; + let mut rows = vec![payload.as_slice(); 8191]; + let last = [b"aaaaaaaa".as_slice(), &[0xFF]].concat(); + rows.push(last.as_slice()); + let input = BinaryArray::from(rows); + let result = cast_binary_to_string::(&input, &cast_options).unwrap(); + let strings = result.as_string::(); + let len = strings.values().len(); + assert_eq!(len, 8191 * 64 + 8 + 3); + assert!( + strings.values().capacity() < len + 128, + "capacity {} for {len} bytes", + strings.values().capacity() + ); + assert_eq!(strings.value(0).as_bytes(), payload); + assert_eq!(strings.value(8191), "aaaaaaaa\u{fffd}"); + + // An ill-formed byte in the first row sizes the buffer for the worst case, which is + // then compacted as well. + let first = [&[0xFF], b"aaaaaaaa".as_slice()].concat(); + let mut rows = vec![first.as_slice()]; + rows.extend(std::iter::repeat_n(payload.as_slice(), 8191)); + let input = BinaryArray::from(rows); + let result = cast_binary_to_string::(&input, &cast_options).unwrap(); + let strings = result.as_string::(); + let len = strings.values().len(); + assert!(strings.values().capacity() < len + 128); + assert_eq!(strings.value(0), "\u{fffd}aaaaaaaa"); + } + #[test] fn test_cast_unsupported_timestamp_to_date() { // Since datafusion uses chrono::Datetime internally not all dates representable by TimestampMicrosecondType are supported diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 8c0cd4a37a6..a182b46eb5b 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -146,6 +146,96 @@ abstract class ParquetReadSuite extends CometTestBase { } } + test("native scan decodes invalid UTF-8 in nested BINARY read as string") { + withTempDir { dir => + val path = new Path(dir.toURI.toString, "nested-binary.parquet") + val parquetSchema = MessageTypeParser.parseMessageType("""message root { + | optional group arr (LIST) { + | repeated group list { + | optional binary element; + | } + | } + | optional group st { + | optional binary name; + | } + | optional group m (MAP) { + | repeated group key_value { + | required binary key; + | optional binary value; + | } + | } + | optional int32 id; + |} + |""".stripMargin) + val valid = Binary.fromString("one") + val invalid = Binary.fromConstantByteArray(Array[Byte](0x41, 0xff.toByte, 0x42)) + // The same bytes with U+FFFD in place of FF: a second map key that decoding would + // collapse onto the invalid one. + val replaced = Binary.fromString("A\uFFFDB") + val writer = createParquetWriter(parquetSchema, path) + try { + // Valid values, then a null list element, struct field and map value, then invalid + // values, then null containers, so one batch mixes nulls, valid and invalid values. + Seq(Some(valid), None, Some(invalid)).zipWithIndex.foreach { case (value, id) => + val row = new SimpleGroup(parquetSchema) + val arr = row.addGroup("arr") + arr.addGroup("list").append("element", valid) + val element = arr.addGroup("list") + val st = row.addGroup("st") + val m = row.addGroup("m") + val entry = m.addGroup("key_value").append("key", value.getOrElse(valid)) + value.foreach { value => + element.append("element", value) + st.append("name", value) + entry.append("value", value) + } + if (value.contains(invalid)) { + m.addGroup("key_value").append("key", replaced).append("value", valid) + } + row.add("id", id) + writer.write(row) + } + val nullContainers = new SimpleGroup(parquetSchema) + nullContainers.add("id", 3) + writer.write(nullContainers) + } finally { + writer.close() + } + + val readSchema = "arr array, st struct, m map, id int" + // Reads with binaryAsString, then with a read schema that declares the strings. + Seq(("true", None), ("false", Some(readSchema))).foreach { case (binaryAsString, schema) => + withSQLConf(SQLConf.PARQUET_BINARY_AS_STRING.key -> binaryAsString) { + val reader = schema.map(spark.read.schema(_)).getOrElse(spark.read) + val df = reader.parquet(path.toString) + assert( + df.schema.catalogString === + "struct,st:struct,m:map,id:int>") + + // Collected values match Spark, and the scan runs natively. + val (_, cometPlan) = checkSparkAnswerAndOperator(df) + assert( + collect(cometPlan) { case scan: CometNativeScanExec => scan }.nonEmpty, + s"Expected a CometNativeScanExec:\n${cometPlan.treeString}") + + // Each invalid nested value holds U+FFFD (EF BF BD) in place of FF, except a map + // key, which keeps its bytes as Spark does: the two keys stay distinct, and looking + // up the replaced one returns its own value. + checkAnswer( + df.where("id = 2") + .selectExpr( + "hex(arr[1])", + "hex(st.name)", + "hex(map_keys(m)[0])", + "hex(map_values(m)[0])", + "size(map_from_entries(map_entries(m)))", + "m['A\uFFFDB']"), + Row("41EFBFBD42", "41EFBFBD42", "41FF42", "41EFBFBD42", 2, "one")) + } + } + } + } + test("basic data types") { Seq(7, 1024).foreach { batchSize => withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {