sunchao commented on code in PR #25342:
URL: https://github.com/apache/datafusion/pull/25342#discussion_r4028317248
##########
datafusion/datasource-parquet/src/schema_coercion.rs:
##########
@@ -51,158 +66,139 @@ pub fn apply_file_schema_type_coercions(
table_schema: &Schema,
file_schema: &Schema,
) -> Option<Schema> {
- let mut needs_view_transform = false;
- let mut needs_string_transform = false;
- let mut needs_nested_transform = false;
+ let fields = coerce_fields_by_name(table_schema.fields(),
file_schema.fields())?;
+ Some(Schema::new_with_metadata(
+ fields,
+ file_schema.metadata.clone(),
+ ))
+}
+/// Coerce `file_fields` towards `table_fields`, matching fields by name.
+///
+/// File fields with no counterpart in `table_fields` are kept unchanged and
+/// table fields missing from the file are ignored. Returns `None` if no field
+/// changed.
+fn coerce_fields_by_name(table_fields: &Fields, file_fields: &Fields) ->
Option<Fields> {
// Create a mapping of table field names to their data types for fast
lookup
- // and simultaneously check if we need any transformations
- let table_fields: HashMap<_, _> = table_schema
- .fields()
+ let table_types: HashMap<_, _> = table_fields
.iter()
- .map(|f| {
- let dt = f.data_type();
- // Check if we need view type transformation
- if matches!(dt, &DataType::Utf8View | &DataType::BinaryView) {
- needs_view_transform = true;
- }
- // Check if we need string type transformation
- if matches!(
- dt,
- &DataType::Utf8 | &DataType::LargeUtf8 | &DataType::Utf8View
- ) {
- needs_string_transform = true;
- }
- // Nested fields can need transformations even when their parent
does not.
- if matches!(
- dt,
- DataType::Struct(_)
- | DataType::List(_)
- | DataType::LargeList(_)
- | DataType::ListView(_)
- | DataType::LargeListView(_)
- | DataType::FixedSizeList(_, _)
- | DataType::Map(_, _)
- ) {
- needs_nested_transform = true;
- }
-
- (f.name(), dt)
- })
+ .map(|f| (f.name(), f.data_type()))
.collect();
- // Early return if no transformation needed
- if !needs_view_transform && !needs_string_transform &&
!needs_nested_transform {
- return None;
- }
+ coerce_fields(file_fields, |_, field| {
+ let table_type = table_types.get(field.name())?;
+ coerce_data_type(table_type, field.data_type())
+ .map(|new_type| field_with_new_type(field, new_type))
+ })
+}
- let transformed_fields: Vec<Arc<Field>> = file_schema
- .fields()
- .iter()
- .map(|field| {
- let field_name = field.name();
- let field_type = field.data_type();
-
- // Look up the corresponding field type in the table schema
- if let Some(table_type) = table_fields.get(field_name) {
- match (table_type, field_type) {
- // table schema uses string type, coerce the file schema
to use string type
- (
- &DataType::Utf8,
- DataType::Binary | DataType::LargeBinary |
DataType::BinaryView,
- ) => {
- return field_with_new_type(field, DataType::Utf8);
- }
- // table schema uses large string type, coerce the file
schema to use large string type
- (
- &DataType::LargeUtf8,
- DataType::Binary | DataType::LargeBinary |
DataType::BinaryView,
- ) => {
- return field_with_new_type(field, DataType::LargeUtf8);
- }
- // table schema uses string view type, coerce the file
schema to use view type
- (
- &DataType::Utf8View,
- DataType::Binary | DataType::LargeBinary |
DataType::BinaryView,
- ) => {
- return field_with_new_type(field, DataType::Utf8View);
- }
- // Handle view type conversions
- (&DataType::Utf8View, DataType::Utf8 |
DataType::LargeUtf8) => {
- return field_with_new_type(field, DataType::Utf8View);
- }
- (&DataType::BinaryView, DataType::Binary |
DataType::LargeBinary) => {
- return field_with_new_type(field,
DataType::BinaryView);
- }
- // Apply the same coercions to matching fields inside
structs.
- (DataType::Struct(table_fields),
DataType::Struct(file_fields)) => {
- if let Some(schema) = apply_file_schema_type_coercions(
- &Schema::new(table_fields.clone()),
- &Schema::new(file_fields.clone()),
- ) {
- return field_with_new_type(
- field,
- DataType::Struct(schema.fields),
- );
- }
- }
- // Container children match by position, regardless of
their names.
- (DataType::List(table_child), DataType::List(file_child))
- | (
- DataType::LargeList(table_child),
- DataType::LargeList(file_child),
- )
- | (DataType::ListView(table_child),
DataType::ListView(file_child))
- | (
- DataType::LargeListView(table_child),
- DataType::LargeListView(file_child),
- )
- | (
- DataType::FixedSizeList(table_child, _),
- DataType::FixedSizeList(file_child, _),
- )
- | (DataType::Map(table_child, _),
DataType::Map(file_child, _)) => {
- if let Some(schema) = apply_file_schema_type_coercions(
- &Schema::new(vec![field_with_new_type(
- file_child,
- table_child.data_type().clone(),
- )]),
- &Schema::new(vec![Arc::clone(file_child)]),
- ) {
- let child = Arc::clone(&schema.fields()[0]);
- let new_type = match field_type {
- DataType::List(_) => DataType::List(child),
- DataType::LargeList(_) =>
DataType::LargeList(child),
- DataType::ListView(_) =>
DataType::ListView(child),
- DataType::LargeListView(_) => {
- DataType::LargeListView(child)
- }
- DataType::FixedSizeList(_, size) => {
- DataType::FixedSizeList(child, *size)
- }
- DataType::Map(_, sorted) =>
DataType::Map(child, *sorted),
- _ => return Arc::clone(field),
- };
- return field_with_new_type(field, new_type);
- }
- }
- _ => {}
+/// Rebuild `file_fields`, replacing every field for which `coerce` returns a
+/// new one. Returns `None` if no field changed.
+///
+/// The output is only allocated once a field actually changes, so schemas
+/// needing no coercion at all (the common case) are walked without allocating
+/// or touching the reference counts of the file fields.
+fn coerce_fields(
+ file_fields: &Fields,
+ mut coerce: impl FnMut(usize, &FieldRef) -> Option<FieldRef>,
+) -> Option<Fields> {
+ let mut coerced: Option<Vec<FieldRef>> = None;
+ for (idx, field) in file_fields.iter().enumerate() {
+ match coerce(idx, field) {
+ Some(new_field) => coerced
+ .get_or_insert_with(|| {
+ // The fields before the first change are carried over as
is
+ let mut fields = Vec::with_capacity(file_fields.len());
+ fields.extend_from_slice(&file_fields[..idx]);
+ fields
+ })
+ .push(new_field),
+ // Unchanged fields are only copied once something else changed
+ None => {
+ if let Some(coerced) = &mut coerced {
+ coerced.push(Arc::clone(field));
}
}
+ }
+ }
- // If no transformation is needed, keep the original field
- Arc::clone(field)
- })
- .collect();
+ coerced.map(Fields::from)
+}
+
+/// Coerce `file_type` towards `table_type`, recursing into nested types.
+///
+/// Returns the new type for the file field, or `None` if no transformation
+/// is needed (including when the two types are unrelated).
+fn coerce_data_type(table_type: &DataType, file_type: &DataType) ->
Option<DataType> {
+ use DataType::*;
+ match (table_type, file_type) {
+ // table schema uses string type, coerce the file schema to use string
type
+ (Utf8, Binary | LargeBinary | BinaryView) => Some(Utf8),
+ // table schema uses large string type, coerce the file schema to use
large string type
+ (LargeUtf8, Binary | LargeBinary | BinaryView) => Some(LargeUtf8),
+ // table schema uses string view type, coerce the file schema to use
view type
+ (Utf8View, Binary | LargeBinary | BinaryView | Utf8 | LargeUtf8) => {
+ Some(Utf8View)
+ }
+ (BinaryView, Binary | LargeBinary) => Some(BinaryView),
+ // Struct children match by name
+ (Struct(table_fields), Struct(file_fields)) => {
+ coerce_fields_by_name(table_fields, file_fields).map(Struct)
+ }
+ // List-like children match by position, regardless of their names.
+ // The container kind and FixedSizeList width always come from the
file.
+ (List(table_child), List(file_child)) => {
+ coerce_child(table_child, file_child).map(List)
+ }
+ (LargeList(table_child), LargeList(file_child)) => {
+ coerce_child(table_child, file_child).map(LargeList)
+ }
+ (ListView(table_child), ListView(file_child)) => {
+ coerce_child(table_child, file_child).map(ListView)
+ }
+ (LargeListView(table_child), LargeListView(file_child)) => {
+ coerce_child(table_child, file_child).map(LargeListView)
+ }
+ (FixedSizeList(table_child, _), FixedSizeList(file_child, size)) => {
+ coerce_child(table_child, file_child).map(|child|
FixedSizeList(child, *size))
+ }
+ // Map keys and values match by position: Parquet always names them
+ // `key`/`value` while Arrow producers commonly use `keys`/`values`.
+ (Map(table_entries, _), Map(file_entries, sorted)) => {
+ coerce_map_entries(table_entries, file_entries)
+ .map(|entries| Map(entries, *sorted))
+ }
+ _ => None,
+ }
+}
+
+/// Coerce a single nested child field, keeping everything but its data type
+/// from `file_child`.
+fn coerce_child(table_child: &FieldRef, file_child: &FieldRef) ->
Option<FieldRef> {
+ coerce_data_type(table_child.data_type(), file_child.data_type())
+ .map(|new_type| field_with_new_type(file_child, new_type))
+}
- if transformed_fields.iter().eq(file_schema.fields().iter()) {
+/// Coerce the `entries` struct of a [`DataType::Map`], matching the key and
+/// value children by position.
+fn coerce_map_entries(
+ table_entries: &FieldRef,
+ file_entries: &FieldRef,
+) -> Option<FieldRef> {
+ let (DataType::Struct(table_fields), DataType::Struct(file_fields)) =
+ (table_entries.data_type(), file_entries.data_type())
+ else {
+ return None;
+ };
+ if table_fields.len() != file_fields.len() {
return None;
}
- Some(Schema::new_with_metadata(
- transformed_fields,
- file_schema.metadata.clone(),
- ))
+ let fields = coerce_fields(file_fields, |idx, file_child| {
+ coerce_child(&table_fields[idx], file_child)
Review Comment:
[P1] Preserve UTF-8 validation for binary map children
This also enables `Binary` → string decoding for renamed map children,
exposing an existing parquet-rs validation limitation to reads that previously
used a validating cast. With Arrow/Parquet 59.3.0, I reproduced this using a
file map named `key_value/key/value`, one key `"k"` and binary value `[0xff]`,
and a table map named `entries/keys/values` with `Utf8View` values (write
without embedded Arrow metadata, as in the new reader test).
On base, coercion returns `None`, decoding produces valid binary data, and
the strict map cast returns `Encountered non UTF-8 data`. On this head, the
reader returns `Ok` with an invalid `StringViewArray`; the subsequent strict
cast to the table map also succeeds, and its `to_data().validate_full()` still
fails with `Encountered non-UTF-8 data at index 0`. The decoder decides whether
to validate from the physical Parquet UTF8 annotation, which a binary column
lacks, then constructs the overridden string-view array unchecked.
`Utf8`/`LargeUtf8` targets instead panic in the debug reader. Valid multibyte
UTF-8 passes on both revisions.
Could we retain a validating conversion for binary-to-string map children,
or ensure reader validation before enabling that conversion, and add this
invalid-byte regression? The `Utf8` → `Utf8View` optimization can still decode
directly.
--
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]
