diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index fdf6ce323b19..1c5b56ffadbc 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/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::>>()?; - 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) => { diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 4bb27976515a..99b4ef6272b2 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/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 { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index e04124a61c96..61b1ea3ff104 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/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(deserializer: D) -> std::result::Result @@ -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__, }) } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 51e1a6fa9271..233b5fee1b29 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1687,6 +1687,8 @@ pub struct PhysicalAggregateExprNode { pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec>, #[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, + /// Output schema preserved by physical optimizer rewrites. + #[prost(message, optional, tag = "14")] + pub schema: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct GlobalLimitExecNode { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index aa10e1c5aa91..41deef4aa271 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/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(), }, )), }) diff --git a/datafusion/proto/tests/cases/plans/aggregates.rs b/datafusion/proto/tests/cases/plans/aggregates.rs index 34b51b57b01c..e57ac9fb5045 100644 --- a/datafusion/proto/tests/cases/plans/aggregates.rs +++ b/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 = 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::() + .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::() + .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 = 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);