fix(proto): preserve AggregateExec schema and reversed state

datafusion/physical-plan/src/aggregates/mod.rs

@@ -2332,6 +2332,7 @@ impl ExecutionPlan for AggregateExec {
                         limit,
                         has_grouping_set: group_by.has_grouping_set(),
                         dynamic_filter,
+                        schema: Some(self.schema.as_ref().try_into()?),
                     },
                 )),
             ),
@@ -2408,6 +2409,7 @@ fn encode_aggregate_expr(
                 ignore_nulls: aggr_expr.ignore_nulls(),
                 fun_definition,
                 human_display,
+                is_reversed: aggr_expr.is_reversed(),
             },
         )),
     })
@@ -2450,6 +2452,7 @@ impl AggregateExec {
             limit,
             has_grouping_set,
             dynamic_filter,
+            schema,
         } = hash_agg.as_ref();
 
         let input =
@@ -2563,6 +2566,7 @@ impl AggregateExec {
                     .with_ignore_nulls(aggregate.ignore_nulls)
                     .with_distinct(aggregate.distinct)
                     .order_by(order_by)
+                    .with_reversed(aggregate.is_reversed)
                     .human_display(human_display);
                 let builder = if let Some(alias) = human_display_alias {
                     builder.human_display_alias(alias)
@@ -2572,14 +2576,29 @@ impl AggregateExec {
                 builder.build().map(Arc::new)
             })
             .collect::<Result<Vec<_>>>()?;
-        let aggregate = AggregateExec::try_new(
-            mode,
-            PhysicalGroupBy::new(group_expr, null_expr, groups, *has_grouping_set),
-            aggr_expr,
-            filter_expr,
-            input,
-            Arc::clone(&input_schema),
-        )?;
+        let group_by =
+            PhysicalGroupBy::new(group_expr, null_expr, groups, *has_grouping_set);
+        let aggregate = if let Some(schema) = schema {
+            let schema = SchemaRef::new(schema.try_into()?);
+            AggregateExec::try_new_with_schema(
+                mode,
+                group_by,
+                aggr_expr,
+                filter_expr,
+                input,
+                Arc::clone(&input_schema),
+                schema,
+            )
+        } else {
+            AggregateExec::try_new(
+                mode,
+                group_by,
+                aggr_expr,
+                filter_expr,
+                input,
+                Arc::clone(&input_schema),
+            )
+        }?;
         let aggregate = if let Some(limit) = limit {
             let options = match limit.descending {
                 Some(descending) => {

datafusion/proto-models/proto/datafusion.proto

@@ -1096,6 +1096,7 @@ message PhysicalAggregateExprNode {
   bool ignore_nulls = 6;
   optional bytes fun_definition = 7;
   string human_display = 8;
+  bool is_reversed = 9;
 }
 
 message PhysicalWindowExprNode {
@@ -1472,6 +1473,8 @@ message AggregateExecNode {
   bool has_grouping_set = 12;
   // Optional dynamic filter expression for pushing down to the child.
   PhysicalExprNode dynamic_filter = 13;
+  // Output schema preserved by physical optimizer rewrites.
+  datafusion_common.Schema schema = 14;
 }
 
 message GlobalLimitExecNode {

datafusion/proto-models/src/generated/pbjson.rs

@@ -157,6 +157,9 @@ impl serde::Serialize for AggregateExecNode {
         if self.dynamic_filter.is_some() {
             len += 1;
         }
+        if self.schema.is_some() {
+            len += 1;
+        }
         let mut struct_ser = serializer.serialize_struct("datafusion.AggregateExecNode", len)?;
         if !self.group_expr.is_empty() {
             struct_ser.serialize_field("groupExpr", &self.group_expr)?;
@@ -199,6 +202,9 @@ impl serde::Serialize for AggregateExecNode {
         if let Some(v) = self.dynamic_filter.as_ref() {
             struct_ser.serialize_field("dynamicFilter", v)?;
         }
+        if let Some(v) = self.schema.as_ref() {
+            struct_ser.serialize_field("schema", v)?;
+        }
         struct_ser.end()
     }
 }
@@ -231,6 +237,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode {
             "hasGroupingSet",
             "dynamic_filter",
             "dynamicFilter",
+            "schema",
         ];
 
         #[allow(clippy::enum_variant_names)]
@@ -248,6 +255,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode {
             Limit,
             HasGroupingSet,
             DynamicFilter,
+            Schema,
         }
         impl<'de> serde::Deserialize<'de> for GeneratedField {
             fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -282,6 +290,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode {
                             "limit" => Ok(GeneratedField::Limit),
                             "hasGroupingSet" | "has_grouping_set" => Ok(GeneratedField::HasGroupingSet),
                             "dynamicFilter" | "dynamic_filter" => Ok(GeneratedField::DynamicFilter),
+                            "schema" => Ok(GeneratedField::Schema),
                             _ => Err(serde::de::Error::unknown_field(value, FIELDS)),
                         }
                     }
@@ -314,6 +323,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode {
                 let mut limit__ = None;
                 let mut has_grouping_set__ = None;
                 let mut dynamic_filter__ = None;
+                let mut schema__ = None;
                 while let Some(k) = map_.next_key()? {
                     match k {
                         GeneratedField::GroupExpr => {
@@ -394,6 +404,12 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode {
                             }
                             dynamic_filter__ = map_.next_value()?;
                         }
+                        GeneratedField::Schema => {
+                            if schema__.is_some() {
+                                return Err(serde::de::Error::duplicate_field("schema"));
+                            }
+                            schema__ = map_.next_value()?;
+                        }
                     }
                 }
                 Ok(AggregateExecNode {
@@ -410,6 +426,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode {
                     limit: limit__,
                     has_grouping_set: has_grouping_set__.unwrap_or_default(),
                     dynamic_filter: dynamic_filter__,
+                    schema: schema__,
                 })
             }
         }
@@ -17293,6 +17310,9 @@ impl serde::Serialize for PhysicalAggregateExprNode {
         if !self.human_display.is_empty() {
             len += 1;
         }
+        if self.is_reversed {
+            len += 1;
+        }
         if self.aggregate_function.is_some() {
             len += 1;
         }
@@ -17317,6 +17337,9 @@ impl serde::Serialize for PhysicalAggregateExprNode {
         if !self.human_display.is_empty() {
             struct_ser.serialize_field("humanDisplay", &self.human_display)?;
         }
+        if self.is_reversed {
+            struct_ser.serialize_field("isReversed", &self.is_reversed)?;
+        }
         if let Some(v) = self.aggregate_function.as_ref() {
             match v {
                 physical_aggregate_expr_node::AggregateFunction::UserDefinedAggrFunction(v) => {
@@ -17344,6 +17367,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode {
             "funDefinition",
             "human_display",
             "humanDisplay",
+            "is_reversed",
+            "isReversed",
             "user_defined_aggr_function",
             "userDefinedAggrFunction",
         ];
@@ -17356,6 +17381,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode {
             IgnoreNulls,
             FunDefinition,
             HumanDisplay,
+            IsReversed,
             UserDefinedAggrFunction,
         }
         impl<'de> serde::Deserialize<'de> for GeneratedField {
@@ -17384,6 +17410,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode {
                             "ignoreNulls" | "ignore_nulls" => Ok(GeneratedField::IgnoreNulls),
                             "funDefinition" | "fun_definition" => Ok(GeneratedField::FunDefinition),
                             "humanDisplay" | "human_display" => Ok(GeneratedField::HumanDisplay),
+                            "isReversed" | "is_reversed" => Ok(GeneratedField::IsReversed),
                             "userDefinedAggrFunction" | "user_defined_aggr_function" => Ok(GeneratedField::UserDefinedAggrFunction),
                             _ => Err(serde::de::Error::unknown_field(value, FIELDS)),
                         }
@@ -17410,6 +17437,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode {
                 let mut ignore_nulls__ = None;
                 let mut fun_definition__ = None;
                 let mut human_display__ = None;
+                let mut is_reversed__ = None;
                 let mut aggregate_function__ = None;
                 while let Some(k) = map_.next_key()? {
                     match k {
@@ -17451,6 +17479,12 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode {
                             }
                             human_display__ = Some(map_.next_value()?);
                         }
+                        GeneratedField::IsReversed => {
+                            if is_reversed__.is_some() {
+                                return Err(serde::de::Error::duplicate_field("isReversed"));
+                            }
+                            is_reversed__ = Some(map_.next_value()?);
+                        }
                         GeneratedField::UserDefinedAggrFunction => {
                             if aggregate_function__.is_some() {
                                 return Err(serde::de::Error::duplicate_field("userDefinedAggrFunction"));
@@ -17466,6 +17500,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode {
                     ignore_nulls: ignore_nulls__.unwrap_or_default(),
                     fun_definition: fun_definition__,
                     human_display: human_display__.unwrap_or_default(),
+                    is_reversed: is_reversed__.unwrap_or_default(),
                     aggregate_function: aggregate_function__,
                 })
             }

datafusion/proto-models/src/generated/prost.rs

@@ -1687,6 +1687,8 @@ pub struct PhysicalAggregateExprNode {
     pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
     #[prost(string, tag = "8")]
     pub human_display: ::prost::alloc::string::String,
+    #[prost(bool, tag = "9")]
+    pub is_reversed: bool,
     #[prost(oneof = "physical_aggregate_expr_node::AggregateFunction", tags = "4")]
     pub aggregate_function: ::core::option::Option<
         physical_aggregate_expr_node::AggregateFunction,
@@ -2230,6 +2232,9 @@ pub struct AggregateExecNode {
     /// Optional dynamic filter expression for pushing down to the child.
     #[prost(message, optional, tag = "13")]
     pub dynamic_filter: ::core::option::Option<PhysicalExprNode>,
+    /// Output schema preserved by physical optimizer rewrites.
+    #[prost(message, optional, tag = "14")]
+    pub schema: ::core::option::Option<super::datafusion_common::Schema>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct GlobalLimitExecNode {

datafusion/proto/src/physical_plan/to_proto.rs

@@ -82,6 +82,7 @@ pub fn serialize_physical_aggr_expr(
                 ignore_nulls: aggr_expr.ignore_nulls(),
                 fun_definition: (!buf.is_empty()).then_some(buf),
                 human_display,
+                is_reversed: aggr_expr.is_reversed(),
             },
         )),
     })

datafusion/proto/tests/cases/plans/aggregates.rs

@@ -22,24 +22,35 @@ use datafusion::arrow::array::ArrayRef;
 use datafusion::arrow::compute::kernels::sort::SortOptions;
 use datafusion::arrow::datatypes::{DataType, Field, Schema};
 use datafusion::logical_expr::Volatility;
+use datafusion::physical_expr::LexOrdering;
 use datafusion::physical_expr::aggregate::AggregateExprBuilder;
-use datafusion::physical_plan::PhysicalExpr;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_optimizer::update_aggr_exprs::OptimizeAggregateOrder;
 use datafusion::physical_plan::aggregates::{
     AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy,
 };
 use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::physical_plan::expressions::{PhysicalSortExpr, col, lit};
+use datafusion::physical_plan::sorts::sort::SortExec;
+use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr};
 use datafusion::prelude::SessionContext;
 use datafusion::scalar::ScalarValue;
-use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::{
     Accumulator, AccumulatorFactoryFunction, AggregateUDF, Signature, SimpleAggregateUDF,
 };
 use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf;
 use datafusion_functions_aggregate::array_agg::array_agg_udaf;
 use datafusion_functions_aggregate::average::avg_udaf;
+use datafusion_functions_aggregate::first_last::first_value_udaf;
 use datafusion_functions_aggregate::nth_value::nth_value_udaf;
 use datafusion_functions_aggregate::string_agg::string_agg_udaf;
+use datafusion_functions_aggregate::sum::sum_udaf;
+use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec};
+use datafusion_proto::protobuf;
+use datafusion_proto::protobuf::PhysicalPlanNode;
+use prost::Message;
 use std::sync::Arc;
 use std::vec;
 
@@ -91,6 +102,100 @@ fn roundtrip_aggregate() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn roundtrip_aggregate_preserves_optimizer_schema_and_reversed_state() -> Result<()> {
+    let input_schema =
+        Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)]));
+    let input_ordering = LexOrdering::new(vec![PhysicalSortExpr {
+        expr: col("b", &input_schema)?,
+        options: SortOptions::new(true, true),
+    }])
+    .expect("single sort expression should form an ordering");
+    let input: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
+        input_ordering,
+        Arc::new(EmptyExec::new(Arc::clone(&input_schema))),
+    ));
+    let original_name = "first_value(b) ORDER BY [b ASC NULLS LAST]";
+    let aggregate_expr = Arc::new(
+        AggregateExprBuilder::new(first_value_udaf(), vec![col("b", &input_schema)?])
+            .order_by(vec![PhysicalSortExpr {
+                expr: col("b", &input_schema)?,
+                options: SortOptions::new(false, false),
+            }])
+            .schema(Arc::clone(&input_schema))
+            .alias(original_name)
+            .build()?,
+    );
+    let aggregate = AggregateExec::try_new(
+        AggregateMode::Single,
+        PhysicalGroupBy::new(vec![], vec![], vec![], false),
+        vec![aggregate_expr],
+        vec![None],
+        input,
+        input_schema,
+    )?;
+
+    let optimized = OptimizeAggregateOrder::new()
+        .optimize(Arc::new(aggregate), &ConfigOptions::new())?;
+    let optimized_aggregate = optimized
+        .downcast_ref::<AggregateExec>()
+        .expect("expected optimized AggregateExec");
+    assert_eq!(optimized.schema().field(0).name(), original_name);
+    assert_eq!(
+        optimized_aggregate.aggr_expr()[0].name(),
+        "last_value(b) ORDER BY [b DESC NULLS FIRST]"
+    );
+    assert!(optimized_aggregate.aggr_expr()[0].is_reversed());
+
+    let codec = DefaultPhysicalExtensionCodec {};
+    let node = PhysicalPlanNode::try_from_physical_plan(Arc::clone(&optimized), &codec)?;
+    let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice())
+        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+    let ctx = SessionContext::new();
+    let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?;
+    let decoded_aggregate = decoded
+        .downcast_ref::<AggregateExec>()
+        .expect("expected decoded AggregateExec");
+
+    assert_eq!(optimized.schema(), decoded.schema());
+    assert!(decoded_aggregate.aggr_expr()[0].is_reversed());
+    Ok(())
+}
+
+#[test]
+fn decode_aggregate_without_output_schema() -> Result<()> {
+    let input_schema =
+        Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)]));
+    let aggregate_expr = Arc::new(
+        AggregateExprBuilder::new(sum_udaf(), vec![col("b", &input_schema)?])
+            .schema(Arc::clone(&input_schema))
+            .alias("SUM(b)")
+            .build()?,
+    );
+    let plan: Arc<dyn ExecutionPlan> = Arc::new(AggregateExec::try_new(
+        AggregateMode::Single,
+        PhysicalGroupBy::new(vec![], vec![], vec![], false),
+        vec![aggregate_expr],
+        vec![None],
+        Arc::new(EmptyExec::new(Arc::clone(&input_schema))),
+        input_schema,
+    )?);
+
+    let codec = DefaultPhysicalExtensionCodec {};
+    let mut node = PhysicalPlanNode::try_from_physical_plan(Arc::clone(&plan), &codec)?;
+    let Some(protobuf::physical_plan_node::PhysicalPlanType::Aggregate(aggregate)) =
+        node.physical_plan_type.as_mut()
+    else {
+        panic!("expected AggregateExecNode");
+    };
+    assert!(aggregate.schema.take().is_some());
+
+    let ctx = SessionContext::new();
+    let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?;
+    assert_eq!(plan.schema(), decoded.schema());
+    Ok(())
+}
+
 #[test]
 fn roundtrip_aggregate_with_limit() -> Result<()> {
     let field_a = Field::new("a", DataType::Int64, false);