zhuqi-lucas commented on code in PR #25098:
URL: https://github.com/apache/datafusion/pull/25098#discussion_r4012608542
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -5312,3 +5339,260 @@ fn
ensure_distribution_reuses_plan_arc_when_no_redistribution_needed() -> Result
);
Ok(())
}
+
+/// Single-child pass-through whose `statistics_from_inputs` increments a
counter
+/// every time it is actually computed (i.e. on a statistics-cache miss). Used
to
+/// observe how often `ensure_distribution` recomputes a node's statistics.
+#[derive(Debug)]
+struct CountingStatsExec {
+ input: Arc<dyn ExecutionPlan>,
+ cache: Arc<PlanProperties>,
+ calls: Arc<AtomicUsize>,
+}
+
+impl CountingStatsExec {
+ fn new(input: Arc<dyn ExecutionPlan>, calls: Arc<AtomicUsize>) -> Self {
+ let cache = PlanProperties::new(
+ input.equivalence_properties().clone(),
+ input.output_partitioning().clone(),
+ input.pipeline_behavior(),
+ input.boundedness(),
+ );
+ Self {
+ input,
+ cache: Arc::new(cache),
+ calls,
+ }
+ }
+}
+
+impl DisplayAs for CountingStatsExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "CountingStatsExec")
+ }
+}
+
+impl ExecutionPlan for CountingStatsExec {
+ fn name(&self) -> &'static str {
+ "CountingStatsExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ assert_eq!(children.len(), 1);
+ Ok(Arc::new(Self::new(
+ children.pop().unwrap(),
+ Arc::clone(&self.calls),
+ )))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) ->
Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion::execution::context::TaskContext>,
+ ) -> Result<datafusion_physical_plan::SendableRecordBatchStream> {
+ unreachable!();
+ }
+
+ fn statistics_from_inputs(
+ &self,
+ _input_stats: &[Arc<Statistics>],
+ _args: &datafusion_physical_plan::statistics::StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ self.calls.fetch_add(1, Ordering::Relaxed);
+ Ok(Arc::new(Statistics::new_unknown(
+ self.input.schema().as_ref(),
+ )))
+ }
+}
+
+/// Regression test for the shared statistics cache in `ensure_distribution`.
+///
+/// A deep stack of pass-through operators sits over a counting leaf. Each
+/// ancestor's distribution enforcement inspects its child's statistics, which
+/// recurse to the leaf. With one `StatisticsContext` shared across the pass
the
+/// leaf is computed once; with a fresh context per node it is recomputed once
+/// per ancestor. This directly detects a regression where the cache is not
+/// actually shared (e.g. reset on every node), which no plan-output assertion
+/// can catch because the optimized plan is identical either way.
+#[test]
+fn ensure_distribution_shares_statistics_cache() -> Result<()> {
+ // Count how many times a leaf's statistics are computed while
+ // `ensure_distribution` runs over a stack of `depth` pass-through
operators
+ // sitting on top of it. Each ancestor's distribution enforcement inspects
+ // its child's statistics, which recurse to the leaf.
+ //
+ // `shared` uses one `StatisticsContext` for the whole pass (what
+ // `EnsureRequirements` does); `fresh` allocates a new context per node
(the
+ // behavior before this change). Returns (shared_computes, fresh_computes).
+ fn run(depth: usize) -> Result<(usize, usize)> {
+ fn deep_plan(depth: usize, calls: &Arc<AtomicUsize>) -> Arc<dyn
ExecutionPlan> {
+ let mut plan: Arc<dyn ExecutionPlan> =
+ Arc::new(CountingStatsExec::new(parquet_exec(),
Arc::clone(calls)));
+ for _ in 0..depth {
+ plan = filter_exec(plan);
+ }
+ plan
+ }
+
+ let mut config = ConfigOptions::new();
+ config.execution.target_partitions = 10;
+ // Keep the plan a fixpoint so no node is rebuilt and the shared cache
is
+ // never reset; statistics are still computed for the round-robin
decision.
+ config.optimizer.enable_round_robin_repartition = false;
+
+ let shared_calls = Arc::new(AtomicUsize::new(0));
+ let stats_ctx =
datafusion_physical_plan::statistics::StatisticsContext::new();
+ DistributionContext::new_default(deep_plan(depth,
&shared_calls)).transform_up(
+ |ctx| {
+ // Reset only when the node's plan pointer actually changed,
exactly
+ // as `EnsureRequirements` does (a rewrite can free a cached
node).
+ let before = Arc::clone(&ctx.plan);
+ let result = ensure_distribution(
+ ctx,
+ &ConfigOnlyContext::new(&config),
+ &stats_ctx,
+ )?;
+ if !Arc::ptr_eq(&before, &result.data.plan) {
+ stats_ctx.reset_cache();
+ }
+ Ok(result)
+ },
+ )?;
Review Comment:
Done — the measured arm now runs `EnsureRequirements::new().optimize(...)`.
I kept the fresh-context-per-node arm as the baseline, since that models the
pre-PR behaviour rather than the code under test.
Counts through the rule come out identical to the old hand-rolled shared arm
(9 at depth 4, 25 at depth 12) against 12/36 for the baseline, so switching
lost no coverage.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1337,14 +1338,16 @@ fn enforce_distribution_relationships(
)]
pub fn ensure_distribution(
dist_context: DistributionContext,
- config: &ConfigOptions,
+ context: &dyn PhysicalOptimizerContext,
+ stats_ctx: &StatisticsContext,
Review Comment:
It can't quite live there: `PhysicalOptimizerContext` is `Send + Sync` and
`StatisticsContext` holds an `Rc<RefCell<StatsCache>>`, so an implementor could
not store one. It also should not outlive a single plan — the cache is keyed by
raw node pointers, which is why it gets reset whenever one changes.
The concern is right though, so I went the other way instead:
`ensure_distribution` only ever used the context for `config_options()`, so it
takes `&ConfigOptions` again plus just the `stats_ctx`. It now never sees a
`PhysicalOptimizerContext` to disagree with, and the registry/cache pairing
happens once, in `optimize_with_context`.
--
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]
