adriangb commented on code in PR #25342:
URL: https://github.com/apache/datafusion/pull/25342#discussion_r4054183278
##########
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:
Filed both halves:
- parquet-rs: https://github.com/apache/arrow-rs/issues/11140.
`ByteArrayColumnValueDecoder::new` takes the UTF-8 validation decision from the
physical annotation only (`desc.converted_type() == ConvertedType::UTF8`), so a
`with_schema` override from `Binary` to `Utf8`, `LargeUtf8` or `Utf8View`
returns an unvalidated array. `OffsetBuffer::into_array` then uses
`build_unchecked` in release builds. The issue has a standalone parquet-only
reproducer.
- DataFusion: https://github.com/apache/datafusion/issues/25509. The
segfault above, with the release and debug results for `Binary`, `LargeBinary`
and `BinaryView`.
To be explicit about the scope: this is pre-existing on `main` for top level
columns, so it is not introduced here. `Utf8View` is the worse of the two
paths, because the debug check in `into_array` does not apply to it.
--
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]
