From b8135139c97918c95f855dd9c865ebdc2525b600 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 11 Jan 2025 11:07:57 +0800 Subject: [PATCH] Minor: Make `group_schema` as `PhysicalGroupBy` method (#14064) * group shema as method Signed-off-by: Jay Zhan * fmt Signed-off-by: Jay Zhan --------- Signed-off-by: Jay Zhan --- datafusion/physical-plan/src/aggregates/mod.rs | 8 ++++---- datafusion/physical-plan/src/aggregates/row_hash.rs | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 52fd6f90e595..cc8d6e74f4b9 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -250,6 +250,10 @@ impl PhysicalGroupBy { } } + pub fn group_schema(&self, schema: &Schema) -> Result { + Ok(Arc::new(Schema::new(self.group_fields(schema)?))) + } + /// Returns the fields that are used as the grouping keys. fn group_fields(&self, input_schema: &Schema) -> Result> { let mut fields = Vec::with_capacity(self.num_group_exprs()); @@ -924,10 +928,6 @@ fn create_schema( )) } -fn group_schema(input_schema: &Schema, group_by: &PhysicalGroupBy) -> Result { - Ok(Arc::new(Schema::new(group_by.group_fields(input_schema)?))) -} - /// Determines the lexical ordering requirement for an aggregate expression. /// /// # Parameters diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs index cdb3b2199cdc..cc95ce51c15b 100644 --- a/datafusion/physical-plan/src/aggregates/row_hash.rs +++ b/datafusion/physical-plan/src/aggregates/row_hash.rs @@ -24,8 +24,8 @@ use std::vec; use crate::aggregates::group_values::{new_group_values, GroupValues}; use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ - create_schema, evaluate_group_by, evaluate_many, evaluate_optional, group_schema, - AggregateMode, PhysicalGroupBy, + create_schema, evaluate_group_by, evaluate_many, evaluate_optional, AggregateMode, + PhysicalGroupBy, }; use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput}; use crate::sorts::sort::sort_batch; @@ -489,7 +489,7 @@ impl GroupedHashAggregateStream { .map(create_group_accumulator) .collect::>()?; - let group_schema = group_schema(&agg.input().schema(), &agg_group_by)?; + let group_schema = agg_group_by.group_schema(&agg.input().schema())?; // fix https://github.com/apache/datafusion/issues/13949 // Builds a **partial aggregation** schema by combining the group columns and