Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions native-engine/datafusion-ext-plans/src/agg_exec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,7 @@ async fn execute_agg_with_grouping_hash(
metrics: ExecutionPlanMetricsSet,
) -> Result<SendableRecordBatchStream> {
// create tables
let baseline_metrics = BaselineMetrics::new(&metrics, partition_id);
let tables = Arc::new(AggTable::new(
partition_id,
agg_ctx.clone(),
Expand All @@ -211,6 +212,7 @@ async fn execute_agg_with_grouping_hash(
.map_err(|err| err.context("agg: polling batches from input error"))?
{
// insert or update rows into in-mem table
let _timer = baseline_metrics.elapsed_compute().timer();
tables.process_input_batch(input_batch).await?;

// stop aggregating if triggered partial skipping
Expand All @@ -224,7 +226,10 @@ async fn execute_agg_with_grouping_hash(
// merge all tables and output
let output_schema = agg_ctx.output_schema.clone();
let output = context.output_with_sender("Agg", output_schema, |sender| async move {
sender.exclude_time(baseline_metrics.elapsed_compute());

// output all aggregated records in table
let _timer = baseline_metrics.elapsed_compute().timer();
tables.output(sender.clone()).await?;

// in partial skipping mode, there might be unconsumed records in input stream
Expand Down