asolimando commented on code in PR #25098:
URL: https://github.com/apache/datafusion/pull/25098#discussion_r4006308153
##########
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:
It think we should replace the closure with an invocation of the rule
itself, something like:
`EnsureRequirements::new().optimize(deep_plan(depth, &shared_calls),
&config)?;`
Otherwise we risk of not catching regressions on the rule itself, unless I
am missing something
##########
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:
`StatisticsContext` is already part of `PhysicalOptimizerContext`, you
simply get it from there, so you avoid a caller passing a different
`StatisticsContext` than that of `PhysicalOptimizerContext` (for which I can't
think of a legal use-case).
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1355,6 +1356,7 @@ fn enforce_distribution_relationships(
pub fn ensure_distribution(
dist_context: DistributionContext,
config: &ConfigOptions,
+ stats_ctx: &StatisticsContext,
) -> Result<Transformed<DistributionContext>> {
Review Comment:
I think this comment is correct, and we should keep the old signature (and
deprecated it), then adding notes on the upgrading guide.
--
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]
