adriangb commented on code in PR #23914:
URL: https://github.com/apache/datafusion/pull/23914#discussion_r3814518790
##########
datafusion/common/src/nested_struct.rs:
##########
@@ -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::<MapArray>().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() {
Review Comment:
Can we add positive allowlist tests? E.g. `Map(Int32,V)→Map(Int64,V)`
##########
datafusion/common/src/nested_struct.rs:
##########
@@ -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)?;
+ }
Review Comment:
This is dropping `field_name` so the error will be of the form `Cannot
change Map sorted flag during schema adaptation` which does not include a
column name. Other errors (e.g. `Dictionary` below) do include it.
##########
datafusion/common/src/nested_struct.rs:
##########
@@ -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<ArrayRef> {
+ 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);
Review Comment:
This can panic w/ `Map(Utf8, non-null Utf8) → Map(Utf8, non-null Int32)` and
`CastOptions { safe: true }`. Can we use
[`try_new(...)?`](https://docs.rs/arrow/latest/arrow/array/struct.StructArray.html#method.try_new)
instead?
The same issue already exists in `cast_struct_column` (pre existing), might
be nice to fix at the same time:
https://github.com/apache/datafusion/blob/c429919c7c8bd81e4e687531a609e2380ed60649/datafusion/common/src/nested_struct.rs#L112-L113
##########
datafusion/common/src/nested_struct.rs:
##########
@@ -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<ArrayRef> {
+ 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<MapArray> {
+ let indices = UInt64Array::from_iter_values(0..map.len() as u64);
+ Ok(take(map, &indices, None)?.as_map().clone())
+}
Review Comment:
Worth noting that this makes Map the only container here that compacts.
`cast_list_column` reuses the source offsets and casts the full backing child,
and `cast_fixed_size_list_column` masks-and-retries only after a failure. The
three containers now handle slice-hidden child data three different ways.
The approach here for `Map` seems like the right approach. `MapArray::slice`
keeps entries whole and `keys()/values()` return the full backing children, so
without this the cast would process every backing entry regardless of the
slice. Measured rows=2 entries=20 for `SELECT CAST(m AS MAP(BIGINT, VARCHAR))
FROM t LIMIT 2`, so compaction is a saving, not an overhead.
The reason to raise it is that the list path has the same exposure and no
protection. On current main:
```sql
CREATE TABLE lt AS SELECT * FROM (VALUES
(1, [struct('1')]), (2, [struct('2')]), (3, [struct('bad')])
) AS t(i, l);
SELECT arrow_cast(l, 'List(Struct("c0": Int32))') FROM lt LIMIT 2;
-- Cast error: Cannot cast string 'bad' to value of Int32 type
```
Row 3 is excluded by the LIMIT, but the cast still sees it. I suggest we
file a tracking issue and defer, but a comment here saying why Map compacts and
the list types don't would stop it reading as an oversight.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
