diff --git a/datafusion/common/src/nested_struct.rs b/datafusion/common/src/nested_struct.rs index e915b91b911c..d54e08701d18 100644 --- a/datafusion/common/src/nested_struct.rs +++ b/datafusion/common/src/nested_struct.rs @@ -19,13 +19,17 @@ use crate::error::{_plan_err, Result}; use arrow::{ array::{ Array, ArrayRef, AsArray, DictionaryArray, FixedSizeListArray, GenericListArray, - GenericListViewArray, StructArray, downcast_integer, make_array, new_null_array, + GenericListViewArray, MapArray, StructArray, UInt64Array, downcast_integer, + make_array, new_null_array, }, buffer::NullBuffer, - compute::{CastOptions, can_cast_types, cast_with_options}, + compute::{CastOptions, can_cast_types, cast_with_options, take}, datatypes::{DataType, DataType::Struct, Field, FieldRef}, }; -use std::{collections::HashSet, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; /// Cast a struct column to match target struct fields, handling nested structs recursively. /// @@ -205,6 +209,14 @@ pub fn cast_column( (DataType::LargeListView(_), DataType::LargeListView(target_inner)) => { cast_list_view_column::(source_col, target_inner, cast_options) } + (DataType::Map(_, _), DataType::Map(target_entries, target_sorted)) => { + cast_map_column( + source_col.as_map(), + target_entries, + *target_sorted, + cast_options, + ) + } ( DataType::Dictionary(source_key_type, _), DataType::Dictionary(target_key_type, target_value_type), @@ -340,6 +352,73 @@ fn mask_array_values( )) } +/// Casts Map children by their semantic positions: key at index 0 and value at +/// index 1. Technical entry field names are taken from the target schema. +/// +/// Nested Struct fields within keys and values are still matched by name. Key +/// evolution is restricted by [`validate_map_key_compatibility`] so it cannot +/// remove identity-bearing fields; sorted Maps require an unchanged key type. +/// Entries hidden by null Map parents are compacted before casting. +fn cast_map_column( + source_map: &MapArray, + target_entries: &FieldRef, + target_sorted: bool, + cast_options: &CastOptions, +) -> Result { + let DataType::Map(source_entries, source_sorted) = source_map.data_type() else { + unreachable!("MapArray data type must be Map") + }; + let (target_key, target_value) = validate_map_compatibility( + source_entries, + *source_sorted, + target_entries, + target_sorted, + )?; + + let offsets = source_map.value_offsets(); + let has_unreachable_entries = offsets[0] != 0 + || offsets[offsets.len() - 1] as usize != source_map.entries().len(); + let needs_compaction = has_unreachable_entries + || source_map.offsets().has_non_empty_nulls(source_map.nulls()); + let compacted_map = if needs_compaction { + Some(compact_map_entries(source_map)?) + } else { + None + }; + let source_map = compacted_map.as_ref().unwrap_or(source_map); + + let cast_keys = cast_column(source_map.keys(), target_key.data_type(), cast_options) + .map_err(|error| error.context("While casting Map keys"))?; + let cast_values = + cast_column(source_map.values(), target_value.data_type(), cast_options) + .map_err(|error| error.context("While casting Map values"))?; + let Struct(target_fields) = target_entries.data_type() else { + unreachable!("validated Map entries must be Struct") + }; + let cast_entries = + StructArray::new(target_fields.clone(), vec![cast_keys, cast_values], None); + + Ok(Arc::new(MapArray::try_new( + Arc::clone(target_entries), + source_map.offsets().clone(), + cast_entries, + source_map.nulls().cloned(), + target_sorted, + )?)) +} + +/// Returns an equivalent MapArray whose entries contain only values reachable +/// from visible Map rows. +/// +/// Arrow Map arrays can contain unreachable entries after slicing, or entries +/// hidden behind null parent rows. An identity `take` rebuilds the Map through +/// Arrow's selection kernel, normalizing offsets and dropping those unreachable +/// child entries before recursive key/value casts are applied. +fn compact_map_entries(map: &MapArray) -> Result { + let indices = UInt64Array::from_iter_values(0..map.len() as u64); + Ok(take(map, &indices, None)?.as_map().clone()) +} + fn cast_dictionary_column( source_col: &ArrayRef, source_key_type: &DataType, @@ -490,8 +569,162 @@ fn validate_field_compatibility( ) } +fn validate_map_compatibility<'a>( + source_entries: &Field, + source_sorted: bool, + target_entries: &'a Field, + target_sorted: bool, +) -> Result<(&'a FieldRef, &'a FieldRef)> { + if source_sorted != target_sorted { + return _plan_err!("Cannot change Map sorted flag during schema adaptation"); + } + let (source_key, source_value) = validate_map_entries_field(source_entries)?; + let (target_key, target_value) = validate_map_entries_field(target_entries)?; + validate_map_key_compatibility(source_key, target_key, target_sorted)?; + validate_field_compatibility(source_value, target_value)?; + Ok((target_key, target_value)) +} + +fn validate_map_key_compatibility( + source_key: &Field, + target_key: &Field, + sorted: bool, +) -> Result<()> { + if sorted && source_key.data_type() != target_key.data_type() { + return _plan_err!("Cannot evolve key type of a sorted Map"); + } + validate_map_key_data_type(source_key.data_type(), target_key.data_type()) +} + +fn validate_map_key_data_type( + source_type: &DataType, + target_type: &DataType, +) -> Result<()> { + if source_type == target_type { + return Ok(()); + } + + match (source_type, target_type) { + (Struct(source_fields), Struct(target_fields)) => { + if !has_one_of_more_common_fields(source_fields, target_fields) { + return _plan_err!( + "Cannot cast Map key Struct because there is no field name overlap" + ); + } + + let target_by_name: HashMap<&str, &FieldRef> = target_fields + .iter() + .map(|field| (field.name().as_str(), field)) + .collect(); + let source_names: HashSet<&str> = source_fields + .iter() + .map(|field| field.name().as_str()) + .collect(); + + for source_field in source_fields { + let Some(target_field) = target_by_name.get(source_field.name().as_str()) + else { + return _plan_err!( + "Cannot remove field '{}' from a Map key Struct", + source_field.name() + ); + }; + if source_field.is_nullable() && !target_field.is_nullable() { + return _plan_err!( + "Cannot cast nullable Map key field '{}' to non-nullable field", + source_field.name() + ); + } + validate_map_key_data_type( + source_field.data_type(), + target_field.data_type(), + )?; + } + for target_field in target_fields { + if !source_names.contains(target_field.name().as_str()) + && !target_field.is_nullable() + { + return _plan_err!( + "Cannot add non-nullable field '{}' to a Map key Struct", + target_field.name() + ); + } + } + Ok(()) + } + _ if is_injective_map_key_cast(source_type, target_type) => Ok(()), + _ => _plan_err!( + "Cannot safely evolve Map key type from {} to {}", + source_type, + target_type + ), + } +} + +/// Returns true when casting a Map key from `source_type` to `target_type` is +/// known to preserve key identity. +/// +/// Keep this list conservative. Map keys define equality/lookup identity, so +/// only representation widenings that cannot merge distinct source keys belong +/// here. Other Arrow-supported casts, including decimal/timestamp widening and +/// dictionaries, should stay rejected until their identity semantics are +/// deliberately reviewed. +fn is_injective_map_key_cast(source_type: &DataType, target_type: &DataType) -> bool { + matches!( + (source_type, target_type), + ( + DataType::Int8, + DataType::Int16 | DataType::Int32 | DataType::Int64 + ) | (DataType::Int16, DataType::Int32 | DataType::Int64) + | (DataType::Int32, DataType::Int64) + | ( + DataType::UInt8, + DataType::UInt16 | DataType::UInt32 | DataType::UInt64 + ) + | (DataType::UInt16, DataType::UInt32 | DataType::UInt64) + | (DataType::UInt32, DataType::UInt64) + | (DataType::Utf8, DataType::LargeUtf8 | DataType::Utf8View) + | ( + DataType::Binary, + DataType::LargeBinary | DataType::BinaryView + ) + ) +} + +fn validate_map_entries_field(entries: &Field) -> Result<(&FieldRef, &FieldRef)> { + if entries.is_nullable() { + return _plan_err!("Map entries field must be non-nullable"); + } + + let Struct(fields) = entries.data_type() else { + return _plan_err!("Map entries field must be a struct"); + }; + if fields.len() != 2 { + return _plan_err!( + "Map entries struct must contain exactly two fields, found {}", + fields.len() + ); + } + if fields[0].is_nullable() { + return _plan_err!("Map key field must be non-nullable"); + } + Ok((&fields[0], &fields[1])) +} + /// Validates that `source_type` can be cast to `target_type`, recursively /// handling container types that wrap structs. +/// +/// # Map evolution +/// +/// Map entry children are matched by position (key, then value), regardless of +/// their technical field names. Nested Struct fields within each child are +/// matched by name. Values use the standard Struct evolution rules: nullable +/// target fields may be added and extra source fields are omitted. +/// +/// Unsorted Struct keys may add nullable fields and use only injective primitive +/// widening casts, but may not remove source fields. Sorted Maps require an +/// unchanged key type. Map entries and keys must remain non-nullable. Source +/// and target sorted flags must match. pub fn validate_data_type_compatibility( field_name: &str, source_type: &DataType, @@ -513,6 +746,9 @@ pub fn validate_data_type_compatibility( | (DataType::LargeListView(s), DataType::LargeListView(t)) => { validate_field_compatibility(s, t)?; } + (DataType::Map(s, source_sorted), DataType::Map(t, target_sorted)) => { + validate_map_compatibility(s, *source_sorted, t, *target_sorted)?; + } (DataType::Dictionary(s_key, s_val), DataType::Dictionary(t_key, t_val)) => { if !can_cast_types(s_key, t_key) { return _plan_err!( @@ -539,13 +775,20 @@ pub fn validate_data_type_compatibility( } /// Returns true if casting from `source_type` to `target_type` requires -/// name-based nested struct casting logic, rather than Arrow's standard cast. +/// DataFusion's specialized nested/container casting logic rather than Arrow's +/// standard cast. /// /// This is the case when both types are struct types, or both are the same /// container type (List, LargeList, equal-width FixedSizeList, ListView, -/// LargeListView, Dictionary) wrapping types that recursively contain structs. +/// LargeListView, Map, Dictionary) wrapping types that recursively contain structs. +/// +/// Maps are always Struct-backed: their entries field is a Struct containing key +/// and value children. Therefore a compatible Map-to-Map adaptation can return +/// true even when the user-visible key and value types are primitive. This +/// deliberately routes Map adaptation through the specialized Map casting path +/// so key/value semantics, sorted flags, and entry compaction are enforced. /// -/// Use this predicate at both planning time (to decide whether to apply struct +/// Use this predicate at both planning time (to decide whether to apply nested /// compatibility validation) and execution time (to decide whether to route /// through [`cast_column`] instead of Arrow's generic cast). pub fn requires_nested_struct_cast( @@ -566,6 +809,9 @@ pub fn requires_nested_struct_cast( | (DataType::LargeListView(s), DataType::LargeListView(t)) => { requires_nested_struct_cast(s.data_type(), t.data_type()) } + (DataType::Map(s, _), DataType::Map(t, _)) => { + requires_nested_struct_cast(s.data_type(), t.data_type()) + } (DataType::Dictionary(_, s_val), DataType::Dictionary(_, t_val)) => { requires_nested_struct_cast(s_val, t_val) } @@ -600,7 +846,7 @@ mod tests { ListArray, ListViewArray, MapArray, MapBuilder, NullArray, StringArray, StringBuilder, }, - buffer::{NullBuffer, ScalarBuffer}, + buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}, datatypes::{DataType, Field, FieldRef, Int32Type}, }; /// Macro to extract and downcast a column from a StructArray @@ -1292,6 +1538,566 @@ mod tests { assert!(b_col.is_null(1)); } + fn map_type(key_type: DataType, value_type: DataType) -> DataType { + map_type_with_entry_names(key_type, value_type, "keys", "values", false) + } + + fn map_type_with_entry_names( + key_type: DataType, + value_type: DataType, + key_name: &str, + value_name: &str, + sorted: bool, + ) -> DataType { + DataType::Map( + Arc::new(non_null_field( + "entries", + struct_type(vec![ + non_null_field(key_name, key_type), + field(value_name, value_type), + ]), + )), + sorted, + ) + } + + fn struct_map_array() -> ArrayRef { + struct_map_array_with_sorted(false) + } + + fn struct_map_array_with_sorted(sorted: bool) -> ArrayRef { + struct_map_array_with_key_fields(sorted, false) + } + + fn struct_map_array_with_key_fields(sorted: bool, include_tenant: bool) -> ArrayRef { + let mut key_fields = vec![( + arc_field("id", DataType::Int32), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]; + if include_tenant { + key_fields.push(( + arc_field("tenant", DataType::Utf8), + Arc::new(StringArray::from(vec!["a"])) as ArrayRef, + )); + } + let keys = StructArray::from(key_fields); + let values = StructArray::from(vec![ + ( + arc_field("amount", DataType::Int32), + Arc::new(Int32Array::from(vec![10])) as ArrayRef, + ), + ( + arc_field("ignored", DataType::Utf8), + Arc::new(StringArray::from(vec!["x"])) as ArrayRef, + ), + ]); + let entries = StructArray::new( + vec![ + Arc::new(non_null_field("keys", keys.data_type().clone())), + arc_field("values", values.data_type().clone()), + ] + .into(), + vec![Arc::new(keys), Arc::new(values)], + None, + ); + Arc::new(MapArray::new( + Arc::new(non_null_field("entries", entries.data_type().clone())), + OffsetBuffer::new(vec![0, 1, 1].into()), + entries, + Some(NullBuffer::from(vec![true, false])), + sorted, + )) + } + + fn nested_struct_map_array( + key_name: &str, + value_name: &str, + sorted: bool, + ) -> ArrayRef { + let key_nested = StructArray::from(vec![( + arc_field("id", DataType::Int32), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]); + let keys = StructArray::from(vec![( + arc_field("nested", key_nested.data_type().clone()), + Arc::new(key_nested) as ArrayRef, + )]); + let value_nested = StructArray::from(vec![( + arc_field("amount", DataType::Int32), + Arc::new(Int32Array::from(vec![10])) as ArrayRef, + )]); + let values = StructArray::from(vec![( + arc_field("nested", value_nested.data_type().clone()), + Arc::new(value_nested) as ArrayRef, + )]); + let entries = StructArray::new( + vec![ + Arc::new(non_null_field(key_name, keys.data_type().clone())), + arc_field(value_name, values.data_type().clone()), + ] + .into(), + vec![Arc::new(keys), Arc::new(values)], + None, + ); + Arc::new(MapArray::new( + Arc::new(non_null_field("entries", entries.data_type().clone())), + OffsetBuffer::new(vec![0, 1].into()), + entries, + None, + sorted, + )) + } + + #[test] + fn test_map_entry_names_match_positionally_and_adapt_nested_structs() { + let source_col = nested_struct_map_array("source_keys", "source_values", false); + let target_type = map_type_with_entry_names( + struct_type(vec![field( + "nested", + struct_type(vec![ + field("id", DataType::Int64), + field("label", DataType::Utf8), + ]), + )]), + struct_type(vec![field( + "nested", + struct_type(vec![ + field("amount", DataType::Int64), + field("currency", DataType::Utf8), + ]), + )]), + "key", + "value", + false, + ); + + assert!( + validate_data_type_compatibility( + "map_col", + source_col.data_type(), + &target_type + ) + .is_ok() + ); + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let map = result.as_any().downcast_ref::().unwrap(); + let (key_field, value_field) = map.entries_fields(); + assert_eq!(key_field.name(), "key"); + assert_eq!(value_field.name(), "value"); + let keys = map.keys().as_struct(); + let key_nested = keys.column_by_name("nested").unwrap().as_struct(); + assert_eq!(get_column_as!(key_nested, "id", Int64Array).value(0), 1); + assert!(get_column_as!(key_nested, "label", StringArray).is_null(0)); + let values = map.values().as_struct(); + let value_nested = values.column_by_name("nested").unwrap().as_struct(); + assert_eq!( + get_column_as!(value_nested, "amount", Int64Array).value(0), + 10 + ); + assert!(get_column_as!(value_nested, "currency", StringArray).is_null(0)); + } + + #[test] + fn test_unsorted_map_key_struct_field_removal_rejected_by_planner_and_runtime() { + let source_col = struct_map_array_with_key_fields(false, true); + let target_type = map_type( + struct_type(vec![field("id", DataType::Int32)]), + struct_type(vec![ + field("amount", DataType::Int32), + field("ignored", DataType::Utf8), + ]), + ); + assert_map_planning_runtime_error( + &source_col, + &target_type, + "Cannot remove field 'tenant' from a Map key Struct", + ); + } + + #[test] + fn test_unsorted_map_non_injective_key_cast_rejected() { + let source_col = struct_map_array(); + let target_type = map_type( + struct_type(vec![field("id", DataType::Float64)]), + struct_type(vec![ + field("amount", DataType::Int32), + field("ignored", DataType::Utf8), + ]), + ); + assert_map_planning_runtime_error( + &source_col, + &target_type, + "Cannot safely evolve Map key type", + ); + } + + #[test] + fn test_unsorted_map_key_struct_no_overlap_rejected_by_planner_and_runtime() { + let keys = StructArray::new_empty_fields(1, Some(NullBuffer::from(vec![true]))); + let values = StructArray::from(vec![( + arc_field("amount", DataType::Int32), + Arc::new(Int32Array::from(vec![10])) as ArrayRef, + )]); + let entries = StructArray::new( + vec![ + Arc::new(non_null_field("keys", keys.data_type().clone())), + arc_field("values", values.data_type().clone()), + ] + .into(), + vec![Arc::new(keys), Arc::new(values)], + None, + ); + let source_col: ArrayRef = Arc::new(MapArray::new( + Arc::new(non_null_field("entries", entries.data_type().clone())), + OffsetBuffer::new(vec![0, 1].into()), + entries, + None, + false, + )); + let target_type = map_type( + struct_type(vec![field("new_id", DataType::Int32)]), + struct_type(vec![field("amount", DataType::Int32)]), + ); + assert_map_planning_runtime_error( + &source_col, + &target_type, + "Cannot cast Map key Struct because there is no field name overlap", + ); + } + + #[test] + fn test_sorted_map_key_struct_schema_evolution_rejected() { + let source_col = struct_map_array_with_sorted(true); + let target_type = map_type_with_entry_names( + struct_type(vec![ + field("id", DataType::Int32), + field("label", DataType::Utf8), + ]), + struct_type(vec![ + field("amount", DataType::Int32), + field("ignored", DataType::Utf8), + ]), + "keys", + "values", + true, + ); + assert_map_planning_runtime_error( + &source_col, + &target_type, + "Cannot evolve key type of a sorted Map", + ); + } + + #[test] + fn test_sorted_map_value_struct_schema_evolution_preserves_sorted() { + let source_col = struct_map_array_with_sorted(true); + let target_type = map_type_with_entry_names( + struct_type(vec![field("id", DataType::Int32)]), + struct_type(vec![ + field("amount", DataType::Int64), + field("currency", DataType::Utf8), + ]), + "keys", + "values", + true, + ); + assert!( + validate_data_type_compatibility( + "map_col", + source_col.data_type(), + &target_type + ) + .is_ok() + ); + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + assert_eq!(result.data_type(), &target_type); + let map = result.as_any().downcast_ref::().unwrap(); + let values = map.values().as_struct(); + assert_eq!(get_column_as!(values, "amount", Int64Array).value(0), 10); + assert!(get_column_as!(values, "currency", StringArray).is_null(0)); + } + + #[test] + fn test_sorted_map_value_struct_schema_evolution_to_unsorted_rejected() { + let source_col = struct_map_array_with_sorted(true); + let target_type = map_type_with_entry_names( + struct_type(vec![ + field("id", DataType::Int32), + field("label", DataType::Utf8), + ]), + struct_type(vec![ + field("amount", DataType::Int64), + field("currency", DataType::Utf8), + ]), + "keys", + "values", + false, + ); + + assert_map_planning_runtime_error( + &source_col, + &target_type, + "Cannot change Map sorted flag", + ); + } + + #[test] + fn test_null_map_parent_hides_invalid_nested_value_cast() { + let values = StructArray::from(vec![( + arc_field("amount", DataType::Utf8), + Arc::new(StringArray::from(vec!["bad", "2"])) as ArrayRef, + )]); + let keys = StringArray::from(vec!["hidden", "visible"]); + let entries = StructArray::new( + vec![ + Arc::new(non_null_field("keys", DataType::Utf8)), + arc_field("values", values.data_type().clone()), + ] + .into(), + vec![Arc::new(keys), Arc::new(values)], + None, + ); + let source_col: ArrayRef = Arc::new(MapArray::new( + Arc::new(non_null_field("entries", entries.data_type().clone())), + OffsetBuffer::new(vec![0, 1, 2].into()), + entries, + Some(NullBuffer::from(vec![false, true])), + false, + )); + let target_type = map_type( + DataType::Utf8, + struct_type(vec![field("amount", DataType::Int32)]), + ); + + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let map = result.as_any().downcast_ref::().unwrap(); + assert!(map.is_null(0)); + assert_eq!(map.value_offsets(), &[0, 0, 1]); + let visible_entries = map.value(1); + let visible_values = visible_entries + .column_by_name("values") + .unwrap() + .as_struct(); + assert_eq!( + get_column_as!(visible_values, "amount", Int32Array).value(0), + 2 + ); + } + + #[test] + fn test_sliced_map_ignores_unreachable_invalid_nested_value() { + let values = StructArray::from(vec![( + arc_field("amount", DataType::Utf8), + Arc::new(StringArray::from(vec!["bad", "2"])) as ArrayRef, + )]); + let entries = StructArray::new( + vec![ + Arc::new(non_null_field("keys", DataType::Utf8)), + arc_field("values", values.data_type().clone()), + ] + .into(), + vec![ + Arc::new(StringArray::from(vec!["unreachable", "visible"])), + Arc::new(values), + ], + None, + ); + let source_col: ArrayRef = Arc::new( + MapArray::new( + Arc::new(non_null_field("entries", entries.data_type().clone())), + OffsetBuffer::new(vec![0, 1, 2].into()), + entries, + None, + false, + ) + .slice(1, 1), + ); + let target_type = map_type( + DataType::Utf8, + struct_type(vec![field("amount", DataType::Int32)]), + ); + + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let map = result.as_any().downcast_ref::().unwrap(); + assert_eq!(map.value_offsets(), &[0, 1]); + let values = map.values().as_struct(); + assert_eq!(get_column_as!(values, "amount", Int32Array).value(0), 2); + } + + #[test] + fn test_cast_map_struct_key_and_value() { + let source_col = struct_map_array(); + let target_type = map_type( + struct_type(vec![ + field("id", DataType::Int64), + field("label", DataType::Utf8), + ]), + struct_type(vec![ + field("amount", DataType::Int64), + field("currency", DataType::Utf8), + ]), + ); + + assert!(requires_nested_struct_cast( + source_col.data_type(), + &target_type + )); + assert!( + validate_data_type_compatibility( + "map_col", + source_col.data_type(), + &target_type + ) + .is_ok() + ); + + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let map = result.as_any().downcast_ref::().unwrap(); + assert!(map.is_valid(0)); + assert!(map.is_null(1)); + + let keys = map.keys().as_any().downcast_ref::().unwrap(); + let key_ids = get_column_as!(keys, "id", Int64Array); + assert_eq!(key_ids.values(), &[1]); + let key_labels = get_column_as!(keys, "label", StringArray); + assert!(key_labels.is_null(0)); + + let values = map.values().as_any().downcast_ref::().unwrap(); + let amounts = get_column_as!(values, "amount", Int64Array); + assert_eq!(amounts.values(), &[10]); + let currencies = get_column_as!(values, "currency", StringArray); + assert!(currencies.is_null(0)); + assert!(values.column_by_name("ignored").is_none()); + } + + #[test] + fn test_cast_all_null_map_with_struct_value() { + let source_type = map_type( + DataType::Utf8, + struct_type(vec![field("amount", DataType::Int32)]), + ); + let target_type = map_type( + DataType::Utf8, + struct_type(vec![ + field("amount", DataType::Int64), + field("currency", DataType::Utf8), + ]), + ); + let source_col = new_null_array(&source_type, 2); + + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let map = result.as_any().downcast_ref::().unwrap(); + assert_eq!(map.null_count(), 2); + assert!(map.entries().is_empty()); + assert_eq!(map.data_type(), &target_type); + } + + #[test] + fn test_map_entry_invariants_rejected_during_validation() { + let source_col = struct_map_array(); + let DataType::Map(target_entries, sorted) = map_type( + struct_type(vec![field("id", DataType::Int32)]), + struct_type(vec![field("amount", DataType::Int32)]), + ) else { + unreachable!() + }; + let nullable_entries = DataType::Map( + Arc::new(target_entries.as_ref().clone().with_nullable(true)), + sorted, + ); + let error = validate_data_type_compatibility( + "map_col", + source_col.data_type(), + &nullable_entries, + ) + .unwrap_err() + .to_string(); + assert_contains!(error, "Map entries field must be non-nullable"); + + let Struct(fields) = target_entries.data_type() else { + unreachable!() + }; + let nullable_key_entries = Arc::new(Field::new( + "entries", + Struct( + vec![ + Arc::new(fields[0].as_ref().clone().with_nullable(true)), + Arc::clone(&fields[1]), + ] + .into(), + ), + false, + )); + let error = validate_data_type_compatibility( + "map_col", + source_col.data_type(), + &DataType::Map(nullable_key_entries, sorted), + ) + .unwrap_err() + .to_string(); + assert_contains!(error, "Map key field must be non-nullable"); + + let sorted_target = DataType::Map(target_entries, !sorted); + assert_map_planning_runtime_error( + &source_col, + &sorted_target, + "Cannot change Map sorted flag", + ); + } + + fn assert_map_planning_runtime_error( + source: &ArrayRef, + target: &DataType, + expected: &str, + ) { + let planning_error = + validate_data_type_compatibility("map_col", source.data_type(), target) + .unwrap_err() + .to_string(); + assert_contains!(planning_error, expected); + + let runtime_error = cast_column(source, target, &DEFAULT_CAST_OPTIONS) + .unwrap_err() + .to_string(); + assert_contains!(runtime_error, expected); + } + + #[test] + fn test_map_struct_planner_runtime_parity_on_invalid_evolution() { + let source_col = struct_map_array(); + let target_type = map_type( + struct_type(vec![field("id", DataType::Int32)]), + struct_type(vec![ + field("amount", DataType::Int32), + non_null_field("currency", DataType::Utf8), + ]), + ); + assert_map_planning_runtime_error( + &source_col, + &target_type, + "target field 'currency' is non-nullable", + ); + + let incompatible_target = map_type( + struct_type(vec![field("id", DataType::Int32)]), + struct_type(vec![field( + "amount", + struct_type(vec![field("value", DataType::Utf8)]), + )]), + ); + assert_map_planning_runtime_error( + &source_col, + &incompatible_target, + "Cannot cast struct field 'amount'", + ); + } + #[test] fn test_validate_dictionary_value_evolution() { let source_inner = struct_type(vec![field("a", DataType::Int32)]); @@ -1679,6 +2485,10 @@ mod tests { &DataType::Dictionary(Box::new(DataType::Int32), Box::new(s1.clone())), &DataType::Dictionary(Box::new(DataType::Int32), Box::new(s2.clone())), )); + assert!(requires_nested_struct_cast( + &map_type(DataType::Int32, DataType::Utf8), + &map_type(DataType::Float64, DataType::Utf8), + )); assert!(requires_nested_struct_cast( &DataType::ListView(arc_field("item", s1.clone())), &DataType::ListView(arc_field("item", s2.clone())), diff --git a/datafusion/core/tests/parquet/expr_adapter.rs b/datafusion/core/tests/parquet/expr_adapter.rs index 535828fa29c2..3bee5e0070c5 100644 --- a/datafusion/core/tests/parquet/expr_adapter.rs +++ b/datafusion/core/tests/parquet/expr_adapter.rs @@ -19,9 +19,10 @@ use std::sync::Arc; use arrow::array::{ Array, ArrayRef, BooleanArray, FixedSizeListArray, Int32Array, Int64Array, - LargeListArray, ListArray, RecordBatch, StringArray, StructArray, record_batch, + LargeListArray, ListArray, MapArray, RecordBatch, StringArray, StructArray, + record_batch, }; -use arrow::buffer::OffsetBuffer; +use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::compute::concat_batches; use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; use bytes::{BufMut, BytesMut}; @@ -928,6 +929,176 @@ async fn test_struct_schema_evolution_projection_and_filter() -> Result<()> { Ok(()) } +fn map_value_struct_evolution_batch() -> Result { + let physical_value_fields: Fields = vec![ + Arc::new(Field::new("amount", DataType::Int32, false)), + Arc::new(Field::new("ignored", DataType::Utf8, true)), + ] + .into(); + let values = StructArray::new( + physical_value_fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![10])) as ArrayRef, + Arc::new(StringArray::from(vec!["x"])) as ArrayRef, + ], + None, + ); + let entry_fields: Fields = vec![ + Arc::new(Field::new("keys", DataType::Utf8, false)), + Arc::new(Field::new( + "values", + DataType::Struct(physical_value_fields), + true, + )), + ] + .into(); + let entries = StructArray::new( + entry_fields.clone(), + vec![ + Arc::new(StringArray::from(vec!["a"])) as ArrayRef, + Arc::new(values) as ArrayRef, + ], + None, + ); + let map = MapArray::new( + Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)), + OffsetBuffer::new(vec![0, 1, 1].into()), + entries, + Some(NullBuffer::from(vec![true, false])), + false, + ); + Ok(RecordBatch::try_from_iter(vec![ + ("row_id", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef), + ("attributes", Arc::new(map) as ArrayRef), + ])?) +} + +fn map_value_struct_table_schema(currency_nullable: bool) -> SchemaRef { + let target_value_fields: Fields = vec![ + Arc::new(Field::new("amount", DataType::Int64, false)), + Arc::new(Field::new("currency", DataType::Utf8, currency_nullable)), + ] + .into(); + let target_entries = Field::new( + "entries", + DataType::Struct( + vec![ + Arc::new(Field::new("key", DataType::Utf8, false)), + Arc::new(Field::new( + "value", + DataType::Struct(target_value_fields), + true, + )), + ] + .into(), + ), + false, + ); + Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::Int32, false), + Field::new( + "attributes", + DataType::Map(Arc::new(target_entries), false), + true, + ), + ])) +} + +fn map_value_struct_nullable_schema() -> SchemaRef { + map_value_struct_table_schema(true) +} + +fn map_value_struct_non_nullable_schema() -> SchemaRef { + map_value_struct_table_schema(false) +} + +async fn setup_map_value_struct_table( + table_path: &str, + table_schema: SchemaRef, +) -> Result { + let store = Arc::new(InMemory::new()) as Arc; + let file_path = format!("{table_path}/data.parquet"); + write_parquet( + map_value_struct_evolution_batch()?, + Arc::clone(&store), + &file_path, + ) + .await; + + let ctx = test_context(); + let table_url = format!("memory:///{table_path}/"); + register_memory_listing_table(&ctx, store, &table_url, table_schema).await; + Ok(ctx) +} + +#[tokio::test] +async fn test_map_value_struct_schema_evolution_end_to_end() -> Result<()> { + let ctx = + setup_map_value_struct_table("map_evolution", map_value_struct_nullable_schema()) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY row_id") + .await? + .collect() + .await?; + let map = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .expect("attributes should be a MapArray"); + assert!(map.is_valid(0)); + assert!(map.is_null(1)); + let (key_field, value_field) = map.entries_fields(); + assert_eq!(key_field.name(), "key"); + assert_eq!(value_field.name(), "value"); + let keys = map + .keys() + .as_any() + .downcast_ref::() + .expect("map keys should be a StringArray"); + assert_eq!(keys.value(0), "a"); + let values = map + .values() + .as_any() + .downcast_ref::() + .expect("map values should be a StructArray"); + let amounts = values + .column_by_name("amount") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(amounts.values(), &[10]); + assert_eq!(values.column_by_name("currency").unwrap().null_count(), 1); + assert!(values.column_by_name("ignored").is_none()); + + Ok(()) +} + +#[tokio::test] +async fn test_map_value_struct_incompatible_schema_evolution_rejected() -> Result<()> { + let ctx = setup_map_value_struct_table( + "map_evolution_rejected", + map_value_struct_non_nullable_schema(), + ) + .await?; + + let error = ctx + .sql("SELECT * FROM t ORDER BY row_id") + .await? + .collect() + .await + .unwrap_err() + .to_string(); + assert!( + error.contains("target field 'currency' is non-nullable"), + "unexpected error: {error}" + ); + + Ok(()) +} + /// Macro to generate schema evolution tests for list-like variants. macro_rules! test_struct_schema_evolution_variants { (