asolimando commented on PR #25098: URL: https://github.com/apache/datafusion/pull/25098#issuecomment-5600485116
Hi @zhuqi-lucas, I can confirm that your use-case fits what the cache in `StatisticsContext` is made for, so it sounds like a great idea. Glad to read that you are seeing a drop in planning time of ~10% due to caching! We can probably push this a little further: instead of just passing `StatisticsContext` through `ensure_distribution` to share the memoization cache across the pass, you could use `EnsureRequirements::optimize_with_context` (instead of `EnsureRequirements::optimize`), so that you receive a `PhysicalOptimizerContext`, including the `StatisticsRegistry` and any other statistics context useful for CBO. `StatisticsContext::new()` builds an empty `StatisticsRegistry`, so today this rule never consults any registered provider, even when the session has one configured. I think it's worth doing in this PR, rather than a follow-up, because we can avoid changing the public `ensure_distribution` twice. I have implemented something similar in #24716 for `JoinSelection` to showcase how I think this should be done for physical optimizer rules. So, concretely, my suggestion is to: - override `optimize_with_context` on `EnsureRequirements` - pass `context` to `ensure_distribution` and `get_repartition_requirement_status` instead of `stats_ctx` - build the shared `StatisticsContext` from `context.statistics_registry()` once per pass This would keep the caching fix from this PR, and also allow the rule to benefit from the present and future statistics context (I will be working on https://github.com/apache/datafusion/issues/21120 as my next task). Note that there is no behavior change by default: a session without registered providers is unaffected, but it gives room for improvement without further breaking changes. I am off this week with very limited access to a computer, but I can surely offer a review or help with a PR from next week! Regarding the cache reset: your proposed solution seems safe to me, but in the future I'd like to make the cache more robust so consumers don't have to think about it. Either by introducing a unique id per constructed `ExecutionPlan` and using that as cache key, or by preventing the freeing of the deleted nodes for the lifetime of the cache (Arc clone of the node). This is only tangential to this PR, but I'd appreciate your opinion on this matter since you have the needed context already. -- 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]
