Skip to content

Commit 44f1022

Browse files
asolimandoasolimando
committed
Pre-compute children with overall stats to avoid wasted partition-specific walks
1 parent c74a90e commit 44f1022

23 files changed

Lines changed: 170 additions & 104 deletions

‎datafusion/physical-optimizer/src/output_requirements.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -245,10 +245,13 @@ impl ExecutionPlan for OutputRequirementExec {
245245

246246
fn partition_statistics_with_context(
247247
&self,
248-
_partition: Option<usize>,
248+
partition: Option<usize>,
249249
ctx: &StatisticsContext,
250250
) -> Result<Arc<Statistics>> {
251-
Ok(Arc::clone(&ctx.child_stats()[0]))
251+
match partition {
252+
Some(_) => ctx.compute_child_statistics(self.input.as_ref(), partition),
253+
None => Ok(Arc::clone(&ctx.child_stats()[0])),
254+
}
252255
}
253256

254257
fn try_swapping_with_projection(

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1646,10 +1646,13 @@ impl ExecutionPlan for AggregateExec {
16461646

16471647
fn partition_statistics_with_context(
16481648
&self,
1649-
_partition: Option<usize>,
1649+
partition: Option<usize>,
16501650
ctx: &StatisticsContext,
16511651
) -> Result<Arc<Statistics>> {
1652-
let child_statistics = Arc::clone(&ctx.child_stats()[0]);
1652+
let child_statistics = match partition {
1653+
Some(_) => ctx.compute_child_statistics(self.input.as_ref(), partition)?,
1654+
None => Arc::clone(&ctx.child_stats()[0]),
1655+
};
16531656
Ok(Arc::new(self.statistics_inner(&child_statistics)?))
16541657
}
16551658

‎datafusion/physical-plan/src/buffer.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -247,10 +247,13 @@ impl ExecutionPlan for BufferExec {
247247

248248
fn partition_statistics_with_context(
249249
&self,
250-
_partition: Option<usize>,
250+
partition: Option<usize>,
251251
ctx: &StatisticsContext,
252252
) -> Result<Arc<Statistics>> {
253-
Ok(Arc::clone(&ctx.child_stats()[0]))
253+
match partition {
254+
Some(_) => ctx.compute_child_statistics(self.input.as_ref(), partition),
255+
None => Ok(Arc::clone(&ctx.child_stats()[0])),
256+
}
254257
}
255258

256259
fn supports_limit_pushdown(&self) -> bool {

‎datafusion/physical-plan/src/coalesce_batches.rs‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -226,10 +226,15 @@ impl ExecutionPlan for CoalesceBatchesExec {
226226

227227
fn partition_statistics_with_context(
228228
&self,
229-
_partition: Option<usize>,
229+
partition: Option<usize>,
230230
ctx: &StatisticsContext,
231231
) -> Result<Arc<Statistics>> {
232-
let stats = Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[0]));
232+
let stats = match partition {
233+
Some(_) => Arc::unwrap_or_clone(
234+
ctx.compute_child_statistics(self.input.as_ref(), partition)?,
235+
),
236+
None => Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[0])),
237+
};
233238
Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
234239
}
235240

‎datafusion/physical-plan/src/coalesce_partitions.rs‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -245,11 +245,7 @@ impl ExecutionPlan for CoalescePartitionsExec {
245245
_partition: Option<usize>,
246246
ctx: &StatisticsContext,
247247
) -> Result<Arc<Statistics>> {
248-
// CoalescePartitions merges all input partitions into one,
249-
// so it always needs overall (None) input stats
250-
let stats = Arc::unwrap_or_clone(
251-
ctx.compute_child_statistics(self.input.as_ref(), None)?,
252-
);
248+
let stats = Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[0]));
253249
Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
254250
}
255251

‎datafusion/physical-plan/src/coop.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -309,10 +309,13 @@ impl ExecutionPlan for CooperativeExec {
309309

310310
fn partition_statistics_with_context(
311311
&self,
312-
_partition: Option<usize>,
312+
partition: Option<usize>,
313313
ctx: &StatisticsContext,
314314
) -> Result<Arc<Statistics>> {
315-
Ok(Arc::clone(&ctx.child_stats()[0]))
315+
match partition {
316+
Some(_) => ctx.compute_child_statistics(self.input.as_ref(), partition),
317+
None => Ok(Arc::clone(&ctx.child_stats()[0])),
318+
}
316319
}
317320

318321
fn supports_limit_pushdown(&self) -> bool {

‎datafusion/physical-plan/src/execution_plan.rs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -574,12 +574,15 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync {
574574
/// (the default), not an error.
575575
/// If `partition` is `None`, it returns statistics for the entire plan.
576576
///
577-
/// The [`StatisticsContext`] carries pre-computed child statistics and
578-
/// additional context for statistics computation. Use
579-
/// [`compute_statistics`] to compute statistics bottom-up, threading
580-
/// child statistics through the context automatically.
577+
/// The [`StatisticsContext`] carries pre-computed overall (`None`) child
578+
/// statistics via [`StatisticsContext::child_stats`] and a shared cache
579+
/// via [`StatisticsContext::compute_child_statistics`]. Operators that
580+
/// need per-partition child stats should call
581+
/// `ctx.compute_child_statistics(child, partition)`.
581582
///
582583
/// [`StatisticsContext`]: crate::statistics_context::StatisticsContext
584+
/// [`StatisticsContext::child_stats`]: crate::statistics_context::StatisticsContext::child_stats
585+
/// [`StatisticsContext::compute_child_statistics`]: crate::statistics_context::StatisticsContext::compute_child_statistics
583586
/// [`compute_statistics`]: crate::statistics_context::compute_statistics
584587
fn partition_statistics_with_context(
585588
&self,

‎datafusion/physical-plan/src/filter.rs‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -580,10 +580,15 @@ impl ExecutionPlan for FilterExec {
580580
/// predicate's selectivity value can be determined for the incoming data.
581581
fn partition_statistics_with_context(
582582
&self,
583-
_partition: Option<usize>,
583+
partition: Option<usize>,
584584
ctx: &StatisticsContext,
585585
) -> Result<Arc<Statistics>> {
586-
let input_stats = Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[0]));
586+
let input_stats = match partition {
587+
Some(_) => Arc::unwrap_or_clone(
588+
ctx.compute_child_statistics(self.input.as_ref(), partition)?,
589+
),
590+
None => Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[0])),
591+
};
587592
let stats = Self::statistics_helper(
588593
&self.input.schema(),
589594
input_stats,

‎datafusion/physical-plan/src/joins/cross_join.rs‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -388,16 +388,16 @@ impl ExecutionPlan for CrossJoinExec {
388388
ctx: &StatisticsContext,
389389
) -> Result<Arc<Statistics>> {
390390
// Left side is always broadcast (collected into a single partition),
391-
// so it needs overall stats. For the Some case, the context has
392-
// partition-specific stats which would be incorrect.
393-
// Right side can have multiple partitions, so the context's
394-
// partition-specific or overall stats are correct as-is.
395-
let left_stats = match partition {
396-
Some(_) => ctx.compute_child_statistics(self.left.as_ref(), None)?,
397-
None => Arc::clone(&ctx.child_stats()[0]),
391+
// so it always needs overall stats (child_stats provides these).
392+
// Right side can have multiple partitions, so it needs per-partition
393+
// stats when a specific partition is requested.
394+
let left_stats = Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[0]));
395+
let right_stats = match partition {
396+
Some(_) => Arc::unwrap_or_clone(
397+
ctx.compute_child_statistics(self.right.as_ref(), partition)?,
398+
),
399+
None => Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[1])),
398400
};
399-
let left_stats = Arc::unwrap_or_clone(left_stats);
400-
let right_stats = Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[1]));
401401

402402
Ok(Arc::new(stats_cartesian_product(left_stats, right_stats)))
403403
}

‎datafusion/physical-plan/src/joins/hash_join/exec.rs‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1468,13 +1468,12 @@ impl ExecutionPlan for HashJoinExec {
14681468
) -> Result<Arc<Statistics>> {
14691469
let stats = match (partition, self.mode) {
14701470
// For CollectLeft mode, the left side is broadcast (collected into
1471-
// a single partition), so it needs overall stats. The context has
1472-
// partition-specific stats which would be incorrect.
1473-
// Right side is partitioned, so the context stats are correct.
1471+
// a single partition), so it needs overall stats (child_stats).
1472+
// Right side is partitioned, so it needs per-partition stats.
14741473
(Some(_), PartitionMode::CollectLeft) => {
1475-
let left_stats =
1476-
ctx.compute_child_statistics(self.left.as_ref(), None)?;
1477-
let right_stats = Arc::clone(&ctx.child_stats()[1]);
1474+
let left_stats = Arc::clone(&ctx.child_stats()[0]);
1475+
let right_stats =
1476+
ctx.compute_child_statistics(self.right.as_ref(), partition)?;
14781477

14791478
estimate_join_statistics(
14801479
Arc::unwrap_or_clone(left_stats),
@@ -1488,8 +1487,10 @@ impl ExecutionPlan for HashJoinExec {
14881487
// For Partitioned mode, both sides are hash-partitioned symmetrically,
14891488
// so each output partition uses the matching partition from both sides.
14901489
(Some(_), PartitionMode::Partitioned) => {
1491-
let left_stats = Arc::clone(&ctx.child_stats()[0]);
1492-
let right_stats = Arc::clone(&ctx.child_stats()[1]);
1490+
let left_stats =
1491+
ctx.compute_child_statistics(self.left.as_ref(), partition)?;
1492+
let right_stats =
1493+
ctx.compute_child_statistics(self.right.as_ref(), partition)?;
14931494

14941495
estimate_join_statistics(
14951496
Arc::unwrap_or_clone(left_stats),
@@ -1505,8 +1506,8 @@ impl ExecutionPlan for HashJoinExec {
15051506
let left_stats = Arc::clone(&ctx.child_stats()[0]);
15061507
let right_stats = Arc::clone(&ctx.child_stats()[1]);
15071508
estimate_join_statistics(
1508-
(*left_stats).clone(),
1509-
(*right_stats).clone(),
1509+
Arc::unwrap_or_clone(left_stats),
1510+
Arc::unwrap_or_clone(right_stats),
15101511
&self.on,
15111512
&self.join_type,
15121513
&self.join_schema,
@@ -1516,10 +1517,8 @@ impl ExecutionPlan for HashJoinExec {
15161517
// Auto mode hasn't decided partitioning yet, so it needs
15171518
// overall stats from both sides.
15181519
(Some(_), PartitionMode::Auto) => {
1519-
let left_stats =
1520-
ctx.compute_child_statistics(self.left.as_ref(), None)?;
1521-
let right_stats =
1522-
ctx.compute_child_statistics(self.right.as_ref(), None)?;
1520+
let left_stats = Arc::clone(&ctx.child_stats()[0]);
1521+
let right_stats = Arc::clone(&ctx.child_stats()[1]);
15231522
estimate_join_statistics(
15241523
Arc::unwrap_or_clone(left_stats),
15251524
Arc::unwrap_or_clone(right_stats),

0 commit comments

Comments
 (0)