From dac55696b45e322ddde737b45a3fcf740adbd767 Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Mon, 17 Nov 2025 12:09:34 +0000 Subject: [PATCH] Update aggregate probe to be locked only if skipping --- datafusion/physical-plan/src/aggregates/row_hash.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs index 2e1b70da284d..94a53bfd5a82 100644 --- a/datafusion/physical-plan/src/aggregates/row_hash.rs +++ b/datafusion/physical-plan/src/aggregates/row_hash.rs @@ -194,7 +194,9 @@ impl SkipAggregationProbe { if self.input_rows >= self.probe_rows_threshold { self.should_skip = self.num_groups as f64 / self.input_rows as f64 >= self.probe_ratio_threshold; - self.is_locked = true; + // Set is_locked to true only if we have decided to skip, otherwise we can try to skip + // during processing the next record_batch. + self.is_locked = self.should_skip; } }