From 9c14234ddf8115e8c2027fb277be54522f830778 Mon Sep 17 00:00:00 2001 From: subotac <73706465+subotac@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:56:58 +0300 Subject: [PATCH] fix: preserve projection field metadata during physical planning --- datafusion/core/src/physical_planner.rs | 55 ++++++++++++++++++++-- datafusion/physical-expr/src/projection.rs | 50 ++++++++++++++++++++ datafusion/physical-plan/src/projection.rs | 30 +++++++++++- 3 files changed, 131 insertions(+), 4 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 4e914556b4cc0..c826018c0dd56 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1239,6 +1239,7 @@ impl DefaultPhysicalPlanner { children.one()?, input, expr, + node.schema(), )?, LogicalPlan::Filter(Filter { predicate, input, .. @@ -1462,6 +1463,7 @@ impl DefaultPhysicalPlanner { physical_left, input, expr, + left.schema(), )?, _ => physical_left, }; @@ -1476,6 +1478,7 @@ impl DefaultPhysicalPlanner { physical_right, input, expr, + right.schema(), )?, _ => physical_right, }; @@ -1860,6 +1863,7 @@ impl DefaultPhysicalPlanner { join, input, expr, + new_logical.schema(), )? } else { join @@ -3088,6 +3092,7 @@ impl DefaultPhysicalPlanner { input_exec: Arc, input: &Arc, expr: &[Expr], + output_schema: &DFSchema, ) -> Result> { let input_logical_schema = input.as_ref().schema(); let input_physical_schema = input_exec.schema(); @@ -3145,7 +3150,11 @@ impl DefaultPhysicalPlanner { .into_iter() .map(|(expr, alias)| ProjectionExpr { expr, alias }) .collect(); - Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input_exec)?)) + Ok(Arc::new(ProjectionExec::try_new_with_schema_metadata( + proj_exprs, + input_exec, + output_schema.as_arrow(), + )?)) } PlanAsyncExpr::Async( async_map, @@ -3157,8 +3166,11 @@ impl DefaultPhysicalPlanner { .into_iter() .map(|(expr, alias)| ProjectionExpr { expr, alias }) .collect(); - let new_proj_exec = - ProjectionExec::try_new(proj_exprs, Arc::new(async_exec))?; + let new_proj_exec = ProjectionExec::try_new_with_schema_metadata( + proj_exprs, + Arc::new(async_exec), + output_schema.as_arrow(), + )?; Ok(Arc::new(new_proj_exec)) } _ => internal_err!("Unexpected PlanAsyncExpressions variant"), @@ -3505,6 +3517,43 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_projection_preserves_field_metadata_for_aggregate() -> Result<()> { + use datafusion_common::metadata::FieldMetadata; + use datafusion_expr::expr::AggregateFunction; + use datafusion_functions_aggregate::min_max::max_udaf; + + let schema = Schema::new(vec![Field::new("value", DataType::Utf8, false)]); + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(schema.to_dfschema()?), + }); + let metadata = + FieldMetadata::from(HashMap::from([("foo".to_string(), "bar".to_string())])); + let projection = LogicalPlan::Projection(Projection::try_new( + vec![col("value").alias_with_metadata("value", Some(metadata))], + Arc::new(input), + )?); + let aggregate = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::new(projection), + vec![], + vec![Expr::AggregateFunction(AggregateFunction::new_udf( + max_udaf(), + vec![col("value")], + false, + None, + vec![], + None, + ))], + )?); + + DefaultPhysicalPlanner::default() + .create_physical_plan(&aggregate, &SessionContext::new().state()) + .await?; + + Ok(()) + } + #[derive(Debug, Default)] struct NullAccumulator; diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index 1f6a6eb08fb78..656cb614fb62e 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -539,6 +539,56 @@ impl ProjectionExprs { }) } + /// Create a new [`Projector`] using field and schema metadata from + /// `projected_schema`. + /// + /// Field names, data types, and nullability are still derived from the physical + /// projection expressions and `input_schema`; only field and schema metadata are + /// taken from `projected_schema`. + /// + /// # Errors + /// + /// Returns an error if the projection cannot be applied to `input_schema`, or if + /// `projected_schema` has a different number of fields than the projection. + pub fn make_projector_with_schema_metadata( + &self, + input_schema: &Schema, + projected_schema: &Schema, + ) -> Result { + let output_schema = self.project_schema(input_schema)?; + if output_schema.fields().len() != projected_schema.fields().len() { + return Err(internal_datafusion_err!( + "Projection has {} output fields but metadata schema has {} fields", + output_schema.fields().len(), + projected_schema.fields().len() + )); + } + + let fields = output_schema + .fields() + .iter() + .zip(projected_schema.fields()) + .map(|(field, projected_field)| { + Arc::new( + field + .as_ref() + .clone() + .with_metadata(projected_field.metadata().clone()), + ) + }) + .collect::>(); + let output_schema = Arc::new(Schema::new_with_metadata( + fields, + projected_schema.metadata().clone(), + )); + + Ok(Projector { + projection: self.clone(), + output_schema, + expression_metrics: None, + }) + } + pub fn create_expression_metrics( &self, metrics: &ExecutionPlanMetricsSet, diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 42501f22395b4..199467f502626 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -40,7 +40,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{ @@ -143,6 +143,34 @@ impl ProjectionExec { Self::try_from_projector(projector, input) } + /// Create a projection using field and schema metadata from + /// `projected_schema`. + /// + /// Field names, data types, and nullability are still derived from the physical + /// projection expressions and the input plan; only field and schema metadata are + /// taken from `projected_schema`. + /// + /// # Errors + /// + /// Returns an error if the projection cannot be applied to the input plan, or if + /// `projected_schema` has a different number of fields than the projection. + pub fn try_new_with_schema_metadata( + expr: I, + input: Arc, + projected_schema: &Schema, + ) -> Result + where + I: IntoIterator, + E: Into, + { + let input_schema = input.schema(); + let expr_arc = expr.into_iter().map(Into::into).collect::>(); + let projection = ProjectionExprs::from_expressions(expr_arc); + let projector = projection + .make_projector_with_schema_metadata(&input_schema, projected_schema)?; + Self::try_from_projector(projector, input) + } + fn try_from_projector( projector: Projector, input: Arc,