Expose accumulator state to allow prefix scanning

[avantgardnerio]

Summary

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

API

// 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>>> { … }
}

[codecov-commenter]

Codecov Report

❌ Patch coverage is 86.98980% with 51 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.05%. Comparing base (3e3a92d) to head (b980ba4).
⚠️ Report is 32 commits behind head on main.

Files with missing lines	Patch %	Lines
...ysical-plan/src/windows/bounded_window_agg_exec.rs	87.08%	16 Missing and 34 partials ⚠️
datafusion/physical-expr/src/window/window_expr.rs	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24035      +/-   ##
==========================================
+ Coverage   80.91%   81.05%   +0.13%     
==========================================
  Files        1103     1106       +3     
  Lines      377219   382677    +5458     
  Branches   377219   382677    +5458     
==========================================
+ Hits       305244   310182    +4938     
- Misses      53775    54142     +367     
- Partials    18200    18353     +153     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[neilconway]

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.

[avantgardnerio]

@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 approx_distinct() (vs others like AVG, STDDEV via Welford / Chan, etc). Though, the work is certainly not limited to Ballista, it has shown improvement in DataFusion as well.

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 (select my_agg() over unbounded preceding...)

[avantgardnerio]

@gene-bordegaray and @JSOD11 you guys might be interested as well.

[alamb]

Thanks @avantgardnerio and @neilconway -- I left some comments

[alamb]

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) 🤔

[alamb]

FWIW claude claims this doesn't get run with windows like

UNBOUNDED PRECEDING → CURRENT ROW

[alamb]

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

[avantgardnerio]

this doesn't get run with windows like

Thanks @alamb ! That was a critical bug that would have defeated the whole point. It is now fixed and asserted in test_finalized_state_observer_fires_on_causal_frame()

[avantgardnerio]

access to the window state for the single last row

Yes, this is exactly what is required.

need access to the window state for the last N rows

No, not for prefix scanning. (answer below)

e.g the HALO rows

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.

add some sort of example / test

Which is exactly where (the newly added) test_prefix_scan_across_tasks_matches_single_bwag() comes in. It shows that with only the accumulator state of the very last row of the entire (DF) partition, parallel partitioned prefix-scanned sums produce exactly the same results as a single (DF) partition BWAG. Given the observer exposes Accumulator::state() directly, any function that supports Accumulator::merge_batch - including approx_distinct - can be prefix-scanned the same way, with the downstream consumer merging sketches instead of adding scalars.

Edit: added the qualifier (DF) partition to distinguish between the ambiguously named (SQL window) partition.

[avantgardnerio]

Pseudo code, stripped directly from the test_prefix_scan_across_tasks_matches_single_bwag test, for those who value brevity:

        // Two tasks under range partition on sn:
        let (task1_out, task1_total) = run_running_sum_task(&[1, 1, 2, 2, 3, 3, 4, 4]);
        let (task2_out, task2_total) = run_running_sum_task(&[5, 5, 6, 6, 7, 7, 8, 8]);

        // Local (uncorrected) outputs and totals — first pass.
        assert_eq!(task1_out, vec![1, 2, 4, 6, 9, 12, 16, 20]);
        assert_eq!(task2_out, vec![5, 10, 16, 22, 29, 36, 44, 52]);

        // Prefix scan over per-task totals → carry-in for each task. Task 0's
        // carry-in is 0; task N's carry-in is the sum of tasks [0, N).
        let carry_ins = [0u64, task1_total];

        // Second pass: shift each task's local values by its carry-in.
        let task1_final: Vec<u64> = task1_out.iter().map(|v| v + carry_ins[0]).collect();
        let task2_final: Vec<u64> = task2_out.iter().map(|v| v + carry_ins[1]).collect();
        let parallel_result: Vec<u64> = task1_final.iter().chain(task2_final.iter());

        // Oracle: single BWAG over the full concatenated input.
        let (single_result, single_total) = run_running_sum_task(
            &[1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8]);

        assert_eq!(parallel_result, single_result);
        assert_eq!(single_result,
            vec![1, 2, 4, 6, 9, 12, 16, 20, 25, 30, 36, 42, 49, 56, 64, 72]
        );

[alamb]

Thanks @avantgardnerio

[avantgardnerio]

Hi @alamb , I appreciate your feedback, and I think I addressed all of it. I also threw in a defensive test for merging HLLs just to be sure. If there's anything I missed, please let me know and I'll address that too. I really appreciate you sticking with this through multiple rounds of reviews 🙂