diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index 231f88c507..cb5d529ee1 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -21,7 +21,9 @@ use arrow::{ record_batch::RecordBatch, }; -use crate::parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions}; +use crate::parquet::parquet_support::{ + spark_parquet_convert_with_mapping, FieldMapping, SparkParquetOptions, +}; use datafusion::common::format::DEFAULT_CAST_OPTIONS; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::logical_expr::ColumnarValue; @@ -149,8 +151,11 @@ pub struct CometCastColumnExpr { /// Options forwarded to [`cast_column`]. cast_options: CastOptions<'static>, /// Spark parquet options for complex nested type conversions. - /// When present, enables `spark_parquet_convert` as a fallback. + /// When present, enables the nested conversion as a fallback. parquet_options: Option, + /// Which file field supplies each requested nested field, resolved once per file and + /// reused for every batch. Set together with `parquet_options`. + field_mapping: Option>, } // Manually derive `PartialEq`/`Hash` as `Arc` does not @@ -162,6 +167,7 @@ impl PartialEq for CometCastColumnExpr { && self.target_field.eq(&other.target_field) && self.cast_options.eq(&other.cast_options) && self.parquet_options.eq(&other.parquet_options) + && self.field_mapping.eq(&other.field_mapping) } } @@ -172,6 +178,7 @@ impl Hash for CometCastColumnExpr { self.target_field.hash(state); self.cast_options.hash(state); self.parquet_options.hash(state); + self.field_mapping.hash(state); } } @@ -210,12 +217,19 @@ impl CometCastColumnExpr { target_field, cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), parquet_options: None, + field_mapping: None, }) } - /// Set Spark parquet options to enable complex nested type conversions. - pub fn with_parquet_options(mut self, options: SparkParquetOptions) -> Self { + /// Enable nested type conversions with Spark parquet options and the field mapping + /// resolved for this expression's physical and target types. + pub fn with_parquet_options( + mut self, + options: SparkParquetOptions, + field_mapping: Arc, + ) -> Self { self.parquet_options = Some(options); + self.field_mapping = Some(field_mapping); self } } @@ -254,12 +268,22 @@ impl PhysicalExpr for CometCastColumnExpr { let input_physical_field = self.input_physical_field.data_type(); let target_field = self.target_field.data_type(); + // Relabeling only swaps metadata, so it is right when every requested field reads + // the file field at its own position. A mapping that reorders fields (ids resolved + // to other positions) has to go through the nested conversion below. + let positional = self + .field_mapping + .as_ref() + .is_none_or(|mapping| mapping.is_positional()); + match (input_physical_field, target_field) { // Nested types that differ only in field names (e.g., List element named // "item" vs "element", or Map entries named "key_value" vs "entries"). // Re-label the array so the DataType metadata matches the logical schema. (physical, logical) - if physical != logical && types_differ_only_in_field_names(physical, logical) => + if positional + && physical != logical + && types_differ_only_in_field_names(physical, logical) => { match value { ColumnarValue::Array(array) => { @@ -269,16 +293,17 @@ impl PhysicalExpr for CometCastColumnExpr { other => Ok(other), } } - // Fallback: use spark_parquet_convert for complex nested type conversions - // (e.g., List → List, Map field selection, etc.) - _ => { - if let Some(parquet_options) = &self.parquet_options { - let converted = spark_parquet_convert(value, target_field, parquet_options)?; - Ok(converted) - } else { - Ok(value) - } - } + // Fallback: nested conversion through the resolved mapping + // (e.g., List -> List, Map field selection, etc.) + _ => match (&self.parquet_options, &self.field_mapping) { + (Some(parquet_options), Some(mapping)) => spark_parquet_convert_with_mapping( + value, + target_field, + mapping, + parquet_options, + ), + _ => Ok(value), + }, } } @@ -302,8 +327,8 @@ impl PhysicalExpr for CometCastColumnExpr { Arc::clone(&self.target_field), Some(self.cast_options.clone()), )?; - if let Some(opts) = &self.parquet_options { - new_expr = new_expr.with_parquet_options(opts.clone()); + if let (Some(opts), Some(mapping)) = (&self.parquet_options, &self.field_mapping) { + new_expr = new_expr.with_parquet_options(opts.clone(), Arc::clone(mapping)); } Ok(Arc::new(new_expr)) } @@ -316,12 +341,166 @@ impl PhysicalExpr for CometCastColumnExpr { #[cfg(test)] mod tests { use super::*; + use crate::parquet::parquet_support::resolve_field_mapping; use arrow::array::{ Array, Int32Array, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray, }; use arrow::datatypes::{Field, Fields}; use datafusion::physical_expr::expressions::Column; use datafusion_comet_spark_expr::EvalMode; + use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use std::collections::HashMap; + + fn int_field_with_id(name: &str, id: i32) -> Field { + Field::new(name, DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + id.to_string(), + )])) + } + + /// File struct `x` (id 1) = 42, `y` (id 2) = 43; requested struct names them the same + /// but swaps the ids. Names and types match, so only the positional gate keeps the + /// relabel shortcut from firing: the mapping reads by id and the result must be + /// `x` = 43, `y` = 42. + #[test] + fn test_swapped_field_ids_bypass_relabel_shortcut() { + let physical_fields = + Fields::from(vec![int_field_with_id("x", 1), int_field_with_id("y", 2)]); + let logical_fields = + Fields::from(vec![int_field_with_id("x", 2), int_field_with_id("y", 1)]); + + let input_field = Arc::new(Field::new( + "s", + DataType::Struct(physical_fields.clone()), + true, + )); + let target_field = Arc::new(Field::new( + "s", + DataType::Struct(logical_fields.clone()), + true, + )); + + let columns: Vec = vec![ + Arc::new(Int32Array::from(vec![42])), + Arc::new(Int32Array::from(vec![43])), + ]; + let struct_arr = StructArray::new(physical_fields, columns, None); + let schema = Schema::new(vec![Arc::clone(&input_field)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_arr)]).unwrap(); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + let mapping = Arc::new( + resolve_field_mapping(input_field.data_type(), target_field.data_type(), &opts) + .unwrap(), + ); + assert!(!mapping.is_positional()); + + let col_expr: Arc = Arc::new(Column::new("s", 0)); + let cast_expr = CometCastColumnExpr::try_new(col_expr, input_field, target_field, None) + .unwrap() + .with_parquet_options(opts, mapping); + + let ColumnarValue::Array(arr) = cast_expr.evaluate(&batch).unwrap() else { + panic!("expected array result"); + }; + assert_eq!(arr.data_type(), &DataType::Struct(logical_fields)); + let result = arr.as_any().downcast_ref::().unwrap(); + let x = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let y = result + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(x.value(0), 43); + assert_eq!(y.value(0), 42); + } + + /// Companion guard: without any field ids the relabel shortcut must keep + /// handling name-only differences, whether or not id read mode is enabled. + #[test] + fn test_relabel_shortcut_kept_for_name_only_differences_without_ids() { + // Physical: s { col: List(Field("item", Int32)) } + // Logical: s { col: List(Field("element", Int32)) } + let physical_list_field = Arc::new(Field::new("item", DataType::Int32, true)); + let logical_list_field = Arc::new(Field::new("element", DataType::Int32, true)); + let physical_fields = Fields::from(vec![Field::new( + "col", + DataType::List(Arc::clone(&physical_list_field)), + true, + )]); + let logical_fields = Fields::from(vec![Field::new( + "col", + DataType::List(logical_list_field), + true, + )]); + + let input_field = Arc::new(Field::new( + "s", + DataType::Struct(physical_fields.clone()), + true, + )); + let target_field = Arc::new(Field::new( + "s", + DataType::Struct(logical_fields.clone()), + true, + )); + + let values = Int32Array::from(vec![1, 2, 3]); + let list = ListArray::new( + physical_list_field, + arrow::buffer::OffsetBuffer::new(vec![0, 2, 3].into()), + Arc::new(values), + None, + ); + let struct_arr = StructArray::new(physical_fields, vec![Arc::new(list) as ArrayRef], None); + let schema = Schema::new(vec![Arc::clone(&input_field)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_arr)]).unwrap(); + + let col_expr: Arc = Arc::new(Column::new("s", 0)); + + // Without parquet options the fallback arm would return the value + // unchanged, so a relabeled result proves the shortcut itself fired. + let plain_expr = CometCastColumnExpr::try_new( + Arc::clone(&col_expr), + Arc::clone(&input_field), + Arc::clone(&target_field), + None, + ) + .unwrap(); + + // Enabling id read mode without any id metadata must not disable the + // shortcut either: the resolved mapping is positional. + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + let mapping = Arc::new( + resolve_field_mapping(input_field.data_type(), target_field.data_type(), &opts) + .unwrap(), + ); + assert!(mapping.is_positional()); + let id_mode_expr = CometCastColumnExpr::try_new(col_expr, input_field, target_field, None) + .unwrap() + .with_parquet_options(opts, mapping); + + for cast_expr in [plain_expr, id_mode_expr] { + let result = cast_expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(arr) = result else { + panic!("expected array result"); + }; + assert_eq!(arr.data_type(), &DataType::Struct(logical_fields.clone())); + let result_struct = arr.as_any().downcast_ref::().unwrap(); + let result_list = result_struct + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result_list.len(), 2); + } + } #[test] fn test_rejects_millisecond_logical_timestamp() { @@ -369,7 +548,10 @@ mod tests { let expr: Arc = Arc::new(Column::new("ts", 0)); let cast_expr = CometCastColumnExpr::try_new(expr, input_field, target_field, None) .unwrap() - .with_parquet_options(SparkParquetOptions::new(eval_mode, "UTC", false)); + .with_parquet_options( + SparkParquetOptions::new(eval_mode, "UTC", false), + Arc::new(FieldMapping::Leaf), + ); let input = TimestampMillisecondArray::from(vec![Some(1_234), Some(-1_234), None]) .with_timezone_opt(source_tz.clone()); diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index fe16e099b7..eb4d23430d 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -17,7 +17,7 @@ use crate::execution::operators::ExecutionError; use crate::parquet::name_fold::fold_names; -use arrow::array::{FixedSizeBinaryArray, ListArray, MapArray, StringArray}; +use arrow::array::{FixedSizeBinaryArray, LargeListArray, ListArray, MapArray, StringArray}; use arrow::buffer::NullBuffer; use arrow::compute::can_cast_types; use arrow::datatypes::{FieldRef, Fields}; @@ -147,17 +147,32 @@ impl SparkParquetOptions { /// Spark-compatible cast implementation. Defers to DataFusion's cast where that is known /// to be compatible, and returns an error when a not supported and not DF-compatible cast -/// is requested. +/// is requested. Resolves the nested field mapping for this one value; a per-file caller +/// resolves once and uses [`spark_parquet_convert_with_mapping`] for every batch. pub fn spark_parquet_convert( arg: ColumnarValue, data_type: &DataType, parquet_options: &SparkParquetOptions, +) -> DataFusionResult { + let mapping = + resolve_field_mapping(&arg.data_type(), data_type, parquet_options).map_err(spark_error)?; + spark_parquet_convert_with_mapping(arg, data_type, &mapping, parquet_options) +} + +/// [`spark_parquet_convert`] with a mapping already resolved for the value's type. +pub(crate) fn spark_parquet_convert_with_mapping( + arg: ColumnarValue, + data_type: &DataType, + mapping: &FieldMapping, + parquet_options: &SparkParquetOptions, ) -> DataFusionResult { match arg { - ColumnarValue::Array(array) => Ok(ColumnarValue::Array(parquet_convert_array( + ColumnarValue::Array(array) => Ok(ColumnarValue::Array(convert_array( array, data_type, + mapping, parquet_options, + true, )?)), ColumnarValue::Scalar(scalar) => { // Note that normally CAST(scalar) should be fold in Spark JVM side. However, for @@ -165,7 +180,7 @@ pub fn spark_parquet_convert( // here. let array = scalar.to_array()?; let scalar = ScalarValue::try_from_array( - &parquet_convert_array(array, data_type, parquet_options)?, + &convert_array(array, data_type, mapping, parquet_options, true)?, 0, )?; Ok(ColumnarValue::Scalar(scalar)) @@ -173,17 +188,230 @@ pub fn spark_parquet_convert( } } -fn parquet_convert_array( - array: ArrayRef, +/// Wrap a [`SparkError`] the way every native operator surfaces it to the JVM. +pub(crate) fn spark_error(error: SparkError) -> DataFusionError { + DataFusionError::External(Box::new(error)) +} + +/// Outcome of matching one requested id or name against a struct's file fields: the last +/// file field that matched and whether more than one did. A plain `Copy` value, so resolving +/// a wide struct allocates nothing per id or per name; the matched names are only gathered +/// when an ambiguity is reported. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct FieldMatch { + pub(crate) index: usize, + pub(crate) ambiguous: bool, +} + +impl FieldMatch { + pub(crate) fn new(index: usize, ambiguous: bool) -> Self { + Self { index, ambiguous } + } + + /// The first file field carrying this id or name. + pub(crate) fn first(index: usize) -> Self { + Self::new(index, false) + } + + /// A further file field carrying the same id or name: the later index wins, as Spark's + /// `toMap` does for exact names, and the entry turns ambiguous. + pub(crate) fn also(self, index: usize) -> Self { + Self::new(index, true) + } +} + +/// Record file field `index` under `key`, keeping the entry `Copy`-sized however many fields +/// share the key. +pub(crate) fn record_field_match( + matches: &mut HashMap, + key: K, + index: usize, +) { + matches + .entry(key) + .and_modify(|m| *m = m.also(index)) + .or_insert_with(|| FieldMatch::first(index)); +} + +/// Comma-joined names of the fields carrying `id`, for the duplicate-id error message. +pub(crate) fn field_names_with_id(fields: &Fields, id: i32) -> String { + fields + .iter() + .filter(|f| field_id(f) == Some(id)) + .map(|f| f.name().as_str()) + .collect::>() + .join(", ") +} + +/// Which file field supplies each requested field, resolved once per file and reused for +/// every batch. Follows the requested type as Spark's `clipParquetSchema` does: a struct +/// lists one source per requested field, a list (large or not) or map carries the mapping +/// of its element or key and value types, and anything else is a leaf converted by type. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) enum FieldMapping { + Struct(Vec), + List(Box), + Map(Box, Box), + Leaf, +} + +/// The file field behind one requested struct field. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct StructFieldSource { + /// Index of the file field supplying the requested field; `None` null-fills it. + pub(crate) from_index: Option, + /// Mapping of the requested field's own type. + pub(crate) nested: FieldMapping, +} + +impl FieldMapping { + /// True when every requested field reads the file field at its own position, so a + /// metadata-only relabel of the file array already yields the requested layout. + pub(crate) fn is_positional(&self) -> bool { + match self { + FieldMapping::Struct(sources) => sources + .iter() + .enumerate() + .all(|(i, s)| s.from_index == Some(i) && s.nested.is_positional()), + FieldMapping::List(inner) => inner.is_positional(), + FieldMapping::Map(key, value) => key.is_positional() && value.is_positional(), + FieldMapping::Leaf => true, + } + } +} + +/// Resolve how `to_type` reads from `from_type`, recursing through struct, list, and map +/// types. Raises the ambiguity Spark reports from `clipParquetGroupFields` when a requested +/// id or case-insensitive name matches more than one file field at any level. +pub(crate) fn resolve_field_mapping( + from_type: &DataType, to_type: &DataType, parquet_options: &SparkParquetOptions, -) -> DataFusionResult { - parquet_convert_array_impl(array, to_type, parquet_options, true) +) -> Result { + use DataType::*; + match (from_type, to_type) { + (Struct(from_fields), Struct(to_fields)) => { + resolve_struct_mapping(from_fields, to_fields, parquet_options) + } + (List(from_item), List(to_item)) | (LargeList(from_item), LargeList(to_item)) => { + Ok(FieldMapping::List(Box::new(resolve_field_mapping( + from_item.data_type(), + to_item.data_type(), + parquet_options, + )?))) + } + (Map(from_entries, from_ordered), Map(to_entries, to_ordered)) + if from_ordered == to_ordered => + { + match (from_entries.data_type(), to_entries.data_type()) { + (Struct(from_kv), Struct(to_kv)) if from_kv.len() == 2 && to_kv.len() == 2 => { + let key = resolve_field_mapping( + from_kv[0].data_type(), + to_kv[0].data_type(), + parquet_options, + )?; + let value = resolve_field_mapping( + from_kv[1].data_type(), + to_kv[1].data_type(), + parquet_options, + )?; + Ok(FieldMapping::Map(Box::new(key), Box::new(value))) + } + _ => Ok(FieldMapping::Leaf), + } + } + _ => Ok(FieldMapping::Leaf), + } } -fn parquet_convert_array_impl( +/// Match `to` (requested) struct fields to `from` (file) fields. Mirrors Spark's +/// `clipParquetGroupFields`: when the requested struct carries Parquet field ids anywhere, +/// id-bearing requested fields match only by id and the rest by name; otherwise every field +/// matches by name. +fn resolve_struct_mapping( + from_fields: &Fields, + to_fields: &Fields, + parquet_options: &SparkParquetOptions, +) -> Result { + let should_match_by_id = + parquet_options.use_field_id && to_fields.iter().any(|f| field_id(f).is_some()); + + let mut id_matches: HashMap = HashMap::new(); + if should_match_by_id { + for (i, field) in from_fields.iter().enumerate() { + if let Some(id) = field_id(field) { + record_field_match(&mut id_matches, id, i); + } + } + } + + // Fold the file and requested names once via the same `toLowerCase(Locale.ROOT)` the + // top-level schema adapter uses, so nested case-insensitive matching agrees with it. + let mut all_names: Vec<&str> = Vec::with_capacity(from_fields.len() + to_fields.len()); + all_names.extend(from_fields.iter().map(|f| f.name().as_str())); + all_names.extend(to_fields.iter().map(|f| f.name().as_str())); + let all_folded = fold_names(&all_names, parquet_options.case_sensitive); + let (from_folded, to_folded) = all_folded.split_at(from_fields.len()); + + let mut name_matches: HashMap<&str, FieldMatch> = HashMap::new(); + for (i, folded) in from_folded.iter().enumerate() { + record_field_match(&mut name_matches, folded.as_str(), i); + } + + let mut sources = Vec::with_capacity(to_fields.len()); + for (to_pos, to_field) in to_fields.iter().enumerate() { + let from_index = match (should_match_by_id, field_id(to_field)) { + // A missing id match is a missing column, never a name match. + (true, Some(id)) => match id_matches.get(&id) { + Some(m) if m.ambiguous => { + return Err(SparkError::DuplicateFieldByFieldId { + required_id: id, + matched_fields: field_names_with_id(from_fields, id), + }); + } + Some(m) => Some(m.index), + None => None, + }, + _ => match name_matches.get(to_folded[to_pos].as_str()) { + // Spark's `caseInsensitiveParquetFieldMap` rejects a requested name that folds + // onto more than one file field. In case-sensitive mode the fold is identity, so + // a collision means byte-identical siblings and the later one wins silently, + // as with Spark's `caseSensitiveParquetFieldMap` built by `toMap`. + Some(m) if m.ambiguous && !parquet_options.case_sensitive => { + let matched: Vec<&str> = from_folded + .iter() + .zip(from_fields.iter()) + .filter(|(folded, _)| *folded == &to_folded[to_pos]) + .map(|(_, f)| f.name().as_str()) + .collect(); + return Err(SparkError::duplicate_field_case_insensitive( + to_field.name(), + &matched, + )); + } + Some(m) => Some(m.index), + None => None, + }, + }; + let nested = match from_index { + Some(i) => resolve_field_mapping( + from_fields[i].data_type(), + to_field.data_type(), + parquet_options, + )?, + None => FieldMapping::Leaf, + }; + sources.push(StructFieldSource { from_index, nested }); + } + Ok(FieldMapping::Struct(sources)) +} + +/// Convert `array` to `to_type` through its resolved `mapping`. `top_level` is true only for +/// the column itself, never for a struct field, list element, or map entry beneath it. +fn convert_array( array: ArrayRef, to_type: &DataType, + mapping: &FieldMapping, parquet_options: &SparkParquetOptions, top_level: bool, ) -> DataFusionResult { @@ -192,18 +420,16 @@ fn parquet_convert_array_impl( // Try Comet specific handlers first, then arrow-rs cast if supported, // return uncasted data otherwise - match (from_type, to_type) { - (Struct(_), Struct(_)) => Ok(parquet_convert_struct_to_struct( - array.as_struct(), - from_type, - to_type, - parquet_options, - )?), - (List(_), List(to_inner_type)) => { + match (from_type, to_type, mapping) { + (Struct(_), Struct(to_fields), FieldMapping::Struct(sources)) => { + convert_struct(array.as_struct(), to_fields, sources, parquet_options) + } + (List(_), List(to_inner_type), FieldMapping::List(inner)) => { let list_arr: &ListArray = array.as_list(); - let cast_field = parquet_convert_array_impl( + let cast_field = convert_array( Arc::clone(list_arr.values()), to_inner_type.data_type(), + inner, parquet_options, false, )?; @@ -215,10 +441,26 @@ fn parquet_convert_array_impl( list_arr.nulls().cloned(), ))) } - ( - Timestamp(TimeUnit::Millisecond, _), - Timestamp(TimeUnit::Microsecond, target_tz), - ) if top_level && parquet_options.checked_timestamp_overflow => { + (LargeList(_), LargeList(to_inner_type), FieldMapping::List(inner)) => { + let list_arr: &LargeListArray = array.as_list(); + let cast_field = convert_array( + Arc::clone(list_arr.values()), + to_inner_type.data_type(), + inner, + parquet_options, + false, + )?; + + Ok(Arc::new(LargeListArray::new( + Arc::clone(to_inner_type), + list_arr.offsets().clone(), + cast_field, + list_arr.nulls().cloned(), + ))) + } + (Timestamp(TimeUnit::Millisecond, _), Timestamp(TimeUnit::Microsecond, target_tz), _) + if top_level && parquet_options.checked_timestamp_overflow => + { // Spark's Parquet reader calls the checked `millisToMicros` conversion for both // direct and dictionary values, independent of CAST evaluation mode: // https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorUpdaterFactory.java#L817-L833 @@ -237,7 +479,7 @@ fn parquet_convert_array_impl( .with_timezone_opt(target_tz.clone()); Ok(Arc::new(micros)) } - (Timestamp(TimeUnit::Microsecond, None), Timestamp(TimeUnit::Microsecond, Some(tz))) => { + (Timestamp(TimeUnit::Microsecond, None), Timestamp(TimeUnit::Microsecond, Some(tz)), _) => { Ok(Arc::new( array .as_primitive::() @@ -245,12 +487,21 @@ fn parquet_convert_array_impl( .with_timezone(Arc::clone(tz)), )) } - (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) - , + (Map(_, ordered_from), Map(_, ordered_to), FieldMapping::Map(key, value)) + if ordered_from == ordered_to => + { + parquet_convert_map_to_map( + array.as_map(), + to_type, + key, + value, + parquet_options, + *ordered_to, + ) + } // 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. - (FixedSizeBinary(16), Utf8) => { + (FixedSizeBinary(16), Utf8, _) => { let binary_array = array .as_any() .downcast_ref::() @@ -260,9 +511,8 @@ fn parquet_convert_array_impl( .iter() .map(|opt_bytes| { opt_bytes.map(|bytes| { - let uuid = uuid::Uuid::from_bytes( - bytes.try_into().expect("Expected 16 bytes") - ); + let uuid = + uuid::Uuid::from_bytes(bytes.try_into().expect("Expected 16 bytes")); uuid.to_string() }) }) @@ -279,135 +529,74 @@ fn parquet_convert_array_impl( } /// Read the Parquet field id stored under arrow-rs's `PARQUET_FIELD_ID_META_KEY`. -fn field_id(field: &arrow::datatypes::Field) -> Option { +pub(crate) fn field_id(field: &arrow::datatypes::Field) -> Option { field .metadata() .get(PARQUET_FIELD_ID_META_KEY) .and_then(|v| v.parse::().ok()) } -/// Cast between struct types based on logic in +/// Build the requested struct from the file struct, reading each requested field from the +/// file field at its resolved source index. Based on /// `org.apache.spark.sql.catalyst.expressions.Cast#castStruct`. -fn parquet_convert_struct_to_struct( +fn convert_struct( array: &StructArray, - from_type: &DataType, - to_type: &DataType, + to_fields: &Fields, + sources: &[StructFieldSource], parquet_options: &SparkParquetOptions, ) -> DataFusionResult { - match (from_type, to_type) { - (DataType::Struct(from_fields), DataType::Struct(to_fields)) => { - // Match `from` (file) fields to `to` (logical) fields. Mirrors Spark's - // `clipParquetGroupFields`: when the logical struct carries Parquet field IDs - // anywhere, ID-bearing logical fields match ONLY by ID; non-ID-bearing fields - // fall back to name match. When no logical field carries an ID, fall back to - // name match across the board. - let should_match_by_id = - parquet_options.use_field_id && to_fields.iter().any(|f| field_id(f).is_some()); - - let from_id_to_index: HashMap = if should_match_by_id { - let mut map = HashMap::new(); - for (i, field) in from_fields.iter().enumerate() { - if let Some(id) = field_id(field) { - map.entry(id).or_insert(i); - } - } - map - } else { - HashMap::new() - }; - - // Fold the file (`from`) and requested (`to`) field names once via the JVM's - // `toLowerCase(Locale.ROOT)` (the same fold the top-level schema adapter uses), so - // nested case-insensitive matching is byte-for-byte consistent with the top level. - let mut all_names: Vec<&str> = Vec::with_capacity(from_fields.len() + to_fields.len()); - all_names.extend(from_fields.iter().map(|f| f.name().as_str())); - all_names.extend(to_fields.iter().map(|f| f.name().as_str())); - let all_folded = fold_names(&all_names, parquet_options.case_sensitive); - let (from_folded, to_folded) = all_folded.split_at(from_fields.len()); - - // Group file field indices by folded name so a case-insensitive collision is detected - // (Spark's `caseInsensitiveParquetFieldMap`) rather than silently overwritten. - let mut folded_to_indices: HashMap<&str, Vec> = HashMap::new(); - for (i, folded) in from_folded.iter().enumerate() { - folded_to_indices - .entry(folded.as_str()) - .or_default() - .push(i); - } + if sources.len() != to_fields.len() { + return Err(DataFusionError::Internal(format!( + "struct field mapping has {} sources for {} requested fields", + sources.len(), + to_fields.len() + ))); + } - let mut field_overlap = false; - let mut cast_fields: Vec = Vec::with_capacity(to_fields.len()); - for (to_pos, to_field) in to_fields.iter().enumerate() { - let from_index = match (should_match_by_id, field_id(to_field)) { - // Spark treats a missing ID match as a missing column rather than - // falling back to name match. - (true, Some(id)) => from_id_to_index.get(&id).copied(), - _ => match folded_to_indices.get(to_folded[to_pos].as_str()) { - // Mirror Spark's `foundDuplicateFieldInCaseInsensitiveModeError`: a - // requested field matching more than one file field is ambiguous. Gated on - // case-insensitive mode to match the top-level check (which only runs when - // `!case_sensitive`): when case-sensitive the fold is identity, so a - // collision means byte-identical sibling names, and raising an error whose - // message says "in case-insensitive mode" would be wrong. Fall through to - // the first match in that case. - Some(indices) if indices.len() > 1 && !parquet_options.case_sensitive => { - let matched: Vec<&str> = indices - .iter() - .map(|&i| from_fields[i].name().as_str()) - .collect(); - return Err(DataFusionError::External(Box::new( - SparkError::duplicate_field_case_insensitive( - to_field.name(), - &matched, - ), - ))); - } - Some(indices) => Some(indices[0]), - None => None, - }, - }; - - if let Some(from_index) = from_index { - cast_fields.push(parquet_convert_array_impl( - Arc::clone(array.column(from_index)), - to_field.data_type(), - parquet_options, - false, - )?); - field_overlap = true; - } else { - cast_fields.push(new_null_array(to_field.data_type(), array.len())); - } + let mut field_overlap = false; + let mut cast_fields: Vec = Vec::with_capacity(to_fields.len()); + for (to_field, source) in to_fields.iter().zip(sources) { + match source.from_index { + Some(from_index) => { + cast_fields.push(convert_array( + Arc::clone(array.column(from_index)), + to_field.data_type(), + &source.nested, + parquet_options, + false, + )?); + field_overlap = true; } - - // When the file's struct contains none of the requested fields, the - // returned validity buffer depends on Spark's - // `spark.sql.legacy.parquet.returnNullStructIfAllFieldsMissing` (SPARK-53535, - // Spark 4.1+). Legacy mode marks the whole column null; the new default - // preserves the file's parent-row nullness so non-null parents materialize - // as a struct of all-null fields. - let nulls = - if !field_overlap && parquet_options.return_null_struct_if_all_fields_missing { - Some(NullBuffer::new_null(array.len())) - } else { - array.nulls().cloned() - }; - - Ok(Arc::new(StructArray::new( - to_fields.clone(), - cast_fields, - nulls, - ))) + None => cast_fields.push(new_null_array(to_field.data_type(), array.len())), } - _ => unreachable!(), } + + // When the file's struct contains none of the requested fields, the + // returned validity buffer depends on Spark's + // `spark.sql.legacy.parquet.returnNullStructIfAllFieldsMissing` (SPARK-53535, + // Spark 4.1+). Legacy mode marks the whole column null; the new default + // preserves the file's parent-row nullness so non-null parents materialize + // as a struct of all-null fields. + let nulls = if !field_overlap && parquet_options.return_null_struct_if_all_fields_missing { + Some(NullBuffer::new_null(array.len())) + } else { + array.nulls().cloned() + }; + + Ok(Arc::new(StructArray::new( + to_fields.clone(), + cast_fields, + nulls, + ))) } /// Cast a map type to another map type. The same as arrow-cast except we recursively call our own -/// parquet_convert_array +/// convert_array with the resolved key and value mappings. fn parquet_convert_map_to_map( from: &MapArray, to_data_type: &DataType, + key_mapping: &FieldMapping, + value_mapping: &FieldMapping, parquet_options: &SparkParquetOptions, to_ordered: bool, ) -> Result { @@ -420,15 +609,17 @@ fn parquet_convert_map_to_map( "map is missing value field".to_string(), ))?; - let key_array = parquet_convert_array_impl( + let key_array = convert_array( Arc::clone(from.keys()), key_field.data_type(), + key_mapping, parquet_options, false, )?; - let value_array = parquet_convert_array_impl( + let value_array = convert_array( Arc::clone(from.values()), value_field.data_type(), + value_mapping, parquet_options, false, )?; @@ -729,9 +920,23 @@ mod tests { } } + /// Convert one array through the public entry point, resolving its mapping. + fn parquet_convert_array( + array: arrow::array::ArrayRef, + to_type: &arrow::datatypes::DataType, + parquet_options: &crate::parquet::parquet_support::SparkParquetOptions, + ) -> datafusion::common::Result { + use crate::parquet::parquet_support::spark_parquet_convert; + use datafusion::physical_plan::ColumnarValue; + match spark_parquet_convert(ColumnarValue::Array(array), to_type, parquet_options)? { + ColumnarValue::Array(array) => Ok(array), + ColumnarValue::Scalar(_) => unreachable!("array input yields an array"), + } + } + #[test] fn test_millis_to_micros_overflow_checked_only_at_top_level() { - use crate::parquet::parquet_support::{parquet_convert_array, SparkParquetOptions}; + use crate::parquet::parquet_support::SparkParquetOptions; use arrow::array::{Array, ArrayRef, StructArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType, Field, Fields, TimeUnit}; use datafusion_comet_spark_expr::EvalMode; @@ -795,4 +1000,251 @@ mod tests { assert!(converted_child.is_null(0), "overflow must become NULL"); assert!(converted_child.is_null(1)); } + + mod struct_field_matching { + use super::parquet_convert_array; + use crate::parquet::parquet_support::{ + resolve_field_mapping, FieldMapping, FieldMatch, SparkParquetOptions, + }; + use arrow::array::{Array, ArrayRef, Int32Array, LargeListArray, StructArray}; + use arrow::datatypes::{DataType, Field, Fields}; + use datafusion_comet_spark_expr::EvalMode; + + /// The per-id lookup entry is a plain `Copy` value: the second field sharing an id + /// only flips the ambiguity flag, so resolving a wide struct allocates no vector + /// per id. + #[test] + fn field_match_records_ambiguity_without_allocating() { + fn assert_copy() {} + assert_copy::(); + + let first = FieldMatch::first(3); + assert_eq!(first, FieldMatch::new(3, false)); + let again = first.also(5); + assert_eq!(again, FieldMatch::new(5, true)); + assert!(again.ambiguous); + } + + /// Every requested id resolves to exactly one file field: the resolved mapping is + /// positional and carries one source per requested field. + #[test] + fn resolve_mapping_by_id_is_positional_for_unique_ids() { + let fields: Vec = (0..256) + .map(|i| field_with_id(&format!("c{i}"), i)) + .collect(); + let from_type = DataType::Struct(Fields::from(fields.clone())); + let to_type = DataType::Struct(Fields::from(fields)); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + + let mapping = resolve_field_mapping(&from_type, &to_type, &opts).unwrap(); + assert!(mapping.is_positional()); + let FieldMapping::Struct(sources) = &mapping else { + panic!("expected a struct mapping"); + }; + assert_eq!(sources.len(), 256); + assert!(sources + .iter() + .enumerate() + .all(|(i, s)| s.from_index == Some(i))); + } + + /// Requested ids in a different order than the file resolve by id, so the mapping + /// is not positional and a metadata-only relabel would read the wrong columns. + #[test] + fn resolve_mapping_by_id_reorders_swapped_ids() { + let from_type = DataType::Struct(Fields::from(vec![ + field_with_id("x", 1), + field_with_id("y", 2), + ])); + let to_type = DataType::Struct(Fields::from(vec![ + field_with_id("x", 2), + field_with_id("y", 1), + ])); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + + let mapping = resolve_field_mapping(&from_type, &to_type, &opts).unwrap(); + assert!(!mapping.is_positional()); + let FieldMapping::Struct(sources) = &mapping else { + panic!("expected a struct mapping"); + }; + assert_eq!(sources[0].from_index, Some(1)); + assert_eq!(sources[1].from_index, Some(0)); + } + + /// A large list element resolves like a list element: swapped ids inside it make the + /// mapping non-positional and the conversion reads each field by id. + #[test] + fn resolve_mapping_recurses_into_large_list_element() { + let from_elem = Fields::from(vec![field_with_id("x", 1), field_with_id("y", 2)]); + let to_elem = Fields::from(vec![field_with_id("x", 2), field_with_id("y", 1)]); + let from_field = Arc::new(Field::new("item", DataType::Struct(from_elem), true)); + let to_field = Arc::new(Field::new("item", DataType::Struct(to_elem), true)); + let from_type = DataType::LargeList(Arc::clone(&from_field)); + let to_type = DataType::LargeList(to_field); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + + let mapping = resolve_field_mapping(&from_type, &to_type, &opts).unwrap(); + assert!(!mapping.is_positional()); + + let element = struct_of( + vec![field_with_id("x", 1), field_with_id("y", 2)], + vec![42, 43], + ); + let list = LargeListArray::new( + from_field, + arrow::buffer::OffsetBuffer::new(vec![0i64, 1].into()), + element, + None, + ); + let result = parquet_convert_array(Arc::new(list), &to_type, &opts).unwrap(); + assert_eq!(result.data_type(), &to_type); + let values = result + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let x = values + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let y = values + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(x.value(0), 43); + assert_eq!(y.value(0), 42); + } + + /// A duplicated requested id nested under a list element is rejected at resolution + /// time, mirroring Spark's `clipParquetListType` recursing into `matchIdField`. + #[test] + fn resolve_mapping_rejects_duplicate_id_inside_list_element() { + let from_elem = DataType::Struct(Fields::from(vec![ + field_with_id("x", 1), + field_with_id("y", 1), + ])); + let to_elem = DataType::Struct(Fields::from(vec![field_with_id("x", 1)])); + let from_type = DataType::List(Arc::new(Field::new("item", from_elem, true))); + let to_type = DataType::List(Arc::new(Field::new("element", to_elem, true))); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + + let err = resolve_field_mapping(&from_type, &to_type, &opts).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("_LEGACY_ERROR_TEMP_2094") && msg.contains("[x, y]"), + "unexpected error: {msg}" + ); + } + use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use std::collections::HashMap; + use std::sync::Arc; + + fn field_with_id(name: &str, id: i32) -> Field { + Field::new(name, DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + id.to_string(), + )])) + } + + fn struct_of(fields: Vec, values: Vec) -> ArrayRef { + let arrays: Vec = values + .into_iter() + .map(|v| Arc::new(Int32Array::from(vec![Some(v)])) as ArrayRef) + .collect(); + Arc::new(StructArray::new(Fields::from(fields), arrays, None)) + } + + /// Two physical struct fields share field ID 1 and the logical struct requests that + /// ID: Spark's `matchIdField` raises `foundDuplicateFieldInFieldIdLookupModeError` + /// (`_LEGACY_ERROR_TEMP_2094`) rather than silently reading the first match. + #[test] + fn requested_duplicate_field_id_errors() { + let from = struct_of( + vec![field_with_id("x", 1), field_with_id("y", 1)], + vec![42, 43], + ); + let to_type = DataType::Struct(Fields::from(vec![field_with_id("f", 1)])); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + + let err = parquet_convert_array(from, &to_type, &opts).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("_LEGACY_ERROR_TEMP_2094") && msg.contains("id=1"), + "unexpected error: {msg}" + ); + } + + /// Companion to `requested_duplicate_field_id_errors`: a duplicated file ID that no + /// requested field looks up must stay harmless (Spark only raises inside + /// `matchIdField`, i.e. for requested IDs). + #[test] + fn unrequested_duplicate_field_id_reads_fine() { + let from = struct_of( + vec![ + field_with_id("x", 1), + field_with_id("y", 1), + field_with_id("z", 2), + ], + vec![42, 43, 44], + ); + let to_type = DataType::Struct(Fields::from(vec![field_with_id("f", 2)])); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + + let result = parquet_convert_array(from, &to_type, &opts).unwrap(); + let result_struct = result.as_any().downcast_ref::().unwrap(); + let col = result_struct + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.value(0), 44); + } + + /// Two physical struct fields carry the IDENTICAL name in case-sensitive mode. + /// Spark's `caseSensitiveParquetFieldMap` is built with `.toMap`, where the later + /// entry wins silently; the exact-name lookup here must do the same rather than + /// return the first field. + #[test] + fn duplicate_exact_names_resolve_to_the_last_field() { + let from = struct_of( + vec![ + Field::new("d", DataType::Int32, true), + Field::new("d", DataType::Int32, true), + ], + vec![1, 2], + ); + let to_type = + DataType::Struct(Fields::from(vec![Field::new("d", DataType::Int32, true)])); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.case_sensitive = true; + + let result = parquet_convert_array(from, &to_type, &opts).unwrap(); + let result_struct = result.as_any().downcast_ref::().unwrap(); + let col = result_struct + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.value(0), 2); + } + } } diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index 8b36261822..562c527207 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -17,7 +17,10 @@ use crate::parquet::cast_column::CometCastColumnExpr; use crate::parquet::name_fold::{fold_name, fold_names, fold_schema_names}; -use crate::parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions}; +use crate::parquet::parquet_support::{ + field_names_with_id, record_field_match, resolve_field_mapping, spark_error, + spark_parquet_convert, FieldMapping, FieldMatch, SparkParquetOptions, +}; use arrow::array::new_empty_array; use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; @@ -102,30 +105,23 @@ fn remap_physical_schema( ))); } - // Build id -> all matching physical field names. We need the full list so we can mirror - // Spark's `_LEGACY_ERROR_TEMP_2094` "Found duplicate field(s)" error when an ID-bearing - // logical field would resolve to more than one physical field. - let mut id_to_phys_names: HashMap> = HashMap::new(); + // Index every physical field id once. Spark's `matchIdField` raises + // `_LEGACY_ERROR_TEMP_2094` "Found duplicate field(s)" when an ID-bearing logical field + // resolves to more than one physical field; the matched names are only gathered then. if should_match_by_id { - for pf in physical_schema.fields() { + let mut id_matches: HashMap = HashMap::new(); + for (i, pf) in physical_schema.fields().iter().enumerate() { if let Some(id) = parse_field_id(pf) { - id_to_phys_names - .entry(id) - .or_default() - .push(pf.name().clone()); + record_field_match(&mut id_matches, id, i); } } for lf in logical_schema.fields() { if let Some(id) = parse_field_id(lf) { - if let Some(matches) = id_to_phys_names.get(&id) { - if matches.len() > 1 { - return Err(DataFusionError::External(Box::new( - SparkError::DuplicateFieldByFieldId { - required_id: id, - matched_fields: matches.join(", "), - }, - ))); - } + if id_matches.get(&id).is_some_and(|m| m.ambiguous) { + return Err(spark_error(SparkError::DuplicateFieldByFieldId { + required_id: id, + matched_fields: field_names_with_id(physical_schema.fields(), id), + })); } } } @@ -151,30 +147,47 @@ fn remap_physical_schema( let logical_folded = fold_schema_names(logical_schema, case_sensitive); let physical_folded = fold_schema_names(physical_schema, case_sensitive); - // Folded names of ID-bearing logical fields whose ID is not present in the file. Any physical - // field that shares one of these names must be renamed to something the - // `DefaultPhysicalExprAdapter` cannot name-match, otherwise the read would silently fall - // through to a name match. Spark's `matchIdField` solves the same problem with - // `generateFakeColumnName` (see `ParquetReadSupport.scala`). - let unmatched_id_logical_folded: HashSet = if should_match_by_id { + // Folded names of ID-bearing logical fields. Spark's `matchIdField` resolves these + // strictly by ID and never falls back to a name match, so a physical field that carries + // such a name WITHOUT being the ID match (its ID is absent, different, or the logical ID + // matched a different physical field) must be renamed to something the + // `DefaultPhysicalExprAdapter` cannot name-match; otherwise the read would silently + // resolve the wrong column instead of null-filling. Spark's `matchIdField` solves the + // same problem with `generateFakeColumnName` (see `ParquetReadSupport.scala`). + let id_logical_folded: HashSet = if should_match_by_id { logical_schema .fields() .iter() .enumerate() - .filter_map(|(j, lf)| { - parse_field_id(lf).and_then(|id| { - if id_to_phys_names.contains_key(&id) { - None - } else { - Some(logical_folded[j].clone()) - } - }) - }) + .filter(|(_, lf)| parse_field_id(lf).is_some()) + .map(|(j, _)| logical_folded[j].clone()) .collect() } else { HashSet::new() }; + + // Fake names must never collide with a real column from either schema on the folded names + // downstream lookups use, or a requested column differing only by case counts as present + // and loses its default. The lowercase candidate is its own fold; the counter is bumped + // past reserved names, and the set is built on the first fake name so flat reads skip it. + let mut reserved_names: Option> = None; let mut fake_counter: usize = 0; + let mut next_fake_name = || { + let reserved = reserved_names.get_or_insert_with(|| { + logical_folded + .iter() + .chain(physical_folded.iter()) + .map(String::as_str) + .collect() + }); + loop { + fake_counter += 1; + let candidate = format!("__comet_unmatched_field_id_{}", fake_counter); + if !reserved.contains(candidate.as_str()) { + return candidate; + } + } + }; let mut name_map: HashMap = HashMap::new(); let remapped_fields: Vec = physical_schema @@ -202,21 +215,13 @@ fn remap_physical_schema( } } - // Block accidental name match for ID-bearing logical fields whose ID is missing - // from the file. Mirrors Spark's `generateFakeColumnName` in `matchIdField`. - if should_match_by_id - && unmatched_id_logical_folded.contains(&physical_folded[phys_idx]) - { - fake_counter += 1; - let fake_name = format!("__comet_unmatched_field_id_{}", fake_counter); - return Arc::new( - Field::new(fake_name, field.data_type().clone(), field.is_nullable()) - .with_metadata(field.metadata().clone()), - ); - } - - // Name match. Spark's `matchIdField` does not fall through to a name match for - // ID-bearing logical fields, so skip those when the schema is ID-bearing. + // Name match. Spark resolves every non-ID-bearing logical field by name + // (`matchCaseSensitiveField` / `matchCaseInsensitiveField` in + // `clipParquetGroupFields`) even when field-ID matching is on; only ID-bearing + // logical fields skip the name fallback. Case-sensitive mode needs no rename + // here (the downstream adapter's exact-name lookup already hits); the + // case-insensitive lookup rewrites the physical name, and a successful match + // claims the field before the shield below can hide it. if !case_sensitive { let logical_field = logical_schema .fields() @@ -239,9 +244,28 @@ fn remap_physical_schema( .with_metadata(field.metadata().clone()), ); } + return Arc::clone(field); } } + // Shield: any remaining physical field whose name would hit an ID-bearing + // logical field downstream gets a fake name (Spark's `generateFakeColumnName` + // equivalent). ID-bearing logical fields resolve strictly by ID, so a name hit + // on one would read the wrong column instead of null-filling it or leaving it + // to its real ID match. The folded comparison mirrors the matcher that would + // otherwise hit: identity fold in case-sensitive mode, the JVM lowercase fold + // otherwise. + if should_match_by_id && id_logical_folded.contains(&physical_folded[phys_idx]) { + return Arc::new( + Field::new( + next_fake_name(), + field.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ); + } + Arc::clone(field) }) .collect(); @@ -452,6 +476,16 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { None }; + // Resolve every nested struct once per file, the way Spark's `clipParquetSchema` + // recurses through struct, list, and map types while clipping the file schema. Each + // batch reuses the result, and an ambiguity inside it is raised from `rewrite` for the + // columns a read references, whether or not a cast is emitted for them. + let nested_mappings = resolve_nested_mappings( + &logical_file_schema, + &adapted_physical_schema, + &self.parquet_options, + ); + let default_factory = DefaultPhysicalExprAdapterFactory; let default_adapter = default_factory.create( Arc::clone(&logical_file_schema), @@ -469,10 +503,74 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { id_resolved_logical_folded, logical_folded, physical_folded, + nested_mappings, })) } } +/// Per logical field name, the mapping of its nested type against its physical counterpart, +/// or the ambiguity Spark reports for it. Only fields whose type holds a struct are listed. +type NestedMappings = HashMap, SparkError>>; + +fn type_holds_struct(data_type: &DataType) -> bool { + match data_type { + DataType::Struct(_) => true, + DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => { + type_holds_struct(f.data_type()) + } + _ => false, + } +} + +/// Resolve the nested mapping of every logical field whose type holds a struct and that has +/// a physical counterpart. Returns `None` when no field qualifies, so flat reads build +/// nothing here. Ambiguities are kept per field rather than raised: Spark validates only +/// the fields a read requests, and `rewrite` sees which ones those are. +fn resolve_nested_mappings( + logical_schema: &SchemaRef, + physical_schema: &SchemaRef, + parquet_options: &SparkParquetOptions, +) -> Option { + let mut physical_by_name: Option> = None; + let mut mappings = NestedMappings::new(); + for logical_field in logical_schema.fields() { + if !type_holds_struct(logical_field.data_type()) { + continue; + } + // First physical field wins an exact-name tie, as the default adapter's lookup does. + let by_name = physical_by_name.get_or_insert_with(|| { + let mut by_name = HashMap::with_capacity(physical_schema.fields().len()); + for (i, pf) in physical_schema.fields().iter().enumerate() { + by_name.entry(pf.name().as_str()).or_insert(i); + } + by_name + }); + let Some(&physical_index) = by_name.get(logical_field.name().as_str()) else { + continue; + }; + let resolved = resolve_field_mapping( + physical_schema.field(physical_index).data_type(), + logical_field.data_type(), + parquet_options, + ) + .map(Arc::new); + mappings.insert(logical_field.name().clone(), resolved); + } + (!mappings.is_empty()).then_some(mappings) +} + +/// Names of every `Column` referenced by `expr`, in traversal order. +fn referenced_column_names(expr: &Arc) -> Vec { + let mut names: Vec = Vec::new(); + let _ = Arc::clone(expr).transform(|e| { + if let Some(col) = e.downcast_ref::() { + names.push(col.name().to_string()); + } + Ok(Transformed::no(e)) + }); + names +} + /// Spark-compatible physical expression adapter. /// /// This adapter rewrites expressions at planning time to: @@ -515,40 +613,50 @@ struct SparkPhysicalExprAdapter { /// `physical_file_schema` field names pre-folded once, parallel to /// `physical_file_schema.fields()`. See `logical_folded`. physical_folded: Vec, + /// Nested field mappings resolved once in `create` (see `resolve_nested_mappings`), + /// handed to every `CometCastColumnExpr` built here. `None` for schemas without structs. + nested_mappings: Option, } impl PhysicalExprAdapter for SparkPhysicalExprAdapter { fn rewrite(&self, expr: Arc) -> DataFusionResult> { - // In case-insensitive mode, check if any Column in this expression references - // a field with multiple case-insensitive matches in the physical schema. - // Only the columns actually referenced trigger the error (not the whole schema). - if let Some((orig_physical, folded_to_indices)) = &self.original_physical_dup_check { - // Collect referenced column names, then fold them in one JVM crossing rather than one - // per Column node. Physical names were already folded once in `create()`. - let mut col_names: Vec = Vec::new(); - let _ = Arc::::clone(&expr).transform(|e| { - if let Some(col) = e.downcast_ref::() { - col_names.push(col.name().to_string()); - } - Ok(Transformed::no(e)) - }); - let col_refs: Vec<&str> = col_names.iter().map(|s| s.as_str()).collect(); - let col_folded = fold_names(&col_refs, false); - for (name, folded) in col_names.iter().zip(&col_folded) { - // Fields resolved by Parquet field id are selected by id before names are - // compared, so an id-resolved column must not trip the name-ambiguity check - // (mirrors Spark's `matchIdField`, which never raises the duplicate-field error). - if self - .id_resolved_logical_folded - .as_ref() - .is_some_and(|ids| ids.contains(folded)) - { - continue; + // Only the columns this expression references are checked, as Spark validates only + // the fields a read requests. Referenced names are collected once for both checks. + if self.original_physical_dup_check.is_some() || self.nested_mappings.is_some() { + let col_names = referenced_column_names(&expr); + + // In case-insensitive mode, a referenced column with more than one + // case-insensitive match in the physical schema is ambiguous. Names are folded in + // one JVM crossing; physical names were already folded once in `create()`. + if let Some((orig_physical, folded_to_indices)) = &self.original_physical_dup_check { + let col_refs: Vec<&str> = col_names.iter().map(|s| s.as_str()).collect(); + let col_folded = fold_names(&col_refs, false); + for (name, folded) in col_names.iter().zip(&col_folded) { + // Fields resolved by Parquet field id are selected by id before names are + // compared, so an id-resolved column must not trip the name-ambiguity check + // (mirrors Spark's `matchIdField`, which never raises the duplicate-field error). + if self + .id_resolved_logical_folded + .as_ref() + .is_some_and(|ids| ids.contains(folded)) + { + continue; + } + if let Some(err) = + check_column_duplicate(name, folded, folded_to_indices, orig_physical) + { + return Err(spark_error(err)); + } } - if let Some(err) = - check_column_duplicate(name, folded, folded_to_indices, orig_physical) - { - return Err(DataFusionError::External(Box::new(err))); + } + + // An ambiguity inside a referenced column's nested type surfaces here, so it is + // raised for every read of that column and not only when a cast is emitted. + if let Some(nested) = &self.nested_mappings { + for name in &col_names { + if let Some(Err(err)) = nested.get(name.as_str()) { + return Err(spark_error(err.clone())); + } } } } @@ -605,6 +713,20 @@ impl PhysicalExprAdapter for SparkPhysicalExprAdapter { } impl SparkPhysicalExprAdapter { + /// The once-per-file mapping for the logical field named `logical_name`, or a leaf + /// mapping for a field whose type holds no struct. + fn field_mapping_for(&self, logical_name: &str) -> DataFusionResult> { + match self + .nested_mappings + .as_ref() + .and_then(|mappings| mappings.get(logical_name)) + { + Some(Ok(mapping)) => Ok(Arc::clone(mapping)), + Some(Err(err)) => Err(spark_error(err.clone())), + None => Ok(Arc::new(FieldMapping::Leaf)), + } + } + /// Wrap ALL Column expressions that have type mismatches with CometCastColumnExpr. /// This is the fallback path when the default adapter fails (e.g., for complex /// nested type casts like List or Map). Uses `spark_parquet_convert` @@ -673,7 +795,10 @@ impl SparkPhysicalExprAdapter { Arc::clone(logical_field), None, )? - .with_parquet_options(self.parquet_options.clone()), + .with_parquet_options( + self.parquet_options.clone(), + self.field_mapping_for(logical_field.name())?, + ), ); return Ok(Transformed::yes(cast_expr)); } else if column.index() != phys_idx { @@ -955,6 +1080,7 @@ impl SparkPhysicalExprAdapter { | (DataType::Timestamp(_, _), DataType::Timestamp(_, _)) | (DataType::Timestamp(_, _), DataType::Int64) ) { + let field_mapping = self.field_mapping_for(cast.target_field().name())?; let comet_cast: Arc = Arc::new( CometCastColumnExpr::try_new( child, @@ -962,7 +1088,7 @@ impl SparkPhysicalExprAdapter { Arc::clone(cast.target_field()), None, )? - .with_parquet_options(self.parquet_options.clone()), + .with_parquet_options(self.parquet_options.clone(), field_mapping), ); return Ok(Transformed::yes(comet_cast)); } @@ -1165,7 +1291,7 @@ mod test { use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory; use arrow::array::UInt32Array; use arrow::array::{ - BinaryArray, Date32Array, Decimal128Array, Float32Array, Float64Array, Int32Array, + Array, BinaryArray, Date32Array, Decimal128Array, Float32Array, Float64Array, Int32Array, Int64Array, StringArray, TimestampMicrosecondArray, }; use arrow::datatypes::SchemaRef; @@ -1177,13 +1303,17 @@ mod test { use datafusion::datasource::source::DataSourceExec; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::TaskContext; + use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::ExecutionPlan; + use datafusion::scalar::ScalarValue; use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename; use datafusion_comet_spark_expr::EvalMode; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use futures::StreamExt; use parquet::arrow::ArrowWriter; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use parquet::file::metadata::KeyValue; + use parquet::file::properties::WriterProperties; use std::collections::HashMap; use std::fs::File; use std::sync::Arc; @@ -1901,6 +2031,182 @@ mod test { assert_eq!(remapped.field(0).name(), "a"); } + /// Build a nullable Int64 field carrying a Parquet field ID. + fn field_with_id(name: &str, id: i32) -> Field { + Field::new(name, DataType::Int64, true).with_metadata(id_meta(&id.to_string())) + } + + /// Write a Parquet file from `file_schema`/`columns`, then scan it with + /// `required_schema` through the Spark expression adapter and return the first batch. + async fn scan_with_adapter( + file_schema: SchemaRef, + columns: Vec>, + required_schema: SchemaRef, + spark_parquet_options: SparkParquetOptions, + ) -> Result { + scan_with_defaults( + file_schema, + columns, + required_schema, + spark_parquet_options, + None, + ) + .await + } + + /// `scan_with_adapter` with column defaults for fields missing from the file. + async fn scan_with_defaults( + file_schema: SchemaRef, + columns: Vec>, + required_schema: SchemaRef, + spark_parquet_options: SparkParquetOptions, + default_values: Option>, + ) -> Result { + let batch = RecordBatch::try_new(Arc::clone(&file_schema), columns).unwrap(); + + let filename = get_temp_filename(); + let filename = filename.as_path().as_os_str().to_str().unwrap().to_string(); + let file = File::create(&filename).unwrap(); + // Spark stamps key-value metadata into every file it writes. arrow-rs folds that into + // the file schema's metadata, so the file schema never compares equal to the requested + // schema and DataFusion always runs the expression adapter, as it does for real reads. + let props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + "org.apache.spark.version".to_string(), + "3.5.0".to_string(), + )])) + .build(); + let mut writer = ArrowWriter::try_new(file, file_schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let expr_adapter_factory: Arc = Arc::new( + SparkPhysicalExprAdapterFactory::new(spark_parquet_options, default_values), + ); + + let object_store_url = ObjectStoreUrl::local_filesystem(); + let parquet_source = ParquetSource::new(required_schema); + let files = FileGroup::new(vec![PartitionedFile::from_path(filename).unwrap()]); + let file_scan_config = + FileScanConfigBuilder::new(object_store_url, Arc::new(parquet_source)) + .with_file_groups(vec![files]) + .with_expr_adapter(Some(expr_adapter_factory)) + .build(); + + let parquet_exec = DataSourceExec::new(Arc::new(file_scan_config)); + let mut stream = parquet_exec + .execute(0, Arc::new(TaskContext::default())) + .unwrap(); + stream.next().await.unwrap() + } + + /// File: one column `κ` (U+03BA) with field ID 2 holding 7. Required: `Κ` (U+039A, + /// field ID 1) and ID-less `κ`; case-sensitive, field-ID reading on. Spark routes `Κ` + /// through `matchIdField` (no ID 1 in the file -> null-filled behind a faked REQUESTED + /// name) and resolves `κ` by exact name through `matchCaseSensitiveField`, reading the + /// real column: the result is (NULL, 7), never (NULL, NULL). + #[tokio::test] + async fn parquet_field_id_miss_null_fills_but_exact_name_sibling_still_reads() { + let file_schema = Arc::new(Schema::new(vec![field_with_id("\u{3BA}", 2)])); + let col = Arc::new(Int64Array::from(vec![7])) as Arc; + let required_schema = Arc::new(Schema::new(vec![ + field_with_id("\u{39A}", 1), + Field::new("\u{3BA}", DataType::Int64, true), + ])); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.case_sensitive = true; + opts.use_field_id = true; + + let batch = scan_with_adapter(file_schema, vec![col], required_schema, opts) + .await + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let capital = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(capital.is_null(0)); + let small = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(!small.is_null(0)); + assert_eq!(small.value(0), 7); + } + + /// Case-insensitive variant of the Kappa scenario. Spark's `matchCaseInsensitiveField` + /// resolves the ID-less requested `κ` through the `toLowerCase(Locale.ROOT)`-keyed map + /// of the file's fields, which holds the physical `κ`; the unmatched-ID requested `Κ` + /// is null-filled and never blocks that lookup. Same (NULL, 7) result as the + /// case-sensitive read. + #[tokio::test] + async fn parquet_field_id_miss_case_insensitive_sibling_still_reads() { + let file_schema = Arc::new(Schema::new(vec![field_with_id("\u{3BA}", 2)])); + let col = Arc::new(Int64Array::from(vec![7])) as Arc; + let required_schema = Arc::new(Schema::new(vec![ + field_with_id("\u{39A}", 1), + Field::new("\u{3BA}", DataType::Int64, true), + ])); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.case_sensitive = false; + opts.use_field_id = true; + + let batch = scan_with_adapter(file_schema, vec![col], required_schema, opts) + .await + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let capital = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(capital.is_null(0)); + let small = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(!small.is_null(0)); + assert_eq!(small.value(0), 7); + } + + /// File: a stray ID-less column literally named `A` = [10, 20] FIRST, then `a` with + /// field ID 1 = [1, 2]. Required: `A` with field ID 1, case-insensitive, field-ID + /// reading on. Spark's `matchIdField` resolves requested `A` to physical `a` by ID; the + /// stray `A` is never requested, and no case-insensitive duplicate error fires because + /// ID-routed requested fields never enter the name lookup. Expect [1, 2] -- neither the + /// stray column's data nor a spurious duplicate-field error. + #[tokio::test] + async fn parquet_field_id_match_beats_stray_column_with_requested_name() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("A", DataType::Int64, true), + field_with_id("a", 1), + ])); + let stray = Arc::new(Int64Array::from(vec![10, 20])) as Arc; + let matched = Arc::new(Int64Array::from(vec![1, 2])) as Arc; + let required_schema = Arc::new(Schema::new(vec![field_with_id("A", 1)])); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.case_sensitive = false; + opts.use_field_id = true; + + let batch = scan_with_adapter(file_schema, vec![stray, matched], required_schema, opts) + .await + .unwrap(); + assert_eq!(batch.num_rows(), 2); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), 1); + assert_eq!(values.value(1), 2); + } + /// Field-id precedence in the case-insensitive duplicate check: an explicit `ω` (id 2) /// reading a file that holds both `ω` (id 2) and `Ω` (id 1) must resolve by id (Spark's /// `matchIdField` selects the id before ever comparing names) rather than raising a @@ -1935,4 +2241,88 @@ mod test { rewritten.err() ); } + + /// A requested column named like the shield's placeholder must still receive its + /// configured default. File: `k` (id 2). Required: `k` (id 1) and an id-less + /// `__COMET_UNMATCHED_FIELD_ID_1` with default 7, case-insensitive, field-id reading on. + /// The file's `k` is not the id match for requested `k`, so it is hidden behind a + /// placeholder name; that placeholder must not fold onto the requested column, or the + /// missing-column check treats it as present and the default is lost. + #[tokio::test] + async fn parquet_shield_placeholder_never_folds_onto_requested_column() { + let file_schema = Arc::new(Schema::new(vec![field_with_id("k", 2)])); + let col = Arc::new(Int64Array::from(vec![1])) as Arc; + let required_schema = Arc::new(Schema::new(vec![ + field_with_id("k", 1), + Field::new("__COMET_UNMATCHED_FIELD_ID_1", DataType::Int64, true), + ])); + let defaults = HashMap::from([( + Column::new("__COMET_UNMATCHED_FIELD_ID_1", 1), + ScalarValue::Int64(Some(7)), + )]); + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.case_sensitive = false; + opts.use_field_id = true; + + let batch = scan_with_defaults( + file_schema, + vec![col], + required_schema, + opts, + Some(defaults), + ) + .await + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let k = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!( + k.is_null(0), + "requested k (id 1) has no id match in the file" + ); + let defaulted = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(!defaulted.is_null(0), "configured default must apply"); + assert_eq!(defaulted.value(0), 7); + } + + /// File and requested schema are identical: `s` holding `x` and `y` that both carry + /// field id 1. No column needs conversion, so no cast is ever emitted, yet Spark's + /// `clipParquetSchema` rejects the read because requested id 1 resolves to two file + /// fields. The validation must therefore run when the file schema is mapped, not + /// only inside a cast. + #[tokio::test] + async fn parquet_duplicate_struct_field_id_rejected_without_cast() { + let child_fields = + arrow::datatypes::Fields::from(vec![field_with_id("x", 1), field_with_id("y", 1)]); + let struct_field = Field::new("s", DataType::Struct(child_fields.clone()), true) + .with_metadata(id_meta("10")); + let file_schema = Arc::new(Schema::new(vec![struct_field.clone()])); + let required_schema = Arc::new(Schema::new(vec![struct_field])); + let children: Vec> = vec![ + Arc::new(Int64Array::from(vec![42])), + Arc::new(Int64Array::from(vec![43])), + ]; + let col = Arc::new(arrow::array::StructArray::new(child_fields, children, None)) + as Arc; + + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + opts.use_field_id = true; + + let err = scan_with_adapter(file_schema, vec![col], required_schema, opts) + .await + .expect_err("requested id 1 matches two file fields and must error"); + let msg = err.to_string(); + assert!( + msg.contains("_LEGACY_ERROR_TEMP_2094") && msg.contains("id=1"), + "expected duplicate field id error, got: {msg}" + ); + } } 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 799197bf33..dbe306b42b 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -2002,6 +2002,39 @@ abstract class ParquetReadSuite extends CometTestBase { } } + // Spark's `clipParquetSchema` runs `matchIdField` at every nesting level while clipping the + // file schema, so a struct child id duplicated in the file is rejected even when the read + // schema is identical to the file schema and no column needs any conversion. + test("duplicate field id inside a struct is rejected without a cast") { + withSQLConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") { + withTempPath { dir => + val schema = + new StructType() + .add( + "s", + new StructType() + .add("x", LongType, true, withId(1)) + .add("y", LongType, true, withId(1)), + true, + withId(2)) + + val writeData = Seq(Row(Row(42L, 43L))) + spark + .createDataFrame(spark.sparkContext.parallelize(writeData), schema) + .write + .mode("overwrite") + .parquet(dir.getCanonicalPath) + + val cause = intercept[SparkException] { + spark.read.schema(schema).parquet(dir.getCanonicalPath).collect() + }.getCause + assert( + cause.isInstanceOf[RuntimeException] && + cause.getMessage.contains("Found duplicate field(s)")) + } + } + } + // Verbatim port of Spark `ParquetFieldIdIOSuite.test("read parquet file without ids")`, // for the same reason as the duplicate-id test above. test("read parquet file without ids") {