fix: preserve projection field metadata during physical planning

Author: subotac

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<dyn ExecutionPlan>,
         input: &Arc<LogicalPlan>,
         expr: &[Expr],
+        output_schema: &DFSchema,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         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;
 

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<Projector> {
+        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::<Vec<_>>();
+        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,

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<I, E>(
+        expr: I,
+        input: Arc<dyn ExecutionPlan>,
+        projected_schema: &Schema,
+    ) -> Result<Self>
+    where
+        I: IntoIterator<Item = E>,
+        E: Into<ProjectionExpr>,
+    {
+        let input_schema = input.schema();
+        let expr_arc = expr.into_iter().map(Into::into).collect::<Arc<_>>();
+        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<dyn ExecutionPlan>,