alamb commented on code in PR #23181: URL: https://github.com/apache/datafusion/pull/23181#discussion_r3500632005
########## datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs: ########## @@ -0,0 +1,214 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Common utilities for aggregate tables used in aggregations that inputs are ordered +//! by the groups. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +use crate::InputOrderMode; +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::row_hash::create_group_accumulator; +use crate::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, +}; + +use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; + +/// Marker for raw rows -> partial state aggregation on ordered input. +pub(in crate::aggregates) struct OrderedPartialMarker; +/// Marker for partial state -> final value aggregation on ordered input. +pub(in crate::aggregates) struct OrderedFinalMarker; + +/// Aggregate table shared by the ordered partial and final paths. +/// +/// The table consumes input batches while `GroupOrdering` tracks which groups +/// are proven complete. Completed groups can be emitted before the input stream +/// ends, which keeps memory bounded by the active ordered key range. +/// +/// # Marker Type +/// +/// `OrderedAggrMode` selects the aggregate semantics. For example, +/// `OrderedAggregateTable::<OrderedPartialMarker>::new(...)` consumes raw rows +/// and emits partial states, while +/// `OrderedAggregateTable::<OrderedFinalMarker>::new_with_input_order(...)` +/// consumes partial states and emits final values. +/// +/// Shared methods live on `impl<T>`; partial/final behavior lives on +/// marker-specific impls. +pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> { + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Group keys, ordering state, and accumulator states. + pub(super) buffer: OrderedAggregateTableBuffer, + + _mode: PhantomData<OrderedAggrMode>, +} + +/// Buffer for the ordered aggregate table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits output rows as soon as the Review Comment: these are really nice, easy to understand, comments ########## datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs: ########## @@ -16,9 +16,15 @@ // under the License. mod common; +mod common_ordered; mod final_table; +mod ordered_final_table; +mod ordered_partial_table; mod partial_table; pub(super) use common::{ AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, Review Comment: I wonder if there is any fundamental difference between `OrderedFinalMarker` and `FinalMarker`? Likewise for `OrderedPartialMarker` and `PartialMarker` if not, perhaps we can consolidate them into one set of markers in a future PR ########## datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs: ########## @@ -0,0 +1,127 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Aggregate table for final aggregation when partial-state input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::InputOrderMode; +use crate::aggregates::{AggregateExec, AggregateMode}; + +use super::common_ordered::{ + OrderedAggregateTable, OrderedFinalMarker, remove_emitted_groups, +}; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. Review Comment: It might also help to add some comments with: 1. Add a link to the `OrderedAggregateTable<OrderedPartialMarker>` 2. Briefly summarize how if differes from the Final table (this one) -- the aggregate mode, no filtering, etc A similar comment on OrderedAggregateTable<OrderedPartialMarker> would be helpful ########## datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs: ########## @@ -0,0 +1,214 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Common utilities for aggregate tables used in aggregations that inputs are ordered +//! by the groups. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +use crate::InputOrderMode; +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::row_hash::create_group_accumulator; +use crate::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, +}; + +use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; + +/// Marker for raw rows -> partial state aggregation on ordered input. +pub(in crate::aggregates) struct OrderedPartialMarker; +/// Marker for partial state -> final value aggregation on ordered input. +pub(in crate::aggregates) struct OrderedFinalMarker; + +/// Aggregate table shared by the ordered partial and final paths. +/// +/// The table consumes input batches while `GroupOrdering` tracks which groups +/// are proven complete. Completed groups can be emitted before the input stream +/// ends, which keeps memory bounded by the active ordered key range. +/// +/// # Marker Type +/// +/// `OrderedAggrMode` selects the aggregate semantics. For example, +/// `OrderedAggregateTable::<OrderedPartialMarker>::new(...)` consumes raw rows +/// and emits partial states, while +/// `OrderedAggregateTable::<OrderedFinalMarker>::new_with_input_order(...)` +/// consumes partial states and emits final values. +/// +/// Shared methods live on `impl<T>`; partial/final behavior lives on +/// marker-specific impls. +pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> { + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Group keys, ordering state, and accumulator states. + pub(super) buffer: OrderedAggregateTableBuffer, + + _mode: PhantomData<OrderedAggrMode>, +} + +/// Buffer for the ordered aggregate table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits output rows as soon as the +/// input ordering proves those groups are complete. +/// +/// [`GroupOrdering`] tracks when and how to do early emit. +/// [`GroupValues`] stores the physical group-key layout, while +/// [`datafusion_expr::GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct OrderedAggregateTableBuffer { + /// GROUP BY expressions evaluated against input batches. + pub(super) group_by: Arc<PhysicalGroupBy>, + + /// Tracks how far ordered input allows this table to drain safely. + pub(super) group_ordering: GroupOrdering, + + /// Interned group keys, in the same group-id order used by accumulators. + pub(super) group_values: Box<dyn GroupValues>, + + /// Scratch group id vector for the current input batch. + pub(super) group_indices: Vec<usize>, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec<AggregateAccumulator>, +} + +/// Methods shared by all aggregate modes +impl<AggrMode> OrderedAggregateTable<AggrMode> { + pub(super) fn new_for_mode( Review Comment: One thought I had (for a future PR) is to potentially avoid some of this duplication by using a builder style here instead -- so something like ```rust OrderedAggregateTableBuilder::new() .with_agg(..) ... .build<AggrMode>() ``` ########## datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs: ########## @@ -373,9 +378,6 @@ async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str .await .unwrap(); assert!(collected_running.len() > 2); - // Running should produce more chunk than the usual AggregateExec. - // Otherwise it means that we cannot generate result in running mode. - assert!(collected_running.len() > collected_usual.len()); Review Comment: I think this is actually because the output is not being clamped to the specified record_batch size I made a PR that targets your branch with a proposed fix here: - https://github.com/2010YOUY01/arrow-datafusion/pull/3 ########## datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs: ########## @@ -0,0 +1,127 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Aggregate table for final aggregation when partial-state input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::InputOrderMode; +use crate::aggregates::{AggregateExec, AggregateMode}; + +use super::common_ordered::{ + OrderedAggregateTable, OrderedFinalMarker, remove_emitted_groups, +}; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +impl OrderedAggregateTable<OrderedFinalMarker> { + pub(in crate::aggregates) fn new_with_input_order( + agg: &AggregateExec, + partition: usize, + input_schema: &SchemaRef, + output_schema: SchemaRef, + input_order_mode: &InputOrderMode, + ) -> Result<Self> { + Self::new_for_mode( + agg, + partition, + input_schema, + output_schema, + input_order_mode, + &AggregateMode::Final, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Merges one partial-state input batch and updates ordering information for + /// any newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + // `PhysicalGroupBy::as_final()` ensures it removes grouping set when + // planning final aggregate, so it's safe to reuse here. + debug_assert_eq!(evaluated_batch.grouping_set_args.len(), 1); + + for group_values in &evaluated_batch.grouping_set_args { + let starting_num_groups = self.buffer.group_values.len(); + self.buffer + .group_values + .intern(group_values, &mut self.buffer.group_indices)?; + let total_num_groups = self.buffer.group_values.len(); + if total_num_groups > starting_num_groups { + self.buffer.group_ordering.new_groups( + group_values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + + let timer = self.group_by_metrics.aggregation_time.timer(); + for (acc, values) in self + .buffer + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + acc.merge_batch(values, &self.buffer.group_indices, total_num_groups)?; Review Comment: likewise the call here to 'merge_batch' is most of the difference -- it would be great to find some way to avoid this duplication I think ########## datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs: ########## @@ -0,0 +1,127 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Aggregate table for final aggregation when partial-state input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::InputOrderMode; +use crate::aggregates::{AggregateExec, AggregateMode}; + +use super::common_ordered::{ + OrderedAggregateTable, OrderedFinalMarker, remove_emitted_groups, +}; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +impl OrderedAggregateTable<OrderedFinalMarker> { + pub(in crate::aggregates) fn new_with_input_order( + agg: &AggregateExec, + partition: usize, + input_schema: &SchemaRef, + output_schema: SchemaRef, + input_order_mode: &InputOrderMode, + ) -> Result<Self> { + Self::new_for_mode( + agg, + partition, + input_schema, + output_schema, + input_order_mode, + &AggregateMode::Final, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Merges one partial-state input batch and updates ordering information for + /// any newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + // `PhysicalGroupBy::as_final()` ensures it removes grouping set when + // planning final aggregate, so it's safe to reuse here. + debug_assert_eq!(evaluated_batch.grouping_set_args.len(), 1); + + for group_values in &evaluated_batch.grouping_set_args { + let starting_num_groups = self.buffer.group_values.len(); + self.buffer + .group_values + .intern(group_values, &mut self.buffer.group_indices)?; + let total_num_groups = self.buffer.group_values.len(); + if total_num_groups > starting_num_groups { + self.buffer.group_ordering.new_groups( + group_values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + + let timer = self.group_by_metrics.aggregation_time.timer(); + for (acc, values) in self + .buffer + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + acc.merge_batch(values, &self.buffer.group_indices, total_num_groups)?; + } + drop(timer); + } + + Ok(()) + } + + /// See comments in `ordered_partial_stream::next_output_batch` + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let Some(emit_to) = self.buffer.group_ordering.emit_to() else { + return Ok(None); + }; + + let timer = self.group_by_metrics.emitting_time.timer(); + let mut output = self.buffer.group_values.emit(emit_to)?; + remove_emitted_groups(&mut self.buffer.group_ordering, emit_to); + + for acc in &mut self.buffer.accumulators { Review Comment: I think this is the only difference between this code and ` OrderedAggregateTable<OrderedPartialMarker>::next_output_batch`: which does ```rust for acc in &mut self.buffer.accumulators { output.extend(acc.state(emit_to)?); } ``` I know that some copying is planned / expected but I do worry there is so much. Maybe we can as a follow on PR consolidate most of the code into a shared implementation on `OrderedAggregateTable` and isolate the differences into a small closure / function ########## datafusion/physical-plan/src/aggregates/mod.rs: ########## @@ -1022,12 +1034,24 @@ impl AggregateExec { .execution .enable_migration_aggregate { + if self.should_use_ordered_partial_aggregate_stream() { Review Comment: It might help a future reader to leave a comment here explaining what is going on Something like ```rust // pick which specialized implementation ``` ########## datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs: ########## @@ -0,0 +1,214 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Common utilities for aggregate tables used in aggregations that inputs are ordered +//! by the groups. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +use crate::InputOrderMode; +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::row_hash::create_group_accumulator; +use crate::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, +}; + +use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; + +/// Marker for raw rows -> partial state aggregation on ordered input. +pub(in crate::aggregates) struct OrderedPartialMarker; +/// Marker for partial state -> final value aggregation on ordered input. +pub(in crate::aggregates) struct OrderedFinalMarker; + +/// Aggregate table shared by the ordered partial and final paths. +/// +/// The table consumes input batches while `GroupOrdering` tracks which groups +/// are proven complete. Completed groups can be emitted before the input stream +/// ends, which keeps memory bounded by the active ordered key range. +/// +/// # Marker Type +/// +/// `OrderedAggrMode` selects the aggregate semantics. For example, +/// `OrderedAggregateTable::<OrderedPartialMarker>::new(...)` consumes raw rows +/// and emits partial states, while +/// `OrderedAggregateTable::<OrderedFinalMarker>::new_with_input_order(...)` +/// consumes partial states and emits final values. +/// +/// Shared methods live on `impl<T>`; partial/final behavior lives on +/// marker-specific impls. +pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> { + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Group keys, ordering state, and accumulator states. + pub(super) buffer: OrderedAggregateTableBuffer, + + _mode: PhantomData<OrderedAggrMode>, +} + +/// Buffer for the ordered aggregate table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits output rows as soon as the +/// input ordering proves those groups are complete. +/// +/// [`GroupOrdering`] tracks when and how to do early emit. +/// [`GroupValues`] stores the physical group-key layout, while +/// [`datafusion_expr::GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct OrderedAggregateTableBuffer { + /// GROUP BY expressions evaluated against input batches. + pub(super) group_by: Arc<PhysicalGroupBy>, + + /// Tracks how far ordered input allows this table to drain safely. + pub(super) group_ordering: GroupOrdering, + + /// Interned group keys, in the same group-id order used by accumulators. + pub(super) group_values: Box<dyn GroupValues>, + + /// Scratch group id vector for the current input batch. + pub(super) group_indices: Vec<usize>, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec<AggregateAccumulator>, +} + +/// Methods shared by all aggregate modes +impl<AggrMode> OrderedAggregateTable<AggrMode> { + pub(super) fn new_for_mode( + agg: &AggregateExec, + partition: usize, + input_schema: &SchemaRef, + output_schema: SchemaRef, + input_order_mode: &InputOrderMode, + aggregate_mode: &AggregateMode, + filters: Vec<Option<Arc<dyn PhysicalExpr>>>, + ) -> Result<Self> { + let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_schema = agg.group_by.group_schema(input_schema)?; + let group_values = new_group_values(group_schema, &group_ordering)?; + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + aggregate_mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(AggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::<Result<_>>()?; + + Ok(Self { + output_schema, + group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), + buffer: OrderedAggregateTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_ordering, + group_values, + group_indices: vec![], + accumulators, + }, + _mode: PhantomData, + }) + } + + /// Evaluates all group by keys and accumulator args. + /// + /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function + /// evaluates `k+1`, `v*v`. + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result<EvaluatedAggregateBatch> { + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + let grouping_set_args = evaluate_group_by(&self.buffer.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + let accumulator_args = self + .buffer + .accumulators + .iter() + .map(|acc| acc.evaluate_acc_args(batch)) + .collect::<Result<Vec<_>>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Called after the input stream is exhausted and the last batch has been + /// aggregated. + /// + /// Updates the internal `GroupOrdering` so it can continue emitting until + /// the buffer is empty. + pub(in crate::aggregates) fn input_done(&mut self) { + self.buffer.group_ordering.input_done(); + } + + /// Check if there is zero groups accumulated so far. + pub(in crate::aggregates) fn is_empty(&self) -> bool { + self.buffer.group_values.is_empty() + } + + /// All internal buffer's memory size. + pub(in crate::aggregates) fn memory_size(&self) -> usize { + self.buffer + .accumulators + .iter() + .map(|acc| acc.size()) + .sum::<usize>() + + self.buffer.group_values.size() + + self.buffer.group_ordering.size() + + self.buffer.group_indices.allocated_size() + } +} + +pub(super) fn remove_emitted_groups(group_ordering: &mut GroupOrdering, emit_to: EmitTo) { Review Comment: It is strange to have a method that masks / aliases `remove_groups` 🤔 Maybe it would be clearer to directly inline it at the callsites ```rust if let EmitTo::First(n) => { group_ordering.remove_groups(n) } ``` ########## datafusion/physical-plan/src/aggregates/mod.rs: ########## @@ -3326,6 +3378,236 @@ mod tests { Ok(()) } + /// Ensures for ordered input, `OrderedPartilAggregateStream` is used. Review Comment: typo: ```suggestion /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. ``` ########## datafusion/physical-plan/src/aggregates/mod.rs: ########## @@ -1069,6 +1101,16 @@ impl AggregateExec { && self.group_by.is_single() } + fn should_use_ordered_final_aggregate_stream(&self) -> bool { + matches!( + self.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned Review Comment: codex also notes that the exising code doesn't use the partial group by operator when memory limited, as the new tables don't seem to handle the OOM path (though maybe I missed that) This is somewhat unobviously implemented by `OutOfMemoryMode` choice here - https://github.com/apache/datafusion/blob/59864f2b751ce94ec786ca7104411ca12a95b044/datafusion/physical-plan/src/aggregates/row_hash.rs#L492-L511 And a proposed PR to fix - https://github.com/2010YOUY01/arrow-datafusion/pull/4 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
