kumarUjjawal commented on code in PR #24227:
URL: https://github.com/apache/datafusion/pull/24227#discussion_r4119051027
##########
datafusion/datasource-parquet/src/schema_coercion.rs:
##########
@@ -84,6 +84,183 @@ pub fn apply_file_schema_type_coercions(
))
}
+/// Like [`apply_file_schema_type_coercions`], but also coerces compatible
+/// string/binary file fields to dictionary types already present in the table
+/// schema.
+pub(crate) fn apply_file_schema_type_coercions_with_rle(
+ table_schema: &Schema,
+ file_schema: &Schema,
+ enable_rle_to_dictionary: bool,
+) -> Option<Schema> {
+ let mut needs_view_transform = false;
+ let mut needs_string_transform = false;
+ let mut needs_nested_transform = false;
+ let mut needs_dict_transform = false;
+
+ // 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()
+ .iter()
+ .map(|field| {
+ let data_type = field.data_type();
+ // Check if we need view type transformation
+ if matches!(data_type, &DataType::Utf8View |
&DataType::BinaryView) {
+ needs_view_transform = true;
+ }
+ // Check if we need string type transformation
+ if matches!(
+ data_type,
+ &DataType::Utf8 | &DataType::LargeUtf8 | &DataType::Utf8View
+ ) {
+ needs_string_transform = true;
+ }
+ // Nested fields can need transformations even when their parent
does not.
+ if matches!(
+ data_type,
+ DataType::Struct(_)
+ | DataType::List(_)
+ | DataType::LargeList(_)
+ | DataType::ListView(_)
+ | DataType::LargeListView(_)
+ | DataType::FixedSizeList(_, _)
+ | DataType::Map(_, _)
+ ) {
+ needs_nested_transform = true;
+ }
+ if enable_rle_to_dictionary
+ && matches!(data_type, &DataType::Dictionary(_, _))
+ {
+ needs_dict_transform = true;
+ }
+
+ (field.name(), data_type)
+ })
+ .collect();
+
+ // Early return if no transformation needed
+ if !needs_view_transform
+ && !needs_string_transform
+ && !needs_nested_transform
+ && !needs_dict_transform
+ {
+ return None;
+ }
+
+ let 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, _)) => {
Review Comment:
Please keep the existing map-specific coercion here. This call treats map
entries as structs and matches their children by name. Parquet often uses `key`
and `value`, while the table schema can use `keys` and `values`. The current
`coerce_map_entries` helper matches them by position. It also keeps binary map
values on the validating cast path. This function runs when
`enable_rle_to_dictionary` is false, so the rebase changes existing behavior.
Please reuse the existing helper pipeline and add dictionary promotion 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]
