Skip to content

Commit e2efece

Browse files
adragomiradragomir
committed
[HSTACK] fix: fix: make Parquet file schema type coercion work on nested schemas
- Upstream PR / issue: apache#25192, https://github.com/apache/datafusion/pull/25193/changes
1 parent 582c831 commit e2efece

1 file changed

Lines changed: 151 additions & 0 deletions

File tree

‎datafusion/datasource-parquet/src/schema_coercion.rs‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@ pub fn apply_file_schema_type_coercions(
5151
table_schema: &Schema,
5252
file_schema: &Schema,
5353
) -> Option<Schema> {
54+
return apply_file_schema_type_coercions_nested(table_schema, file_schema);
55+
5456
let mut needs_view_transform = false;
5557
let mut needs_string_transform = false;
5658

@@ -129,6 +131,155 @@ pub fn apply_file_schema_type_coercions(
129131
})
130132
.collect();
131133

134+
Some(Schema::new_with_metadata(
135+
transformed_fields,
136+
file_schema.metadata.clone(),
137+
))
138+
//
139+
}
140+
141+
/// Apply string and binary type coercions recursively within structs, lists, and maps.
142+
///
143+
/// Struct fields match by name. List elements and map entries retain their physical
144+
/// container types. File field order, nullability, and metadata remain unchanged.
145+
/// Missing fields are not added; expression adaptation handles them separately.
146+
/// Returns `None` if no file field types change.
147+
pub fn apply_file_schema_type_coercions_nested(
148+
table_schema: &Schema,
149+
file_schema: &Schema,
150+
) -> Option<Schema> {
151+
let mut needs_view_transform = false;
152+
let mut needs_string_transform = false;
153+
let mut needs_nested_transform = false;
154+
155+
// Create a mapping of table field names to their data types for fast lookup
156+
// and simultaneously check if we need any transformations
157+
let table_fields: HashMap<_, _> = table_schema
158+
.fields()
159+
.iter()
160+
.map(|f| {
161+
let dt = f.data_type();
162+
// Check if we need view type transformation
163+
if matches!(dt, &DataType::Utf8View | &DataType::BinaryView) {
164+
needs_view_transform = true;
165+
}
166+
// Check if we need string type transformation
167+
if matches!(
168+
dt,
169+
&DataType::Utf8 | &DataType::LargeUtf8 | &DataType::Utf8View
170+
) {
171+
needs_string_transform = true;
172+
}
173+
// Nested fields can need transformations even when their parent does not.
174+
if matches!(
175+
dt,
176+
DataType::Struct(_)
177+
| DataType::List(_)
178+
| DataType::LargeList(_)
179+
| DataType::ListView(_)
180+
| DataType::LargeListView(_)
181+
| DataType::FixedSizeList(_, _)
182+
| DataType::Map(_, _)
183+
) {
184+
needs_nested_transform = true;
185+
}
186+
187+
(f.name(), dt)
188+
})
189+
.collect();
190+
191+
// Early return if no transformation needed
192+
if !needs_view_transform && !needs_string_transform && !needs_nested_transform {
193+
return None;
194+
}
195+
196+
let transformed_fields: Vec<Arc<Field>> = file_schema
197+
.fields()
198+
.iter()
199+
.map(|field| {
200+
let field_name = field.name();
201+
let field_type = field.data_type();
202+
203+
// Look up the corresponding field type in the table schema
204+
if let Some(table_type) = table_fields.get(field_name) {
205+
match (table_type, field_type) {
206+
// table schema uses string type, coerce the file schema to use string type
207+
(
208+
&DataType::Utf8,
209+
DataType::Binary | DataType::LargeBinary | DataType::BinaryView,
210+
) => {
211+
return field_with_new_type(field, DataType::Utf8);
212+
}
213+
// table schema uses large string type, coerce the file schema to use large string type
214+
(
215+
&DataType::LargeUtf8,
216+
DataType::Binary | DataType::LargeBinary | DataType::BinaryView,
217+
) => {
218+
return field_with_new_type(field, DataType::LargeUtf8);
219+
}
220+
// table schema uses string view type, coerce the file schema to use view type
221+
(
222+
&DataType::Utf8View,
223+
DataType::Binary | DataType::LargeBinary | DataType::BinaryView,
224+
) => {
225+
return field_with_new_type(field, DataType::Utf8View);
226+
}
227+
// Handle view type conversions
228+
(&DataType::Utf8View, DataType::Utf8 | DataType::LargeUtf8) => {
229+
return field_with_new_type(field, DataType::Utf8View);
230+
}
231+
(&DataType::BinaryView, DataType::Binary | DataType::LargeBinary) => {
232+
return field_with_new_type(field, DataType::BinaryView);
233+
}
234+
// Apply the same coercions to matching fields inside structs.
235+
(DataType::Struct(table_fields), DataType::Struct(file_fields)) => {
236+
if let Some(schema) = apply_file_schema_type_coercions_nested(
237+
&Schema::new(table_fields.clone()),
238+
&Schema::new(file_fields.clone()),
239+
) {
240+
return field_with_new_type(field, DataType::Struct(schema.fields));
241+
}
242+
}
243+
// Container children match by position, regardless of their names.
244+
(DataType::List(table_child), DataType::List(file_child))
245+
| (DataType::LargeList(table_child), DataType::LargeList(file_child))
246+
| (DataType::ListView(table_child), DataType::ListView(file_child))
247+
| (DataType::LargeListView(table_child), DataType::LargeListView(file_child))
248+
| (DataType::FixedSizeList(table_child, _), DataType::FixedSizeList(file_child, _))
249+
| (DataType::Map(table_child, _), DataType::Map(file_child, _)) => {
250+
if let Some(schema) = apply_file_schema_type_coercions_nested(
251+
&Schema::new(vec![field_with_new_type(
252+
file_child,
253+
table_child.data_type().clone(),
254+
)]),
255+
&Schema::new(vec![Arc::clone(file_child)]),
256+
) {
257+
let child = Arc::clone(&schema.fields()[0]);
258+
let new_type = match field_type {
259+
DataType::List(_) => DataType::List(child),
260+
DataType::LargeList(_) => DataType::LargeList(child),
261+
DataType::ListView(_) => DataType::ListView(child),
262+
DataType::LargeListView(_) => DataType::LargeListView(child),
263+
DataType::FixedSizeList(_, size) => DataType::FixedSizeList(child, *size),
264+
DataType::Map(_, sorted) => DataType::Map(child, *sorted),
265+
_ => return Arc::clone(field),
266+
};
267+
return field_with_new_type(field, new_type);
268+
}
269+
}
270+
_ => {}
271+
}
272+
}
273+
274+
// If no transformation is needed, keep the original field
275+
Arc::clone(field)
276+
})
277+
.collect();
278+
279+
if transformed_fields.iter().eq(file_schema.fields().iter()) {
280+
return None;
281+
}
282+
132283
Some(Schema::new_with_metadata(
133284
transformed_fields,
134285
file_schema.metadata.clone(),

0 commit comments

Comments
 (0)