kosiew opened a new pull request, #25929:
URL: https://github.com/apache/datafusion/pull/25929
---
```markdown
# Title: Make StatisticsContext cache entries own their node for
cross-rewrite safety
## Changes
- **Update `StatsCache` entries to carry the owning `Arc` to the node**
whose address is the key. The `CacheEntry` now includes `_plan: Arc<dyn
ExecutionPlan>`, and the value is `Arc<Statistics>` or `Arc<Extensions>` based
on the map. The `CacheKey` remains address-based.
- **Modify `store_statistics` and `store_extensions` to require an owning
`Arc`** for the node, preventing any entry from being written without a parent.
The `debug_assert!` ensures the owner’s data pointer matches the key.
- **Thread `owner: Option<&Arc<dyn ExecutionPlan>>` through `compute_base`
and `try_provider_stats`**. `resolve_children` now passes `Some(child)`; only
the root node is stored when `owner` is `Some`.
- **Add `pub fn compute_arc(&self, plan: &Arc<dyn ExecutionPlan>, args:
&StatisticsArgs) -> Result<Arc<Statistics>>`** with Rustdoc and a doctest. This
function is used to cache the root node with its owner.
- **Update documentation and type docs** to reflect the ownership invariant
in `StatsCache`, `StatisticsContext`, and function signatures. `reset_cache`
remains optional; context now holds every cached node until dropped or reset.
- **Ensure all calls to `compute` on borrowed roots (children or parents)**
return the cached value via `Arc::ptr_eq`, but do not store the root. The root
is only stored via `compute_arc`.
## Testing
New tests in `statistics.rs`:
- `context_cache_retains_plan_lifetime`:
- Build `parent -> leaf` as `Arc`s, call `ctx.compute_arc(&parent, ...)`,
and take `Weak`s of both.
- Drop the caller's `Arc`s and assert both `Weak`s upgrade.
- Drop `ctx` and assert both are dead.
- Repeat with `reset_cache()` in place of drop.
- `borrowed_root_is_not_stored`:
- Borrowed `compute(leaf.as_ref())` on a childless leaf leaves
`cache.statistics` empty.
- `compute_arc(&leaf)` populates it.
- A later borrowed `compute` returns the `Arc::ptr_eq` cached value.
Existing tests:
- Update `context_caches_within_walk` (`statistics.rs:603-613`) to use
`compute_arc`. Its borrowed-root expectation is altered as per this change.
- The extension/provider tests pass unchanged.
## Acceptance Criteria
- No code path inserts a `statistics` or `extensions` entry without an
owning `Arc` to the keyed node. Enforced by the `store_*` signatures.
- A node cached by a context stays alive until that context is dropped or
`reset_cache` is called, and no longer after that.
- `compute(&dyn)` on a node whose entry already exists (cached as a child or
via `compute_arc`) returns the cached `Arc<Statistics>` (`Arc::ptr_eq`).
- `compute` / `compute_extended` signatures are unchanged; `reset_cache` is
retained; `compute_arc` is the only new public item.
## Risks
- **Memory**: retained nodes live as long as the context. Every in-tree
context is either per-call or per-pass.
- **Behavior change**: Repeated borrowed `compute` calls on the same root
now recompute that one node from cached children. Mention this in the upgrade
guide only if the bench shows a measurable cost.
## Dependencies / Blockers
- SUB_ISSUE_01 (root extensions returned from the walk rather than re-read
from the cache).
```
--
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]
