Skip to content

Commit ec74746

Browse files
dwsmith1983ryzhyk
andauthored
[branch-55] Preserve field metadata when coercing INT96 timestamps (#24790) (#25895)
## Which issue does this PR close? - Backport of #24790 to `branch-55`, for the 55.2.0 patch release (#25758). - Closes #24786 on `branch-55`. ## Rationale for this change With `datafusion.execution.parquet.coerce_int96` set, a file with an INT96 column comes back with the metadata of every struct, list and map field emptied, while leaf fields keep theirs. Readers that match fields by Parquet field id lose the id of any container whose id sits only on the container. Apache DataFusion Comet always enables the coercion, so with Spark's field id reads on it null fills such a container where Spark reads it (apache/datafusion-comet#6131). ## What changes are included in this PR? A clean cherry-pick of #24790 (with `-x`): the three container constructors in `schema_coercion.rs` keep the original field's metadata, plus that PR's tests. ## What is the testing strategy for this PR? The tests from #24790 come with the cherry-pick. The `datafusion-datasource-parquet` crate's tests, `cargo fmt` and `cargo clippy` pass on this branch. ## Are there any user-facing changes? Struct, list and map fields read from files with INT96 columns keep their field metadata when `coerce_int96` is set. No API changes. Signed-off-by: Leonid Ryzhyk <[email protected]> Co-authored-by: Leonid Ryzhyk <[email protected]>
1 parent c1514b5 commit ec74746

1 file changed

Lines changed: 147 additions & 3 deletions

File tree

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

Lines changed: 147 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -328,7 +328,8 @@ fn coerce_int96_to_resolution_impl(
328328
current_field.name(),
329329
processed_children.as_slice(),
330330
current_field.is_nullable(),
331-
);
331+
)
332+
.with_metadata(current_field.metadata().clone());
332333
parent_fields.borrow_mut().push(Arc::new(processed_struct));
333334
}
334335
(DataType::List(unprocessed_child), None) => {
@@ -361,7 +362,8 @@ fn coerce_int96_to_resolution_impl(
361362
current_field.name(),
362363
Arc::clone(&processed_children[0]),
363364
current_field.is_nullable(),
364-
);
365+
)
366+
.with_metadata(current_field.metadata().clone());
365367
parent_fields.borrow_mut().push(Arc::new(processed_list));
366368
}
367369
(DataType::Map(unprocessed_child, _), None) => {
@@ -391,7 +393,8 @@ fn coerce_int96_to_resolution_impl(
391393
current_field.name(),
392394
DataType::Map(Arc::clone(&processed_children[0]), *sorted),
393395
current_field.is_nullable(),
394-
);
396+
)
397+
.with_metadata(current_field.metadata().clone());
395398
parent_fields.borrow_mut().push(Arc::new(processed_map));
396399
}
397400
(DataType::Timestamp(TimeUnit::Nanosecond, None), None)
@@ -460,6 +463,7 @@ pub fn transform_binary_to_string(schema: &Schema) -> Schema {
460463
}
461464
#[cfg(test)]
462465
mod tests {
466+
use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
463467
use parquet::arrow::parquet_to_arrow_schema;
464468

465469
use super::*;
@@ -728,4 +732,144 @@ mod tests {
728732

729733
assert_eq!(result, expected_schema);
730734
}
735+
736+
#[test]
737+
fn coerce_int96_to_resolution_preserves_field_metadata() {
738+
// Spark and Delta Lake stamp a field id on every field. Coercion must
739+
// carry that metadata across on nested fields as well as leaves:
740+
// formats that identify a column by id rather than by name, such as
741+
// Delta Lake column mapping, cannot resolve a column that loses it.
742+
let spark_schema = "
743+
message spark_schema {
744+
REQUIRED INT64 c0 = 1;
745+
OPTIONAL group c1 = 2 {
746+
OPTIONAL INT96 c2 = 3;
747+
}
748+
OPTIONAL group c3 (LIST) = 4 {
749+
REPEATED group list {
750+
OPTIONAL INT96 element = 5;
751+
}
752+
}
753+
OPTIONAL group c4 (MAP) = 6 {
754+
REPEATED group key_value {
755+
REQUIRED BYTE_ARRAY key (STRING) = 7;
756+
OPTIONAL INT96 value = 8;
757+
}
758+
}
759+
}
760+
";
761+
762+
let schema = parse_message_type(spark_schema).expect("should parse schema");
763+
let descr = SchemaDescriptor::new(Arc::new(schema));
764+
// Arrow derives a field id for every field but a map's entry struct, and
765+
// metadata holds more than ids, so stamp each field with a key of its
766+
// own: coercion has to preserve arbitrary metadata at every level.
767+
let arrow_schema =
768+
stamp_every_field(&parquet_to_arrow_schema(&descr, None).unwrap());
769+
770+
let result = Int96Coercer::new(&descr, &arrow_schema, &TimeUnit::Microsecond)
771+
.coerce()
772+
.unwrap();
773+
774+
assert_eq!(result.fields().len(), arrow_schema.fields().len());
775+
for (original, coerced) in arrow_schema.fields().iter().zip(result.fields()) {
776+
assert_field_metadata_preserved(original, coerced, original.name());
777+
}
778+
779+
// The containers are the fields the coercion rebuilds, so spell out that
780+
// they kept their ids: the recursion above compares the two schemas
781+
// against each other and would also pass on a schema carrying no ids.
782+
for (name, field_id) in [("c1", "2"), ("c3", "4"), ("c4", "6")] {
783+
assert_eq!(
784+
result.field_with_name(name).unwrap().metadata()
785+
[PARQUET_FIELD_ID_META_KEY],
786+
field_id,
787+
);
788+
}
789+
}
790+
791+
/// Metadata key stamped by [`stamp_every_field`], holding the field's path.
792+
const TEST_METADATA_KEY: &str = "test.field_path";
793+
794+
/// Add [`TEST_METADATA_KEY`] to every field of `schema`, nested fields
795+
/// included, keeping the metadata each field already has.
796+
fn stamp_every_field(schema: &Schema) -> Schema {
797+
fn stamp(field: &FieldRef, path: &str) -> FieldRef {
798+
let child_path = |child: &FieldRef| format!("{path}.{}", child.name());
799+
let data_type = match field.data_type() {
800+
DataType::Struct(children) => DataType::Struct(
801+
children
802+
.iter()
803+
.map(|child| stamp(child, &child_path(child)))
804+
.collect(),
805+
),
806+
DataType::List(child) => DataType::List(stamp(child, &child_path(child))),
807+
DataType::Map(child, sorted) => {
808+
DataType::Map(stamp(child, &child_path(child)), *sorted)
809+
}
810+
other => other.clone(),
811+
};
812+
let mut metadata = field.metadata().clone();
813+
metadata.insert(TEST_METADATA_KEY.to_string(), path.to_string());
814+
Arc::new(
815+
field
816+
.as_ref()
817+
.clone()
818+
.with_data_type(data_type)
819+
.with_metadata(metadata),
820+
)
821+
}
822+
823+
Schema::new_with_metadata(
824+
schema
825+
.fields()
826+
.iter()
827+
.map(|field| stamp(field, field.name()))
828+
.collect::<Vec<_>>(),
829+
schema.metadata.clone(),
830+
)
831+
}
832+
833+
/// Assert `coerced` carries the metadata of `original`, and so does every
834+
/// one of its children: struct fields, list elements, map entry structs.
835+
fn assert_field_metadata_preserved(
836+
original: &FieldRef,
837+
coerced: &FieldRef,
838+
path: &str,
839+
) {
840+
assert_eq!(
841+
original.metadata(),
842+
coerced.metadata(),
843+
"field {path} lost its metadata",
844+
);
845+
846+
match (original.data_type(), coerced.data_type()) {
847+
(DataType::Struct(original_children), DataType::Struct(coerced_children)) => {
848+
assert_eq!(original_children.len(), coerced_children.len());
849+
for (original_child, coerced_child) in
850+
original_children.iter().zip(coerced_children)
851+
{
852+
assert_field_metadata_preserved(
853+
original_child,
854+
coerced_child,
855+
&format!("{path}.{}", original_child.name()),
856+
);
857+
}
858+
}
859+
(DataType::List(original_child), DataType::List(coerced_child))
860+
| (DataType::Map(original_child, _), DataType::Map(coerced_child, _)) => {
861+
assert_field_metadata_preserved(
862+
original_child,
863+
coerced_child,
864+
&format!("{path}.{}", original_child.name()),
865+
);
866+
}
867+
// Only the arms above recurse, so a nested type here means the
868+
// coercion changed the field's shape and children went unchecked.
869+
(original_type, coerced_type) => assert!(
870+
!original_type.is_nested(),
871+
"field {path} changed shape: {original_type} -> {coerced_type}",
872+
),
873+
}
874+
}
731875
}

0 commit comments

Comments
 (0)