Expose accumulator state to allow prefix scanning - #24035
Conversation
ae37c0b to
65014d5
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24035 +/- ##
==========================================
+ Coverage 80.99% 81.01% +0.02%
==========================================
Files 1106 1106
Lines 383352 384686 +1334
Branches 383352 384686 +1334
==========================================
+ Hits 310488 311665 +1177
- Misses 54544 54620 +76
- Partials 18320 18401 +81 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Thanks for this contribution @avantgardnerio ! I'd find it helpful if you'd elaborate a little bit more about the motivation for this change in the PR description. For example, some intended use-cases, what kind of performance improvement this unlocks, etc. |
|
@neilconway I'm trying to speed up window functions using parallel prefix scans. I am presently incubating this in Ballista, and this is the minimum API exposure that I need to do it for non-decomposable operations like The jury is still out about re-partition cost vs performance benefit, but the signs are hopeful:
And at least from a big-O time perspective (table 1) it should be optimal for some queries ( |
0767f73 to
13023b2
Compare
|
@gene-bordegaray and @JSOD11 you guys might be interested as well. |
alamb
left a comment
There was a problem hiding this comment.
Thanks @avantgardnerio and @neilconway -- I left some comments
alamb
left a comment
There was a problem hiding this comment.
I may not fully understand prefix scanning, but it seems to me like this API will only give you access to the window state for the single last row in each partition.
Don't you potentially need access to the window state for the last N rows in a partition (e.g the HALO rows) 🤔
|
FWIW claude claims this doesn't get run with windows like UNBOUNDED PRECEDING → CURRENT ROW |
|
I think it would also be super useful to add some sort of example / test that shows how you intend to use this API (for exmple some simple example for computing a window function in parallel or something 🤔 that way we could see the API in action |
Thanks @alamb ! That was a critical bug that would have defeated the whole point. It is now fixed and asserted in |
Yes, this is exactly what is required.
No, not for prefix scanning. (answer below)
The HPC "halo" term is a good fit for bounded preceding/following (surrounding cells, in 1D) but doesn't extend cleanly to "last row of every other partition." Regardless of the name, this PR doesn't take that approach - because although it works for SUM, and decomposes for AVG (sum+count), it fails by the time you get to arbitrary accumulators like approx_distinct.
Which is exactly where (the newly added) Edit: added the qualifier (DF) partition to distinguish between the ambiguously named (SQL window) partition. |
|
Pseudo code, stripped directly from the |
Match the field's type so with_new_children collapses to a single chained call and the setter can also clear a previously-installed observer. Addresses apache#24035 review comment 3738803158. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The method now always mutates when called and takes the observer as a required argument; the "is observer installed?" check moves to the caller in `compute_aggregates`. Removes the "&mut self that only mutates when observer is set" shape. Addresses apache#24035 review comment 3738822701. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Rename the trait method to `finalize_window_aggregate` and split its signature so the callback fires once per aggregate window expression per closing PARTITION BY group, receiving that expression's Arc and its own `Accumulator::state` directly. Non-aggregate window functions no longer fire the callback at all. Removes the per-partition-key `Vec<Option<Vec<ScalarValue>>>` wrapper allocation, and gives the observer the window-expression context needed to disambiguate calls when the exec carries multiple window expressions. Leaves room to add a peer `finalize_window_function` later for built-in (non-aggregate) window functions. Addresses apache#24035 review comments 3738790091 and 3738816488. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…close tests The two `test_finalized_state_observer_*` tests were structurally identical apart from the window frame. Fold their common setup and assertions into a single async helper that takes the frame, so each test body is now just the frame construction + a one-line comment explaining which causality regime it exercises. Addresses apache#24035 review comment 3738845164. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
@timsaucer thank you for the thorough review. These are all legitimate concerns:
I think all of the above stem from this trying to fire per-partition-key, rather than once, at the end, when the motivating case is prefix scanning. What would you think about a getter shape where a downstream node reads |
timsaucer
left a comment
There was a problem hiding this comment.
Thank you for all the work on this!
|
Thank you @avantgardnerio and @timsaucer |
|
run benchmark bounded_window |
|
run benchmark window_query_sql |
|
run benchmark h2o_small_window |
|
(just running some benchmarks to make sure this doesn't change performance noticably) |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark window_query_sqlResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowCPU Details (lscpu)Details
Resource Usagebounded_window — base (merge-base)
bounded_window — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowCPU Details (lscpu)Details
Resource Usageh2o_small_window — base (merge-base)
h2o_small_window — branch
File an issue against this benchmark runner |
|
run benchmark h2o_small_window |
|
run benchmark bounded_window |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark bounded_windowCPU Details (lscpu)Details
Resource Usagebounded_window — base (merge-base)
bounded_window — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark h2o_small_windowCPU Details (lscpu)Details
Resource Usageh2o_small_window — base (merge-base)
h2o_small_window — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing brent/bwag-finalized-state-observer (aa4d196) to 33ad1cc (merge-base) diff Run configurationrun benchmark window_query_sqlCPU Details (lscpu)Details
Resource Usagewindow_query_sql — base (merge-base)
window_query_sql — branch
File an issue against this benchmark runner |
## Summary
Expose state of aggregate streams within BWAG so downstream prefix
scanning can take place.
## API
```rust
// physical-plan/src/windows/bounded_window_agg_exec.rs
pub type FinalizedWindowStateObserver = Arc<
dyn Fn(usize, &PartitionKey, &[Option<Vec<ScalarValue>>]) -> Result<()>
+ Send + Sync,
>;
impl BoundedWindowAggExec {
pub fn with_finalized_state_observer(mut self, obs: FinalizedWindowStateObserver) -> Self { … }
}
// physical-expr/src/window/window_expr.rs
impl WindowState {
/// `Accumulator::state()` if this is an aggregate window function, `None` otherwise.
pub fn aggregate_state(&mut self) -> Result<Option<Vec<ScalarValue>>> { … }
}
```

Summary
Expose state of aggregate streams within BWAG so downstream prefix scanning can take place.
API