diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs index b2d4a57431..2539eb5b26 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -18,8 +18,10 @@ mod error; mod query_context; pub mod tracing; +mod utf8; mod utils; pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult}; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; +pub use utf8::decode_string_arrays; pub use utils::{bytes_to_i128, decode_utf8_spark_lossy}; diff --git a/native/common/src/utf8.rs b/native/common/src/utf8.rs new file mode 100644 index 0000000000..d3c038941e --- /dev/null +++ b/native/common/src/utf8.rs @@ -0,0 +1,444 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::decode_utf8_spark_lossy; +use arrow::array::{ + downcast_dictionary_array, Array, ArrayRef, FixedSizeListArray, GenericListArray, + GenericStringArray, GenericStringBuilder, MapArray, OffsetSizeTrait, StructArray, +}; +use arrow::datatypes::DataType; +use arrow::error::ArrowError; +use std::sync::Arc; + +/// Ensure every `Utf8`/`LargeUtf8` array reachable from `array` holds valid UTF-8, decoding invalid +/// bytes the way Spark renders `StringType`. Returns the same `Arc` (zero-copy) when nothing needed +/// decoding. Used at the JVM->native FFI import boundary, where arrow's `from_ffi` builds string +/// arrays via `new_unchecked` and does not validate UTF-8. +pub fn decode_string_arrays(array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Utf8 => decode_generic_string::(array), + DataType::LargeUtf8 => decode_generic_string::(array), + DataType::Dictionary(_, value_type) + if matches!(value_type.as_ref(), DataType::Utf8 | DataType::LargeUtf8) => + { + // Capture the original Arc before `downcast_dictionary_array!` shadows `array`, so the + // unchanged branch returns it verbatim, preserving the zero-copy contract that the + // Struct/List/Map arms rely on via `Arc::ptr_eq`. + let original = Arc::clone(array); + downcast_dictionary_array!( + array => { + let values = array.values(); + let decoded = decode_string_arrays(values)?; + if Arc::ptr_eq(&decoded, values) { + Ok(original) + } else { + Ok(Arc::new(array.with_values(decoded))) + } + } + t => unreachable!("dictionary type checked by guard: {t}"), + ) + } + DataType::Struct(fields) => { + let s = array + .as_any() + .downcast_ref::() + .expect("data type checked by caller"); + let mut changed = false; + let mut columns = Vec::with_capacity(s.num_columns()); + for col in s.columns() { + let decoded = decode_string_arrays(col)?; + changed |= !Arc::ptr_eq(&decoded, col); + columns.push(decoded); + } + if !changed { + return Ok(Arc::clone(array)); + } + Ok(Arc::new(StructArray::new( + fields.clone(), + columns, + s.nulls().cloned(), + ))) + } + DataType::List(_) => decode_list::(array), + DataType::LargeList(_) => decode_list::(array), + DataType::FixedSizeList(field, size) => { + let list = array + .as_any() + .downcast_ref::() + .expect("data type checked by caller"); + let values = list.values(); + let decoded = decode_string_arrays(values)?; + if Arc::ptr_eq(&decoded, values) { + return Ok(Arc::clone(array)); + } + Ok(Arc::new(FixedSizeListArray::try_new( + Arc::clone(field), + *size, + decoded, + list.nulls().cloned(), + )?)) + } + DataType::Map(field, ordered) => { + let map = array + .as_any() + .downcast_ref::() + .expect("data type checked by caller"); + let entries: ArrayRef = Arc::new(map.entries().clone()); + let decoded = decode_string_arrays(&entries)?; + if Arc::ptr_eq(&decoded, &entries) { + return Ok(Arc::clone(array)); + } + let decoded_struct = decoded + .as_any() + .downcast_ref::() + .expect("entries are a struct") + .clone(); + Ok(Arc::new(MapArray::try_new( + Arc::clone(field), + map.offsets().clone(), + decoded_struct, + map.nulls().cloned(), + *ordered, + )?)) + } + // Note: string *view* types (Utf8View, BinaryView, etc.) fall through here undecoded. + // Spark's JVM->native FFI only ever emits StringType as Utf8/LargeUtf8 (plus dictionaries + // and nesting thereof), so this is safe today. If a future producer starts handing a + // view-typed string column across FFI, a decode arm must be added above, or the + // from_utf8_unchecked UB this function exists to prevent will silently return. + _ => Ok(Arc::clone(array)), + } +} + +fn decode_generic_string(array: &ArrayRef) -> Result { + let arr = array + .as_any() + .downcast_ref::>() + .expect("data type checked by caller"); + let len = arr.len(); + if len == 0 { + return Ok(Arc::clone(array)); + } + let values: &[u8] = arr.value_data(); + let offsets = arr.value_offsets(); // &[O], length == len + 1 + let start = offsets[0].as_usize(); + let end = offsets[len].as_usize(); + + // Fast path: the used byte range parses as UTF-8 AND no element boundary splits a codepoint. + // Both are required: a whole-buffer-valid "é" (C3 A9) with per-element offsets [0,1,2] yields + // element slices `C3` and `A9`, each invalid, and `value()` would decode them unchecked (UB). + if std::str::from_utf8(&values[start..end]).is_ok() { + let mut boundaries_ok = true; + for off in &offsets[1..len] { + let o = off.as_usize(); + // A boundary landing on a UTF-8 continuation byte (0b10xx_xxxx) splits a codepoint. + if o < values.len() && (values[o] & 0xC0) == 0x80 { + boundaries_ok = false; + break; + } + } + if boundaries_ok { + return Ok(Arc::clone(array)); + } + } + + // Slow path: rebuild element-by-element via the Spark-lossy decoder. We slice the raw values + // buffer directly rather than calling `arr.value(i)`, which uses `from_utf8_unchecked`. + let mut builder = GenericStringBuilder::::with_capacity(len, end - start); + for i in 0..len { + if arr.is_null(i) { + builder.append_null(); + } else { + let s = offsets[i].as_usize(); + let e = offsets[i + 1].as_usize(); + builder.append_value(decode_utf8_spark_lossy(&values[s..e])); + } + } + Ok(Arc::new(builder.finish())) +} + +fn decode_list(array: &ArrayRef) -> Result { + let list = array + .as_any() + .downcast_ref::>() + .expect("data type checked by caller"); + let values = list.values(); + let decoded = decode_string_arrays(values)?; + if Arc::ptr_eq(&decoded, values) { + return Ok(Arc::clone(array)); + } + let field = match array.data_type() { + DataType::List(f) | DataType::LargeList(f) => Arc::clone(f), + _ => unreachable!("decode_list called on non-list"), + }; + Ok(Arc::new(GenericListArray::::try_new( + field, + list.offsets().clone(), + decoded, + list.nulls().cloned(), + )?)) +} + +#[cfg(test)] +mod walker_tests { + use super::decode_string_arrays; + use arrow::array::{make_array, Array, ArrayData, ArrayRef, LargeStringArray, StringArray}; + use arrow::buffer::Buffer; + use arrow::datatypes::DataType; + use std::sync::Arc; + + /// Build a (possibly invalid) Utf8 array from raw offsets + value bytes, the way an FFI import + /// would deliver it (no validation). `build_unchecked` mirrors arrow's `from_ffi`. + fn utf8_unchecked(offsets: &[i32], values: &[u8], len: usize) -> ArrayRef { + let data = unsafe { + ArrayData::builder(DataType::Utf8) + .len(len) + .add_buffer(Buffer::from_slice_ref(offsets)) + .add_buffer(Buffer::from(values.to_vec())) + .build_unchecked() + }; + make_array(data) + } + + #[test] + fn valid_utf8_is_zero_copy() { + let input: ArrayRef = Arc::new(StringArray::from(vec!["a", "é", "🦀"])); + let out = decode_string_arrays(&input).unwrap(); + assert!( + Arc::ptr_eq(&input, &out), + "valid input must be returned unchanged" + ); + } + + #[test] + fn invalid_bytes_decode_to_replacement() { + // element 0 = [0xFF, 0x41] -> "\u{FFFD}A", element 1 = [0x42] -> "B" + let input = utf8_unchecked(&[0, 2, 3], &[0xFF, 0x41, 0x42], 2); + let out = decode_string_arrays(&input).unwrap(); + let s = out.as_any().downcast_ref::().unwrap(); + assert_eq!(s.value(0), "\u{FFFD}A"); + assert_eq!(s.value(1), "B"); + } + + #[test] + fn split_codepoint_boundary_is_rebuilt() { + // whole buffer "é" (C3 A9) is valid UTF-8, but offsets split it into two invalid slices. + let input = utf8_unchecked(&[0, 1, 2], &[0xC3, 0xA9], 2); + let out = decode_string_arrays(&input).unwrap(); + let s = out.as_any().downcast_ref::().unwrap(); + assert_eq!(s.value(0), "\u{FFFD}"); + assert_eq!(s.value(1), "\u{FFFD}"); + } + + #[test] + fn nulls_are_preserved() { + let input: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), None, Some("b")])); + let out = decode_string_arrays(&input).unwrap(); + let s = out.as_any().downcast_ref::().unwrap(); + assert!(s.is_null(1)); + assert_eq!(s.value(0), "a"); + } + + #[test] + fn large_utf8_invalid_decodes() { + let data = unsafe { + ArrayData::builder(DataType::LargeUtf8) + .len(1) + .add_buffer(Buffer::from_slice_ref([0i64, 1])) + .add_buffer(Buffer::from(vec![0xFFu8])) + .build_unchecked() + }; + let out = decode_string_arrays(&make_array(data)).unwrap(); + let s = out.as_any().downcast_ref::().unwrap(); + assert_eq!(s.value(0), "\u{FFFD}"); + } + + use arrow::array::{ + DictionaryArray, FixedSizeListArray, Int32Array, ListArray, MapArray, StructArray, + }; + use arrow::datatypes::{Field, Fields, Int32Type}; + + /// An invalid Utf8 leaf ["\u{FFFD}"] built from raw bytes. + fn invalid_leaf() -> ArrayRef { + utf8_unchecked(&[0, 1], &[0xFF], 1) + } + + #[test] + fn dictionary_values_are_decoded() { + let values = invalid_leaf(); + let keys = Int32Array::from(vec![0, 0]); + let dict: ArrayRef = Arc::new(DictionaryArray::::new(keys, values)); + let out = decode_string_arrays(&dict).unwrap(); + let d = out + .as_any() + .downcast_ref::>() + .unwrap(); + let vals = d.values().as_any().downcast_ref::().unwrap(); + assert_eq!(vals.value(0), "\u{FFFD}"); + } + + #[test] + fn struct_field_is_decoded() { + let field = Arc::new(Field::new("s", DataType::Utf8, true)); + let input: ArrayRef = Arc::new(StructArray::new( + Fields::from(vec![field]), + vec![invalid_leaf()], + None, + )); + let out = decode_string_arrays(&input).unwrap(); + let s = out.as_any().downcast_ref::().unwrap(); + let col = s.column(0).as_any().downcast_ref::().unwrap(); + assert_eq!(col.value(0), "\u{FFFD}"); + } + + #[test] + fn list_values_are_decoded() { + let field = Arc::new(Field::new("item", DataType::Utf8, true)); + let offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 1].into()); + let input: ArrayRef = + Arc::new(ListArray::try_new(field, offsets, invalid_leaf(), None).unwrap()); + let out = decode_string_arrays(&input).unwrap(); + let l = out.as_any().downcast_ref::().unwrap(); + let vals = l.values().as_any().downcast_ref::().unwrap(); + assert_eq!(vals.value(0), "\u{FFFD}"); + } + + #[test] + fn valid_struct_is_zero_copy() { + let field = Arc::new(Field::new("s", DataType::Utf8, true)); + let leaf: ArrayRef = Arc::new(StringArray::from(vec!["ok"])); + let input: ArrayRef = Arc::new(StructArray::new( + Fields::from(vec![field]), + vec![leaf], + None, + )); + let out = decode_string_arrays(&input).unwrap(); + assert!( + Arc::ptr_eq(&input, &out), + "all-valid nested input must be unchanged" + ); + } + + #[test] + fn dictionary_valid_values_are_zero_copy() { + let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let keys = Int32Array::from(vec![0, 1, 0]); + let input: ArrayRef = Arc::new(DictionaryArray::::new(keys, values)); + let out = decode_string_arrays(&input).unwrap(); + assert!( + Arc::ptr_eq(&input, &out), + "dictionary with all-valid values must be returned as the original Arc" + ); + } + + #[test] + fn fixed_size_list_values_are_decoded() { + let field = Arc::new(Field::new("item", DataType::Utf8, true)); + let input: ArrayRef = + Arc::new(FixedSizeListArray::try_new(field, 1, invalid_leaf(), None).unwrap()); + let out = decode_string_arrays(&input).unwrap(); + let l = out.as_any().downcast_ref::().unwrap(); + let vals = l.values().as_any().downcast_ref::().unwrap(); + assert_eq!(vals.value(0), "\u{FFFD}"); + } + + #[test] + fn fixed_size_list_valid_is_zero_copy() { + let field = Arc::new(Field::new("item", DataType::Utf8, true)); + let leaf: ArrayRef = Arc::new(StringArray::from(vec!["ok"])); + let input: ArrayRef = Arc::new(FixedSizeListArray::try_new(field, 1, leaf, None).unwrap()); + let out = decode_string_arrays(&input).unwrap(); + assert!( + Arc::ptr_eq(&input, &out), + "fixed-size list with all-valid values must be unchanged" + ); + } + + /// Build a Map array whose single entry maps `key` -> `values` (a Utf8 leaf array of length 1). + fn build_map(key: &str, values: ArrayRef) -> ArrayRef { + let entries_fields = Fields::from(vec![ + Arc::new(Field::new("keys", DataType::Utf8, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ]); + let keys: ArrayRef = Arc::new(StringArray::from(vec![key])); + let entries = StructArray::new(entries_fields.clone(), vec![keys, values], None); + let map_field = Arc::new(Field::new( + "entries", + DataType::Struct(entries_fields), + false, + )); + let offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 1].into()); + Arc::new(MapArray::try_new(map_field, offsets, entries, None, false).unwrap()) + } + + #[test] + fn map_values_are_decoded() { + let input = build_map("k", invalid_leaf()); + let out = decode_string_arrays(&input).unwrap(); + let m = out.as_any().downcast_ref::().unwrap(); + let values = m + .entries() + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), "\u{FFFD}"); + } + + #[test] + fn map_valid_is_zero_copy() { + let values: ArrayRef = Arc::new(StringArray::from(vec!["ok"])); + let input = build_map("k", values); + let out = decode_string_arrays(&input).unwrap(); + assert!( + Arc::ptr_eq(&input, &out), + "map with all-valid values must be unchanged" + ); + } + + #[test] + fn trailing_empty_string_offset_is_handled() { + // element 0 = "ab", element 1 = "" (its offset lands exactly at values.len()). + let input = utf8_unchecked(&[0, 2, 2], b"ab", 2); + let out = decode_string_arrays(&input).unwrap(); + assert!( + Arc::ptr_eq(&input, &out), + "valid input with a trailing empty element must be returned unchanged" + ); + let s = out.as_any().downcast_ref::().unwrap(); + assert_eq!(s.value(0), "ab"); + assert_eq!(s.value(1), ""); + } + + #[test] + fn sliced_array_is_decoded() { + // Slicing shifts offsets[0] away from 0; the fast path must still handle this correctly. + let base: ArrayRef = Arc::new(StringArray::from(vec!["aa", "bb", "cc"])); + let sliced: ArrayRef = base.slice(1, 2); + let out = decode_string_arrays(&sliced).unwrap(); + // Slicing already produces a new Arc distinct from `base`, so ptr_eq against `base` isn't + // meaningful here; what matters is that the fast path recognizes the slice as valid and + // returns it unchanged (same Arc as `sliced`) rather than rebuilding it. + assert!( + Arc::ptr_eq(&sliced, &out), + "valid sliced input must be returned unchanged" + ); + let s = out.as_any().downcast_ref::().unwrap(); + assert_eq!(s.value(0), "bb"); + assert_eq!(s.value(1), "cc"); + } +} diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d2c1a5df02..7d9b64b8d2 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -41,6 +41,7 @@ use datafusion::{ physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream}, prelude::{SessionConfig, SessionContext}, }; +use datafusion_comet_common::decode_string_arrays; use datafusion_comet_proto::spark_operator::Operator; use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl}; use datafusion_spark::function::array::array_contains::SparkArrayContains; @@ -1296,7 +1297,10 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowConvert( let array_data = from_ffi(ffi_array, &ffi_schema) .map_err(|e| CometError::Internal(format!("Failed to import array: {}", e)))?; - arrays.push(arrow::array::make_array(array_data)); + let imported = arrow::array::make_array(array_data); + arrays.push(decode_string_arrays(&imported).map_err(|e| { + CometError::Internal(format!("Failed to decode imported string array: {}", e)) + })?); } // Convert columnar to row diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs index 409d064284..9833a121bf 100644 --- a/native/core/src/execution/operators/scan.rs +++ b/native/core/src/execution/operators/scan.rs @@ -30,6 +30,7 @@ use datafusion::{ physical_expr::*, physical_plan::{ExecutionPlan, *}, }; +use datafusion_comet_common::decode_string_arrays; use futures::Stream; use itertools::Itertools; use std::{ @@ -156,7 +157,7 @@ impl ScanExec { let columns = record_batch.columns(); let mut inputs: Vec = Vec::with_capacity(columns.len()); for col in columns { - inputs.push(copy_or_unpack_array(col, &CopyMode::UnpackOrClone)?); + inputs.push(import_column(col)?); } Ok(InputBatch::new(inputs, Some(num_rows))) } @@ -164,6 +165,15 @@ impl ScanExec { } } +/// Transform one FFI-imported column for native execution: first decode any invalid UTF-8 to the +/// Spark-rendered form (arrow's `from_ffi` imports string buffers unchecked), then copy/unpack. +/// Decoding runs before unpacking so a `Dictionary(_, Utf8)` decodes its compact values, not the +/// expanded ones. +fn import_column(col: &ArrayRef) -> Result { + let decoded = decode_string_arrays(col)?; + Ok(copy_or_unpack_array(&decoded, &CopyMode::UnpackOrClone)?) +} + fn schema_from_data_types(data_types: &[DataType]) -> SchemaRef { let fields = data_types .iter() @@ -372,3 +382,28 @@ impl InputBatch { InputBatch::Batch(columns, num_rows) } } + +#[cfg(test)] +mod import_tests { + use super::*; + use arrow::array::{make_array, Array, ArrayData, ArrayRef, StringArray}; + use arrow::buffer::Buffer; + use arrow::datatypes::DataType; + + #[test] + fn import_decodes_invalid_utf8_column() { + // An FFI-imported Utf8 column with invalid bytes [0xFF] -> must be decoded, not passed + // through, so downstream `value()` is sound. + let data = unsafe { + ArrayData::builder(DataType::Utf8) + .len(1) + .add_buffer(Buffer::from_slice_ref([0i32, 1])) + .add_buffer(Buffer::from(vec![0xFFu8])) + .build_unchecked() + }; + let col: ArrayRef = make_array(data); + let out = import_column(&col).unwrap(); + let s = out.as_any().downcast_ref::().unwrap(); + assert_eq!(s.value(0), "\u{FFFD}"); + } +} diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index b0c45353cd..4501d459b5 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -147,3 +147,7 @@ harness = false [[bench]] name = "to_json" harness = false + +[[bench]] +name = "utf8_decode" +harness = false diff --git a/native/spark-expr/benches/utf8_decode.rs b/native/spark-expr/benches/utf8_decode.rs new file mode 100644 index 0000000000..e4026abeff --- /dev/null +++ b/native/spark-expr/benches/utf8_decode.rs @@ -0,0 +1,40 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, StringArray}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion_comet_common::decode_string_arrays; +use std::hint::black_box; +use std::sync::Arc; + +fn bench_valid_fast_path(c: &mut Criterion) { + // 8192-row batch of valid mixed ASCII/multibyte strings, exercising the fast path (from_utf8 + // plus boundary check, zero-copy return) that every scanned string column pays. + let values: Vec = (0..8192).map(|i| format!("row-{i}-café-日本語")).collect(); + let refs: Vec<&str> = values.iter().map(|s| s.as_str()).collect(); + let input: ArrayRef = Arc::new(StringArray::from(refs)); + + c.bench_function("decode_string_arrays/valid_8192", |b| { + b.iter(|| { + let out = black_box(decode_string_arrays(black_box(&input)).unwrap()); + assert!(Arc::ptr_eq(&input, &out)); + }) + }); +} + +criterion_group!(benches, bench_valid_fast_path); +criterion_main!(benches); diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs index 0ca603ac9b..1b2dfec60e 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -28,6 +28,7 @@ use datafusion::common::Result as DFResult; use datafusion::logical_expr::ColumnarValue; use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_common::decode_string_arrays; use datafusion_comet_jni_bridge::errors::{CometError, ExecutionError}; use datafusion_comet_jni_bridge::JVMClasses; use jni::objects::{Global, JObject, JValue}; @@ -238,7 +239,10 @@ impl PhysicalExpr for JvmScalarUdfExpr { // exactly once when the Box drops at end of scope. let result_data = unsafe { from_ffi(*out_array, &out_schema) } .map_err(|e| CometError::Arrow { source: e })?; - Ok(ColumnarValue::Array(make_array(result_data))) + let imported = make_array(result_data); + let decoded = + decode_string_arrays(&imported).map_err(|e| CometError::Arrow { source: e })?; + Ok(ColumnarValue::Array(decoded)) } fn children(&self) -> Vec<&Arc> {