diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 6fcafbf71c2..dca663d9228 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -79,7 +79,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 @@ -98,9 +99,15 @@ 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. Comet's native Parquet scan never -produces such strings (it rejects them, see below), and neither do most string expressions. They can -reach Comet in these ways, which are the ones to check before relying on a result: - +produces such strings from UTF8-annotated columns (it rejects them, see below), and neither do most +string expressions. They can reach Comet in these ways, which are the ones to check before relying on +a result: + +- **Parquet BINARY columns read as strings.** A BINARY column without the UTF8 annotation can hold + bytes that are not valid UTF-8 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; see + [Parquet](../parquet.md#reads-that-differ-from-spark). - **Text files.** Spark's Text reader keeps the bytes of each line as they are, so a Text file in an encoding other than UTF-8, such as Latin-1, holds strings that are not valid UTF-8. Comet decodes them when it converts a Text scan to Arrow, which it does only when diff --git a/docs/source/user-guide/latest/parquet.md b/docs/source/user-guide/latest/parquet.md index 4dd198bdc47..0d6f2ae765e 100644 --- a/docs/source/user-guide/latest/parquet.md +++ b/docs/source/user-guide/latest/parquet.md @@ -280,12 +280,14 @@ for production, so do not turn it off to avoid an entry on this page. read back wrong inside lists written in a legacy layout, as Hive and `spark.sql.parquet.writeLegacyFormat=true` write them ([#7029](https://github.com/apache/datafusion-comet/issues/7029)). -- **Nested BINARY values that are not valid UTF-8, read as STRING.** This applies to a BINARY - column without the UTF8 annotation that is read as `string`, through the read schema or - `spark.sql.parquet.binaryAsString=true`. Spark returns the bytes. Comet returns NULL for each - such nested value. Top-level values match Spark in a release build, because the invalid bytes - reach an Arrow string array unchecked. A debug build panics on them - ([#7032](https://github.com/apache/datafusion-comet/issues/7032), part of +- **BINARY values that are not valid UTF-8, read as STRING.** This applies to a BINARY column + without the UTF8 annotation that is read as `string`, through the read schema or + `spark.sql.parquet.binaryAsString=true`. Spark keeps the bytes. In a struct field, list element + or map value, Comet replaces each ill-formed sequence with `U+FFFD`, as described in + [Strings with non-UTF-8 bytes](compatibility/index.md#strings-with-non-utf-8-bytes). Top-level + values, and map keys at any depth, reach an Arrow string array unchecked, which keeps keys that + differ only in such bytes distinct: they match Spark in a release build, and a debug build + panics on top-level values ([#7032](https://github.com/apache/datafusion-comet/issues/7032), part of [#4764](https://github.com/apache/datafusion-comet/issues/4764)). - **Field-id reads where one file column matches a field by id and a sibling by name.** With `spark.sql.parquet.fieldId.read.enabled=true`, Comet returns NULL for the sibling at the top @@ -354,6 +356,15 @@ With `spark.comet.write.parquet.enabled=true`: ([#7035](https://github.com/apache/datafusion-comet/issues/7035)). - **A `DATE` column read as `INTERVAL YEAR TO MONTH`.** Spark reads it, and Comet fails the read ([#5188](https://github.com/apache/datafusion-comet/issues/5188)). +- **A nested dictionary-encoded BINARY column, read as STRING, whose batch decodes to more than + 2 GiB.** In a struct field, list element, map key or map value, Comet expands the dictionary + into one string array per batch, whose offsets hold at most 2,147,483,647 bytes. A batch of + `spark.comet.batchSize` rows (8192 by default) fails the scan when the values its rows reference + add up to more, counting each reference of a shared entry, even for a query that never reads the + strings. Values that are not valid UTF-8 decode to up to three times their size, since each + ill-formed sequence becomes `U+FFFD`. Spark keeps the dictionary and reads it row by row. Lower + `spark.comet.batchSize` so that a batch holds fewer rows + ([#7254](https://github.com/apache/datafusion-comet/issues/7254)). ### Writes that fail where Spark succeeds diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index d8ffee1aa70..f2e3e1e4b6c 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -834,6 +834,91 @@ mod tests { assert!(err.contains("rebase"), "unexpected error: {err}"); } + /// 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, + true, + "CORRECTED", + "CORRECTED", + ) + .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}" + ); + } + } + /// Writes a metadata-free parquet file with one INT64 `TIMESTAMP_MICROS` column (`ts`) and /// one INT96 column (`ts96`) through the low-level writer (arrow's writer cannot emit /// INT96), then scans it with the given session read modes. `int96_days` is the day count diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 29d9c9d278b..42f0bab544f 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -19,10 +19,11 @@ use crate::execution::operators::ExecutionError; use crate::parquet::field_resolver::FieldResolver; 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::{ @@ -41,7 +42,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}; @@ -207,16 +208,19 @@ fn parquet_convert_array( to_type: &DataType, parquet_options: &SparkParquetOptions, ) -> DataFusionResult { - parquet_convert_array_impl(array, to_type, parquet_options, None).map(|(array, _)| array) + parquet_convert_array_impl(array, to_type, parquet_options, None, false).map(|(array, _)| array) } // The boolean records whether any requested leaf exists in the file schema, independent // of the values' nullness. Spark propagates missing children through nested containers. +// `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<(ArrayRef, bool)> { use DataType::*; let from_type = array.data_type(); @@ -239,6 +243,7 @@ fn parquet_convert_array_impl( to_type, parquet_options, visible.as_ref(), + in_map_key, ), ( List(_) | LargeList(_) | FixedSizeList(_, _) | ListView(_) | LargeListView(_), @@ -256,6 +261,7 @@ fn parquet_convert_array_impl( to_inner_type.data_type(), parquet_options, child_visibility.as_ref(), + in_map_key, )?; if !field_overlap && parquet_options.return_null_struct_if_all_fields_missing { return Ok((new_null_array(to_type, array.len()), false)); @@ -318,7 +324,7 @@ fn parquet_convert_array_impl( )) } (Map(_, ordered_from), Map(_, ordered_to)) if ordered_from == ordered_to => - return parquet_convert_map_to_map(array.as_map(), to_type, parquet_options, *ordered_to, visible.as_ref(), checked_timestamp_overflow) + return 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. @@ -342,6 +348,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), true)); + } + // 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); + 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)?) @@ -487,6 +547,7 @@ fn parquet_convert_struct_to_struct( to_type: &DataType, parquet_options: &SparkParquetOptions, parent_nulls: Option<&NullBuffer>, + in_map_key: bool, ) -> DataFusionResult<(ArrayRef, bool)> { match (from_type, to_type) { (DataType::Struct(from_fields), DataType::Struct(to_fields)) => { @@ -501,6 +562,7 @@ fn parquet_convert_struct_to_struct( to_field.data_type(), parquet_options, parent_nulls, + in_map_key, )?; cast_fields.push(converted); field_overlap |= child_overlap; @@ -541,6 +603,7 @@ fn parquet_convert_map_to_map( to_ordered: bool, parent_nulls: Option<&NullBuffer>, checked_timestamp_overflow: bool, + in_map_key: bool, ) -> Result<(ArrayRef, bool), DataFusionError> { match to_data_type { DataType::Map(entries_field, _) => { @@ -561,12 +624,14 @@ fn parquet_convert_map_to_map( key_field.data_type(), parquet_options, child_visibility.as_ref(), + true, )?; let (value_array, value_overlap) = parquet_convert_array_impl( Arc::clone(from.values()), value_field.data_type(), parquet_options, child_visibility.as_ref(), + in_map_key, )?; let field_overlap = key_overlap || value_overlap; if !field_overlap && parquet_options.return_null_struct_if_all_fields_missing { @@ -594,6 +659,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() { @@ -1110,6 +1216,195 @@ 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"); + + 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"); + + 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"); + 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 884573d01ef..205b184e7ff 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::{format_decimal_str, DataType, GenericBinaryType, Schema}; use arrow::error::ArrowError; use arrow::{ @@ -170,7 +172,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, @@ -858,8 +861,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. @@ -876,12 +884,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), @@ -892,7 +982,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 @@ -904,14 +994,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) => { @@ -930,7 +1037,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)] @@ -1145,6 +1270,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"); + 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"); + 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"); + + 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"); + 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"); + 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 65d306d69ba..56743367a4a 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -147,6 +147,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) {