paleolimbot commented on code in PR #23169:
URL: https://github.com/apache/datafusion/pull/23169#discussion_r3509294400
##########
datafusion/expr/src/expr_schema.rs:
##########
@@ -1208,4 +1229,121 @@ mod tests {
assert_eq!(meta, expr.metadata(&schema).unwrap());
}
+
+ #[test]
+ fn test_cast_extension_type_metadata() {
+ use crate::expr::Cast;
+ use arrow_schema::extension::{
+ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY,
+ };
+
+ // Create a schema with a field that has extension type metadata
+ let mut source_meta = HashMap::new();
+ source_meta.insert(
+ EXTENSION_TYPE_NAME_KEY.to_string(),
+ "arrow.uuid".to_string(),
+ );
+ source_meta.insert("custom_key".to_string(),
"custom_value".to_string());
+
+ let source_field = Field::new("foo", DataType::FixedSizeBinary(16),
false)
+ .with_metadata(source_meta);
+
+ let schema = MockExprSchema::new()
+ .with_data_type(DataType::FixedSizeBinary(16))
+
.with_metadata(FieldMetadata::from(source_field.metadata().clone()));
+
+ // Test 1: Cast to a type without extension metadata strips extension
metadata
+ // but preserves non-extension metadata
+ let cast_expr = Expr::Cast(Cast {
+ expr: Box::new(col("foo")),
+ field: Arc::new(Field::new("", DataType::Utf8, true)),
+ });
+
+ let (_, result_field) = cast_expr.to_field(&schema).unwrap();
+ assert!(
+ result_field
+ .metadata()
+ .get(EXTENSION_TYPE_NAME_KEY)
+ .is_none(),
+ "Extension type name should be stripped when target has no
extension metadata"
+ );
+ assert_eq!(
+ result_field.metadata().get("custom_key"),
+ Some(&"custom_value".to_string()),
+ "Non-extension metadata should be preserved"
+ );
+
+ // Test 2: Cast to a field with different extension type replaces
extension metadata
+ let mut target_meta = HashMap::new();
+ target_meta.insert(
+ EXTENSION_TYPE_NAME_KEY.to_string(),
+ "arrow.json".to_string(),
+ );
+ target_meta.insert(EXTENSION_TYPE_METADATA_KEY.to_string(),
"{}".to_string());
+
+ let target_field =
+ Field::new("", DataType::Utf8, true).with_metadata(target_meta);
+
+ let cast_expr = Expr::Cast(Cast {
+ expr: Box::new(col("foo")),
+ field: Arc::new(target_field),
+ });
+
+ let (_, result_field) = cast_expr.to_field(&schema).unwrap();
+ assert_eq!(
+ result_field.metadata().get(EXTENSION_TYPE_NAME_KEY),
+ Some(&"arrow.json".to_string()),
+ "Extension type name should come from target field"
+ );
+ assert_eq!(
+ result_field.metadata().get(EXTENSION_TYPE_METADATA_KEY),
+ Some(&"{}".to_string()),
+ "Extension type metadata should come from target field"
+ );
+ assert_eq!(
+ result_field.metadata().get("custom_key"),
+ Some(&"custom_value".to_string()),
+ "Non-extension metadata should still be preserved from source"
+ );
+ }
+
+ #[test]
+ fn test_try_cast_extension_type_metadata() {
Review Comment:
I think I did this one!
--
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]
