Minor: Make group_schema as PhysicalGroupBy method (#14064)

* group shema as method Signed-off-by: Jay Zhan <[email protected]> * fmt Signed-off-by: Jay Zhan <[email protected]> --------- Signed-off-by: Jay Zhan <[email protected]>
apache · Jan 11, 2025 · b813513 · b813513
1 parent 4b7281e
commit b813513
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git 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<SchemaRef> {
+        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<Vec<Field>> {
         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<SchemaRef> {
-    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
@@ -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::<Result<_>>()?;
 
-        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