gene-bordegaray commented on code in PR #25009:
URL: https://github.com/apache/datafusion/pull/25009#discussion_r3958383080
##########
datafusion/physical-plan/src/projection.rs:
##########
@@ -673,20 +673,25 @@ impl ExecutionPlan for ProjectionExec {
metrics: _,
// Derived plan properties, recomputed on decode.
cache: _,
- // Derived metadata comparison, recomputed with the projector.
- overrides_metadata: _,
+ overrides_metadata,
} = self;
let projection_exprs = projector.projection().as_ref();
let input = ctx.encode_child(input)?;
let expr = ctx.encode_expressions(projection_exprs.iter().map(|p|
&p.expr))?;
let expr_name = projection_exprs.iter().map(|p|
p.alias.clone()).collect();
+ let schema = if *overrides_metadata {
Review Comment:
This conditioins shouldnt just be if we are overriding metadata I believe.
Rather we should be checking if we have metadata in general.
##########
datafusion/proto/tests/cases/plans/exprs.rs:
##########
@@ -43,9 +43,176 @@ use datafusion_proto::physical_plan::{
};
use datafusion_proto::protobuf;
use datafusion_proto::protobuf::PhysicalPlanNode;
+use prost::Message;
+use std::collections::HashMap;
use std::sync::Arc;
use std::vec;
+#[test]
+fn roundtrip_projection_metadata() -> Result<()> {
+ let input_schema = Arc::new(Schema::new(vec![Field::new(
+ "value",
+ DataType::Int32,
+ false,
+ )]));
+ let projected_schema = Schema::new_with_metadata(
+ vec![Field::new("value", DataType::Int32, false).with_metadata(
+ [("field-key".to_string(), "field-value".to_string())].into(),
+ )],
+ [("schema-key".to_string(), "schema-value".to_string())].into(),
+ );
+ let plan = Arc::new(ProjectionExec::try_new_with_schema_metadata(
+ vec![(col("value", &input_schema)?, "value".to_string())],
+ Arc::new(EmptyExec::new(input_schema)),
+ &projected_schema,
+ )?);
+ let ctx = SessionContext::new();
+ let codec = DefaultPhysicalExtensionCodec {};
+ let converter = DefaultPhysicalProtoConverter {};
+ let decoded = roundtrip_test_and_return(plan, &ctx, &codec, &converter)?;
+ assert_eq!(decoded.schema().as_ref(), &projected_schema);
+ Ok(())
+}
+
+#[test]
+fn roundtrip_projection_metadata_overrides() -> Result<()> {
+ let field_metadata = [("field-key".to_string(),
"field-value".to_string())].into();
+ let extension_metadata = [
+ ("ARROW:extension:name".to_string(), "arrow.uuid".to_string()),
+ ("ARROW:extension:metadata".to_string(), String::new()),
+ ]
+ .into();
+ for (input_field, output_field, input_metadata) in [
+ (
+ Field::new("value", DataType::Int32,
false).with_metadata(field_metadata),
+ Field::new("value", DataType::Int32, false),
+ [("input-schema".to_string(), "input-value".to_string())].into(),
+ ),
+ (
+ Field::new("value", DataType::FixedSizeBinary(16), true),
+ Field::new("value", DataType::FixedSizeBinary(16), true)
+ .with_metadata(extension_metadata),
+ HashMap::new(),
+ ),
+ ] {
+ let input_schema =
+ Arc::new(Schema::new_with_metadata(vec![input_field],
input_metadata));
+ let projected_schema = Schema::new(vec![output_field]);
+ let plan = Arc::new(ProjectionExec::try_new_with_schema_metadata(
+ vec![(col("value", &input_schema)?, "value".to_string())],
+ Arc::new(EmptyExec::new(input_schema)),
+ &projected_schema,
+ )?);
+ let codec = DefaultPhysicalExtensionCodec {};
+ let ctx = SessionContext::new();
+ let node = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?;
+ let
Some(protobuf::physical_plan_node::PhysicalPlanType::Projection(projection)) =
+ node.physical_plan_type.as_ref()
+ else {
+ unreachable!("expected ProjectionExecNode")
+ };
+ assert!(projection.schema.is_some());
+ let node =
PhysicalPlanNode::decode(node.encode_to_vec().as_slice()).unwrap();
+ #[cfg(feature = "json")]
+ let node: PhysicalPlanNode =
+
serde_json::from_str(&serde_json::to_string(&node).unwrap()).unwrap();
+ let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?;
+ assert_eq!(decoded.schema().as_ref(), &projected_schema);
+ }
+ Ok(())
+}
+
+#[test]
+fn roundtrip_projection_without_metadata_override() -> Result<()> {
+ let input_schema = Arc::new(Schema::new_with_metadata(
+ vec![Field::new("value", DataType::Int32, false).with_metadata(
+ [("field-key".to_string(), "field-value".to_string())].into(),
+ )],
+ [("schema-key".to_string(), "schema-value".to_string())].into(),
+ ));
+ let plan = Arc::new(ProjectionExec::try_new(
+ vec![(col("value", &input_schema)?, "value".to_string())],
+ Arc::new(EmptyExec::new(Arc::clone(&input_schema))),
Review Comment:
this is preserving the metadata through the child after decoding, not a
proojection that needs to rederive the metadata completely from itself.
Could we add a test that forces the projection to completely rederive the
metadata from its own proto after roundtrip 👍
--
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]
