Merge branch 'main' into parquet-row-filter-struct-access-tree

Author: SubhamSinghal

datafusion/physical-plan/Cargo.toml

@@ -123,6 +123,7 @@ required-features = ["test_utils"]
 [[bench]]
 harness = false
 name = "aggregate_vectorized"
+required-features = ["test_utils"]
 
 [[bench]]
 harness = false

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -182,6 +182,9 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                 acc + state.group_values.size()
                     + state.batch_group_indices.allocated_size()
             }
+            AggregateHashTableState::OutputtingMaterializedFinal(output) => {
+                output.memory_size()
+            }
             AggregateHashTableState::Done => 0,
         }
     }
@@ -297,11 +300,53 @@ pub(super) struct AggregateHashTableBuffer {
 }
 
 pub(super) enum AggregateHashTableState {
+    /// Accumulating input rows into group keys and aggregate state.
     Building(AggregateHashTableBuffer),
+    /// Emitting results directly from group keys and aggregate state.
     Outputting(AggregateHashTableBuffer),
+    /// Materialize all the output results, and then incrementally output in the `OutputtingMaterializedFinal` state.
+    ///
+    /// Note this is a temporary solution until the `GroupValues` issue is solved:
+    /// Issue: <https://github.com/apache/datafusion/issues/23178>
+    OutputtingMaterializedFinal(MaterializedFinalOutput),
     Done,
 }
 
+/// Fully evaluated final aggregate output and the next row offset to emit.
+///
+/// Final aggregate evaluation consumes accumulator state, so final output is
+/// materialized once and then sliced to honor `batch_size` across output polls.
+pub(super) struct MaterializedFinalOutput {
+    batch: RecordBatch,
+    offset: usize,
+}
+
+impl MaterializedFinalOutput {
+    pub(super) fn new(batch: RecordBatch) -> Self {
+        Self { batch, offset: 0 }
+    }
+
+    pub(super) fn next_batch(&mut self, batch_size: usize) -> Option<RecordBatch> {
+        debug_assert!(batch_size > 0);
+        if self.is_exhausted() {
+            return None;
+        }
+
+        let length = batch_size.min(self.batch.num_rows() - self.offset);
+        let batch = self.batch.slice(self.offset, length);
+        self.offset += length;
+        Some(batch)
+    }
+
+    pub(super) fn is_exhausted(&self) -> bool {
+        self.offset >= self.batch.num_rows()
+    }
+
+    pub(super) fn memory_size(&self) -> usize {
+        self.batch.get_array_memory_size()
+    }
+}
+
 impl HashAggregateAccumulator {
     fn new(
         aggregate_expr: Arc<AggregateFunctionExpr>,
@@ -440,3 +485,44 @@ impl AggregateHashTableState {
         state
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::array::{Array, Int32Array};
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    use super::*;
+
+    #[test]
+    fn materialized_final_output_slices_batches_until_exhausted() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "group_col",
+            DataType::Int32,
+            false,
+        )]));
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
+        )?;
+        let mut output = MaterializedFinalOutput::new(batch);
+
+        assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]);
+        assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]);
+        assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![5]);
+        assert!(output.next_batch(2).is_none());
+        assert!(output.is_exhausted());
+
+        Ok(())
+    }
+
+    fn int32_values(batch: &RecordBatch, column: usize) -> Vec<i32> {
+        let array = batch
+            .column(column)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+        (0..array.len()).map(|idx| array.value(idx)).collect()
+    }
+}

datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs

@@ -20,11 +20,13 @@ use std::sync::Arc;
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::{Result, internal_err};
+use datafusion_expr::EmitTo;
 
 use crate::aggregates::AggregateExec;
 
 use super::common::{
-    AggregateHashTable, AggregateHashTableState, FinalMarker, emit_to_for_batch_size,
+    AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker,
+    MaterializedFinalOutput,
 };
 
 /// Methods specific to the aggregate hash table used in the final aggregation stage.
@@ -55,30 +57,19 @@ impl AggregateHashTable<FinalMarker> {
     ) -> Result<Option<RecordBatch>> {
         let output_schema = Arc::clone(&self.output_schema);
         let batch_size = self.batch_size;
-        match &mut self.state {
+        // Take ownership of the output state. Note `emit_next_materialized_batch`
+        // updates state after it emits a materialized slice.
+        match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
             AggregateHashTableState::Outputting(state) => {
                 if state.group_values.is_empty() {
-                    self.state = AggregateHashTableState::Done;
                     return Ok(None);
                 }
 
-                let emit_to =
-                    emit_to_for_batch_size(batch_size, state.group_values.len());
-                let timer = self.group_by_metrics.emitting_time.timer();
-                let mut output = state.group_values.emit(emit_to)?;
-
-                for acc in state.accumulators.iter_mut() {
-                    output.push(acc.evaluate(emit_to)?);
-                }
-                let done = state.group_values.is_empty();
-                drop(timer);
-
-                let batch = RecordBatch::try_new(output_schema, output)?;
-                debug_assert!(batch.num_rows() > 0);
-                if done {
-                    self.state = AggregateHashTableState::Done;
-                }
-                Ok(Some(batch))
+                let output = self.materialize_final_output(state, output_schema)?;
+                Ok(self.emit_next_materialized_batch(output, batch_size))
+            }
+            AggregateHashTableState::OutputtingMaterializedFinal(output) => {
+                Ok(self.emit_next_materialized_batch(output, batch_size))
             }
             AggregateHashTableState::Done => Ok(None),
             AggregateHashTableState::Building(_) => {
@@ -87,6 +78,41 @@ impl AggregateHashTable<FinalMarker> {
         }
     }
 
+    fn materialize_final_output(
+        &self,
+        mut state: AggregateHashTableBuffer,
+        output_schema: SchemaRef,
+    ) -> Result<MaterializedFinalOutput> {
+        // Final aggregate evaluation consumes accumulator state. Evaluate all
+        // groups once, then slice the materialized batch on subsequent polls.
+        let emit_to = EmitTo::All;
+        let timer = self.group_by_metrics.emitting_time.timer();
+        let mut output = state.group_values.emit(emit_to)?;
+
+        for acc in state.accumulators.iter_mut() {
+            output.push(acc.evaluate(emit_to)?);
+        }
+        drop(timer);
+
+        let batch = RecordBatch::try_new(output_schema, output)?;
+        debug_assert!(batch.num_rows() > 0);
+        Ok(MaterializedFinalOutput::new(batch))
+    }
+
+    fn emit_next_materialized_batch(
+        &mut self,
+        mut output: MaterializedFinalOutput,
+        batch_size: usize,
+    ) -> Option<RecordBatch> {
+        let batch = output.next_batch(batch_size);
+        if output.is_exhausted() {
+            self.state = AggregateHashTableState::Done;
+        } else {
+            self.state = AggregateHashTableState::OutputtingMaterializedFinal(output);
+        }
+        batch
+    }
+
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,

datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs

@@ -91,6 +91,11 @@ impl AggregateHashTable<PartialMarker> {
             AggregateHashTableState::Building(_) => {
                 internal_err!("next_output_batch must be called in the outputting state")
             }
+            AggregateHashTableState::OutputtingMaterializedFinal(_) => {
+                internal_err!(
+                    "partial aggregate output should not materialize final output"
+                )
+            }
         }
     }
 

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

@@ -1009,6 +1009,13 @@ impl AggregateExec {
             ));
         }
 
+        // `GroupedHashAggregateStream` is being incrementally refactored. See the
+        // tracking issue for details.
+        //
+        // New features and improvements should go directly into the new implementation.
+        // Please coordinate through the tracking issue.
+        //
+        // Issue: <https://github.com/apache/datafusion/issues/22710>
         if context
             .session_config()
             .options()
@@ -1028,7 +1035,7 @@ impl AggregateExec {
             }
         }
 
-        // grouping by something else and we need to just materialize all results
+        // Execution paths that have not been migrated use the fallback implementation
         Ok(StreamType::GroupedHash(GroupedHashAggregateStream::new(
             self, context, partition,
         )?))

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

@@ -132,6 +132,16 @@ enum OutOfMemoryMode {
 
 /// HashTable based Grouping Aggregator
 ///
+/// # Development Note
+///
+/// This implementation is being incrementally refactored. See the tracking issue
+/// for details.
+///
+/// New features and improvements should go directly into the new implementation.
+/// Please coordinate through the tracking issue.
+///
+/// Issue: <https://github.com/apache/datafusion/issues/22710>
+///
 /// # Design Goals
 ///
 /// This structure is designed so that updating the aggregates can be