zhuqi-lucas commented on code in PR #25929:
URL: https://github.com/apache/datafusion/pull/25929#discussion_r4163458608
##########
datafusion/physical-plan/src/statistics.rs:
##########
@@ -207,19 +224,80 @@ impl StatisticsContext {
plan: &dyn ExecutionPlan,
args: &StatisticsArgs,
) -> Result<Arc<Statistics>> {
- self.compute_base(plan, args, false)
+ self.compute_base(plan, None, args, false)
+ .map(|(statistics, _)| statistics)
+ }
+
+ /// Like [`Self::compute`], but the cache retains `plan` so the root's own
+ /// statistics are memoized as well as its descendants'.
+ ///
+ /// Prefer this when sharing one context across repeated calls on the same
+ /// nodes, such as an optimizer pass.
+ ///
+ /// # Example
+ ///
+ /// ```
+ /// # use std::sync::Arc;
+ /// # use arrow::datatypes::{DataType, Field, Schema};
+ /// # use datafusion_common::Statistics;
+ /// # use datafusion_common::stats::Precision;
+ /// # use datafusion_physical_plan::ExecutionPlan;
+ /// # use datafusion_physical_plan::statistics::{StatisticsArgs,
StatisticsContext};
+ /// # use datafusion_physical_plan::test::exec::StatisticsExec;
+ ///
+ /// let schema = Schema::new(vec![Field::new("a", DataType::Int32,
false)]);
+ /// let stats =
Statistics::new_unknown(&schema).with_num_rows(Precision::Exact(100));
+ /// let plan: Arc<dyn ExecutionPlan> = Arc::new(StatisticsExec::new(stats,
schema));
+ ///
+ /// let context = StatisticsContext::new();
+ /// let first = context.compute_arc(&plan, &StatisticsArgs::new())?;
+ /// let second = context.compute_arc(&plan, &StatisticsArgs::new())?;
+ ///
+ /// // The second call is a cache hit for the root.
+ /// assert!(Arc::ptr_eq(&first, &second));
+ /// assert_eq!(first.num_rows, Precision::Exact(100));
+ /// # Ok::<(), datafusion_common::DataFusionError>(())
+ /// ```
+ pub fn compute_arc(
Review Comment:
There is a ready-made caller for this: `enforce_distribution.rs:1012` does
`stats_ctx.compute(child.as_ref(), ..)` where `child` is already an `Arc`, so
under the new rule that root is looked up but never inserted — and that call
site is exactly the repeated-subtree case #25098 added the shared context for.
Switching it to `compute_arc(&child, ..)` restores root caching.
Happy to do that as a follow-up, together with dropping the `Arc::ptr_eq` +
`reset_cache()` guard in `ensure_requirements/mod.rs` now that resetting is no
longer required for correctness — unless you would rather fold either into this
PR.
--
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]
