Skip to content

Commit b56d0d8

Browse files
adragomirSabin Floares
adragomir
authored and
Sabin Floares
committed
[HSTACK] fix: make Parquet file schema type coercion work on nested schemas
- Upstream PR / issue: apache#25192, apache#25193 - Backported upstream merge 3a6a18c in place of the main-54 draft - Ported from main-54 e2efece
1 parent 7822cbd commit b56d0d8

1 file changed

Lines changed: 315 additions & 1 deletion

File tree

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

Lines changed: 315 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ pub fn apply_file_schema_type_coercions(
5353
) -> Option<Schema> {
5454
let mut needs_view_transform = false;
5555
let mut needs_string_transform = false;
56+
let mut needs_nested_transform = false;
5657

5758
// Create a mapping of table field names to their data types for fast lookup
5859
// and simultaneously check if we need any transformations
@@ -72,13 +73,26 @@ pub fn apply_file_schema_type_coercions(
7273
) {
7374
needs_string_transform = true;
7475
}
76+
// Nested fields can need transformations even when their parent does not.
77+
if matches!(
78+
dt,
79+
DataType::Struct(_)
80+
| DataType::List(_)
81+
| DataType::LargeList(_)
82+
| DataType::ListView(_)
83+
| DataType::LargeListView(_)
84+
| DataType::FixedSizeList(_, _)
85+
| DataType::Map(_, _)
86+
) {
87+
needs_nested_transform = true;
88+
}
7589

7690
(f.name(), dt)
7791
})
7892
.collect();
7993

8094
// Early return if no transformation needed
81-
if !needs_view_transform && !needs_string_transform {
95+
if !needs_view_transform && !needs_string_transform && !needs_nested_transform {
8296
return None;
8397
}
8498

@@ -120,6 +134,58 @@ pub fn apply_file_schema_type_coercions(
120134
(&DataType::BinaryView, DataType::Binary | DataType::LargeBinary) => {
121135
return field_with_new_type(field, DataType::BinaryView);
122136
}
137+
// Apply the same coercions to matching fields inside structs.
138+
(DataType::Struct(table_fields), DataType::Struct(file_fields)) => {
139+
if let Some(schema) = apply_file_schema_type_coercions(
140+
&Schema::new(table_fields.clone()),
141+
&Schema::new(file_fields.clone()),
142+
) {
143+
return field_with_new_type(
144+
field,
145+
DataType::Struct(schema.fields),
146+
);
147+
}
148+
}
149+
// Container children match by position, regardless of their names.
150+
(DataType::List(table_child), DataType::List(file_child))
151+
| (
152+
DataType::LargeList(table_child),
153+
DataType::LargeList(file_child),
154+
)
155+
| (DataType::ListView(table_child), DataType::ListView(file_child))
156+
| (
157+
DataType::LargeListView(table_child),
158+
DataType::LargeListView(file_child),
159+
)
160+
| (
161+
DataType::FixedSizeList(table_child, _),
162+
DataType::FixedSizeList(file_child, _),
163+
)
164+
| (DataType::Map(table_child, _), DataType::Map(file_child, _)) => {
165+
if let Some(schema) = apply_file_schema_type_coercions(
166+
&Schema::new(vec![field_with_new_type(
167+
file_child,
168+
table_child.data_type().clone(),
169+
)]),
170+
&Schema::new(vec![Arc::clone(file_child)]),
171+
) {
172+
let child = Arc::clone(&schema.fields()[0]);
173+
let new_type = match field_type {
174+
DataType::List(_) => DataType::List(child),
175+
DataType::LargeList(_) => DataType::LargeList(child),
176+
DataType::ListView(_) => DataType::ListView(child),
177+
DataType::LargeListView(_) => {
178+
DataType::LargeListView(child)
179+
}
180+
DataType::FixedSizeList(_, size) => {
181+
DataType::FixedSizeList(child, *size)
182+
}
183+
DataType::Map(_, sorted) => DataType::Map(child, *sorted),
184+
_ => return Arc::clone(field),
185+
};
186+
return field_with_new_type(field, new_type);
187+
}
188+
}
123189
_ => {}
124190
}
125191
}
@@ -129,6 +195,10 @@ pub fn apply_file_schema_type_coercions(
129195
})
130196
.collect();
131197

198+
if transformed_fields.iter().eq(file_schema.fields().iter()) {
199+
return None;
200+
}
201+
132202
Some(Schema::new_with_metadata(
133203
transformed_fields,
134204
file_schema.metadata.clone(),
@@ -470,6 +540,250 @@ mod tests {
470540

471541
use parquet::schema::parser::parse_message_type;
472542

543+
#[test]
544+
fn nested_coercion_preserves_file_schema() {
545+
let field_metadata = HashMap::from([("source".into(), "file".into())]);
546+
let identifiers = Field::new_struct(
547+
"identifiers",
548+
vec![Field::new("cid", DataType::Utf8, true)],
549+
true,
550+
)
551+
.with_metadata(field_metadata.clone());
552+
let struct_field = Field::new_struct(
553+
"_struct",
554+
vec![
555+
Field::new("file_only", DataType::Binary, true),
556+
identifiers.clone(),
557+
],
558+
true,
559+
)
560+
.with_metadata(field_metadata.clone());
561+
let file_schema = Schema::new_with_metadata(
562+
vec![
563+
Field::new("id", DataType::Int64, false),
564+
struct_field.clone(),
565+
],
566+
field_metadata.clone(),
567+
);
568+
let table_schema = Schema::new(vec![
569+
Field::new_struct(
570+
"_struct",
571+
vec![
572+
Field::new_struct(
573+
"identifiers",
574+
vec![Field::new("cid", DataType::Utf8View, false)],
575+
false,
576+
),
577+
Field::new("schemaEvolutionOnly", DataType::Utf8View, true),
578+
],
579+
false,
580+
),
581+
Field::new("missing", DataType::Utf8View, true),
582+
Field::new("id", DataType::Int64, true),
583+
]);
584+
let expected = Schema::new_with_metadata(
585+
vec![
586+
Field::new("id", DataType::Int64, false),
587+
struct_field.with_data_type(DataType::Struct(
588+
vec![
589+
Field::new("file_only", DataType::Binary, true),
590+
identifiers.with_data_type(DataType::Struct(
591+
vec![Field::new("cid", DataType::Utf8View, true)].into(),
592+
)),
593+
]
594+
.into(),
595+
)),
596+
],
597+
field_metadata,
598+
);
599+
600+
let result = apply_file_schema_type_coercions(&table_schema, &file_schema);
601+
assert_eq!(result, Some(expected.clone()));
602+
assert_eq!(
603+
apply_file_schema_type_coercions(&table_schema, &expected),
604+
None
605+
);
606+
}
607+
608+
#[test]
609+
fn nested_coercion_preserves_list_containers() {
610+
let wrap: [fn(FieldRef) -> DataType; 6] = [
611+
DataType::List,
612+
DataType::LargeList,
613+
DataType::ListView,
614+
DataType::LargeListView,
615+
|field| DataType::FixedSizeList(field, 3),
616+
|field| DataType::Map(field, false),
617+
];
618+
for wrap in wrap {
619+
let file_element = Arc::new(Field::new_struct(
620+
"entries",
621+
vec![
622+
Field::new("key", DataType::Utf8, false),
623+
Field::new("value", DataType::Binary, true),
624+
],
625+
false,
626+
));
627+
let table_element = Arc::new(Field::new_struct(
628+
"entries",
629+
vec![
630+
Field::new("value", DataType::Utf8View, true),
631+
Field::new("key", DataType::Utf8View, false),
632+
],
633+
false,
634+
));
635+
let expected_element = Arc::new(Field::new_struct(
636+
"entries",
637+
vec![
638+
Field::new("key", DataType::Utf8View, false),
639+
Field::new("value", DataType::Utf8View, true),
640+
],
641+
false,
642+
));
643+
let file_schema =
644+
Schema::new(vec![Field::new("data", wrap(file_element), true)]);
645+
let table_schema =
646+
Schema::new(vec![Field::new("data", wrap(table_element), true)]);
647+
let expected =
648+
Schema::new(vec![Field::new("data", wrap(expected_element), true)]);
649+
assert_eq!(
650+
apply_file_schema_type_coercions(&table_schema, &file_schema),
651+
Some(expected)
652+
);
653+
}
654+
}
655+
656+
#[test]
657+
fn nested_coercion_preserves_leaf_conversion_rules() {
658+
let nested_schema = |data_type| {
659+
Schema::new(vec![Field::new_struct(
660+
"nested",
661+
vec![Field::new("value", data_type, true)],
662+
true,
663+
)])
664+
};
665+
for (table_type, file_types) in [
666+
(
667+
DataType::Utf8,
668+
vec![
669+
DataType::Binary,
670+
DataType::LargeBinary,
671+
DataType::BinaryView,
672+
],
673+
),
674+
(
675+
DataType::LargeUtf8,
676+
vec![
677+
DataType::Binary,
678+
DataType::LargeBinary,
679+
DataType::BinaryView,
680+
],
681+
),
682+
(
683+
DataType::Utf8View,
684+
vec![
685+
DataType::Binary,
686+
DataType::LargeBinary,
687+
DataType::BinaryView,
688+
DataType::Utf8,
689+
DataType::LargeUtf8,
690+
],
691+
),
692+
(
693+
DataType::BinaryView,
694+
vec![DataType::Binary, DataType::LargeBinary],
695+
),
696+
] {
697+
for file_type in file_types {
698+
let file_schema = nested_schema(file_type);
699+
let table_schema = nested_schema(table_type.clone());
700+
assert_eq!(
701+
apply_file_schema_type_coercions(&table_schema, &file_schema),
702+
Some(table_schema)
703+
);
704+
}
705+
}
706+
let file_schema = nested_schema(DataType::Int32);
707+
let table_schema = nested_schema(DataType::Int64);
708+
assert_eq!(
709+
apply_file_schema_type_coercions(&table_schema, &file_schema),
710+
None
711+
);
712+
}
713+
714+
#[test]
715+
fn nested_coercion_reader_produces_string_views() {
716+
use arrow::array::{Array, ArrayRef, StringArray, StringViewArray, StructArray};
717+
use arrow::record_batch::RecordBatch;
718+
use bytes::Bytes;
719+
use parquet::arrow::ArrowWriter;
720+
use parquet::arrow::arrow_reader::{
721+
ArrowReaderOptions, ParquetRecordBatchReaderBuilder,
722+
};
723+
724+
let identifiers = StructArray::from(vec![(
725+
Arc::new(Field::new("cid", DataType::Utf8, true)),
726+
Arc::new(StringArray::from(vec![Some("customer-03"), None])) as ArrayRef,
727+
)]);
728+
let struct_field = StructArray::from(vec![(
729+
Arc::new(Field::new(
730+
"identifiers",
731+
identifiers.data_type().clone(),
732+
true,
733+
)),
734+
Arc::new(identifiers) as ArrayRef,
735+
)]);
736+
let batch =
737+
RecordBatch::try_from_iter([("_struct", Arc::new(struct_field) as ArrayRef)])
738+
.unwrap();
739+
let table_schema = Schema::new(vec![Field::new_struct(
740+
"_struct",
741+
vec![
742+
Field::new_struct(
743+
"identifiers",
744+
vec![Field::new("cid", DataType::Utf8View, true)],
745+
true,
746+
),
747+
Field::new("schemaEvolutionOnly", DataType::Utf8View, true),
748+
],
749+
true,
750+
)]);
751+
let reader_schema = Arc::new(
752+
apply_file_schema_type_coercions(&table_schema, &batch.schema()).unwrap(),
753+
);
754+
let mut bytes = Vec::new();
755+
let mut writer = ArrowWriter::try_new(&mut bytes, batch.schema(), None).unwrap();
756+
writer.write(&batch).unwrap();
757+
writer.close().unwrap();
758+
let mut reader = ParquetRecordBatchReaderBuilder::try_new_with_options(
759+
Bytes::from(bytes),
760+
ArrowReaderOptions::new().with_schema(reader_schema),
761+
)
762+
.unwrap()
763+
.build()
764+
.unwrap();
765+
let decoded = reader.next().unwrap().unwrap();
766+
let struct_field = decoded
767+
.column(0)
768+
.as_any()
769+
.downcast_ref::<StructArray>()
770+
.unwrap();
771+
let identifiers = struct_field
772+
.column(0)
773+
.as_any()
774+
.downcast_ref::<StructArray>()
775+
.unwrap();
776+
let cust_id = identifiers
777+
.column(0)
778+
.as_any()
779+
.downcast_ref::<StringViewArray>()
780+
.unwrap();
781+
assert_eq!(cust_id.value(0), "customer-03");
782+
assert!(cust_id.is_null(1));
783+
assert_eq!(struct_field.num_columns(), 1);
784+
assert!(reader.next().is_none());
785+
}
786+
473787
#[test]
474788
fn coerce_int96_to_resolution_with_mixed_timestamps() {
475789
// Unclear if Spark (or other writer) could generate a file with mixed timestamps like this,

0 commit comments

Comments
 (0)