Skip to content

Commit 800425c

Browse files
Weijun-Halamb
andauthored
refactor: improve display formatting for Union (#8529)
# Which issue does this PR close? - Parts of #8351 # Rationale for this change Support more human readable display in `Union` # What changes are included in this PR? # Are these changes tested? Yes # Are there any user-facing changes? No --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent 28aaee8 commit 800425c

File tree

1 file changed

+83
-1
lines changed

1 file changed

+83
-1
lines changed

arrow-schema/src/datatype_display.rs

Lines changed: 83 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,24 @@ impl fmt::Display for DataType {
132132
Ok(())
133133
}
134134
Self::Union(union_fields, union_mode) => {
135-
write!(f, "Union({union_fields:?}, {union_mode:?})")
135+
write!(f, "Union({union_mode:?}, ")?;
136+
if !union_fields.is_empty() {
137+
let fields_str = union_fields
138+
.iter()
139+
.map(|v| {
140+
let type_id = v.0;
141+
let field = v.1;
142+
let maybe_nullable = if field.is_nullable() { "nullable " } else { "" };
143+
let data_type = field.data_type();
144+
let metadata_str = format_metadata(field.metadata());
145+
format!("{type_id:?}: {maybe_nullable}{data_type}{metadata_str}")
146+
})
147+
.collect::<Vec<_>>()
148+
.join(", ");
149+
write!(f, "{fields_str}")?;
150+
}
151+
write!(f, ")")?;
152+
Ok(())
136153
}
137154
Self::Dictionary(data_type, data_type1) => {
138155
write!(f, "Dictionary({data_type}, {data_type1})")
@@ -248,4 +265,69 @@ mod tests {
248265
"FixedSizeList(4 x nullable Int32, metadata: {\"key2\": \"value2\"})";
249266
assert_eq!(fixed_size_metadata_string, expected_metadata_string);
250267
}
268+
269+
#[test]
270+
fn test_display_struct() {
271+
let fields = vec![
272+
Field::new("a", DataType::Int32, false),
273+
Field::new("b", DataType::Utf8, true),
274+
];
275+
let struct_data_type = DataType::Struct(fields.into());
276+
let struct_data_type_string = struct_data_type.to_string();
277+
let expected_string = "Struct(\"a\": Int32, \"b\": nullable Utf8)";
278+
assert_eq!(struct_data_type_string, expected_string);
279+
280+
// Test with metadata
281+
let mut field_with_metadata = Field::new("b", DataType::Utf8, true);
282+
let metadata = HashMap::from([("key".to_string(), "value".to_string())]);
283+
field_with_metadata.set_metadata(metadata);
284+
let struct_fields_with_metadata =
285+
vec![Field::new("a", DataType::Int32, false), field_with_metadata];
286+
let struct_data_type_with_metadata = DataType::Struct(struct_fields_with_metadata.into());
287+
let struct_data_type_with_metadata_string = struct_data_type_with_metadata.to_string();
288+
let expected_string_with_metadata =
289+
"Struct(\"a\": Int32, \"b\": nullable Utf8, metadata: {\"key\": \"value\"})";
290+
assert_eq!(
291+
struct_data_type_with_metadata_string,
292+
expected_string_with_metadata
293+
);
294+
}
295+
296+
#[test]
297+
fn test_display_union() {
298+
let fields = vec![
299+
Field::new("a", DataType::Int32, false),
300+
Field::new("b", DataType::Utf8, true),
301+
];
302+
let type_ids = vec![0, 1];
303+
let union_fields = type_ids
304+
.into_iter()
305+
.zip(fields.into_iter().map(Arc::new))
306+
.collect();
307+
308+
let union_data_type = DataType::Union(union_fields, crate::UnionMode::Sparse);
309+
let union_data_type_string = union_data_type.to_string();
310+
let expected_string = "Union(Sparse, 0: Int32, 1: nullable Utf8)";
311+
assert_eq!(union_data_type_string, expected_string);
312+
313+
// Test with metadata
314+
let mut field_with_metadata = Field::new("b", DataType::Utf8, true);
315+
let metadata = HashMap::from([("key".to_string(), "value".to_string())]);
316+
field_with_metadata.set_metadata(metadata);
317+
let union_fields_with_metadata = vec![
318+
(0, Arc::new(Field::new("a", DataType::Int32, false))),
319+
(1, Arc::new(field_with_metadata)),
320+
]
321+
.into_iter()
322+
.collect();
323+
let union_data_type_with_metadata =
324+
DataType::Union(union_fields_with_metadata, crate::UnionMode::Sparse);
325+
let union_data_type_with_metadata_string = union_data_type_with_metadata.to_string();
326+
let expected_string_with_metadata =
327+
"Union(Sparse, 0: Int32, 1: nullable Utf8, metadata: {\"key\": \"value\"})";
328+
assert_eq!(
329+
union_data_type_with_metadata_string,
330+
expected_string_with_metadata
331+
);
332+
}
251333
}

0 commit comments

Comments
 (0)