GitHub user Subodh-17 created a discussion: [RFC/Discussion]
Partition-at-a-Time Streaming Evaluation in WindowAggExec
# [RFC/Discussion] Partition-at-a-Time Streaming Evaluation in WindowAggExec
- **Status:** Early Design Proposal — Seeking Maintainer & Community Feedback
- **Target Area:** `datafusion-physical-plan` (`WindowAggExec`,
`WindowAggStream`)
- **Author:** Community Contributor
---
## 1. Problem Statement
In Apache DataFusion, window functions that require full partition awareness
(such as `ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING`, `LEAD`,
`LAG`, and `RANK`) are executed by `WindowAggExec`.
Currently, `WindowAggStream` acts as a pipeline-breaking operator:
1. It buffers every incoming `RecordBatch` into `Vec<RecordBatch>` until the
upstream stream yields `None` (End-of-Stream / EOS).
2. It blocks downstream operators from consuming even a single output row until
the entire upstream pipeline terminates.
3. At stream exhaustion, it invokes `concat_batches` across all accumulated
batches, materializing a monolithic batch of $N$ rows before evaluating window
expressions.
```text
CURRENT BEHAVIOR (Whole-Stream Buffering Stall):
Input Batches: [B0] ---> [B1] ---> [B2] ---> [B3] ---> ... ---> [Bn] ---> EOS
Buffering: |----------------- FULL BUFFER (All N rows) ------------------|
Evaluation:
|--> concat_batches(N rows)
Emission:
|--> [Single Monolithic Batch]
Pipeline State: |<--------------------- BLOCKED ----------------------------->|
```
### Operational Drawbacks
- **$O(N)$ Memory Retention:** Working memory scales linearly with total stream
length ($N$), even when partition cardinality is high and individual partitions
are small ($|P_i| \ll N$). In local benchmarks on a 2.5-million-row dataset
containing ~16.6k partitions, baseline DataFusion retained all 2.5M rows in
memory, reaching a peak resident set size (Peak RSS) of ~724 MB.
- **Pipeline Latency:** Downstream operators that could stream data
incrementally (such as `LimitExec`, streaming network sinks, or interactive
displays) are starved of rows until the full upstream input is scanned.
---
## 2. Proposed Approach: Partition-at-a-Time Streaming
DataFusion's physical planner requires inputs to `WindowAggExec` to be
pre-sorted by `PARTITION BY` keys (enforced via `required_input_ordering()`).
Consequently, all rows belonging to any partition $P_k$ arrive contiguously in
the stream.
Once the first row with partition key $P_{k+1}$ is observed, partition $P_k$ is
guaranteed to be complete and will receive no further rows.
We propose updating `WindowAggStream` to evaluate and emit completed partitions
incrementally:
1. **Active Partition Buffer:** Maintain `active_partition_batches:
Vec<RecordBatch>` and `active_partition_key: Option<Vec<ScalarValue>>` for the
currently open partition.
2. **Boundary Detection (`process_batch`):** Inspect partition keys in incoming
batches using `evaluate_partition_ranges`:
- *Intra-batch partitions:* Evaluate partitions that start and end within
the batch directly on `batch.slice(start, len)` without cross-batch
concatenation.
- *Cross-batch partitions:* Accumulate slices until a boundary key
transition occurs, concatenate only that partition's slices, evaluate, and emit.
3. **Output FIFO & Backpressure:** Enqueue evaluated partitions into a
`pending_batches: VecDeque<RecordBatch>`. The stream drains `pending_batches`
before polling upstream input, ensuring upstream ingestion pauses when
downstream consumption is slow.
4. **Unpartitioned Fallback:** Queries without a `PARTITION BY` clause (`OVER
(ORDER BY ...)`) take an explicit code branch (`self.is_unpartitioned()`) that
falls back to the existing whole-stream buffering path, preserving current
behavior. (In contrast, partitioned queries where all rows happen to share a
single partition key run through the streaming state machine, buffering slices
in `active_partition_batches` until EOF).
```text
PROPOSED PARTITION-STREAMING PIPELINE:
Input Batch B0: [ P0 (tail) ][ P1 (complete) ][ P2 (head) ]
| | |
Finalize P0: |---------> |-> Eval & Enqueue -> [pending_batches] ->
Emitted Downstream!
Evaluate P1: |-> Eval & Enqueue -> [pending_batches] ->
Emitted Downstream!
Buffer P2: |-> [active_partition_batches]
Working Memory: Bounded by max(|P_i|) + BatchSize
```
---
## 3. Working Memory Model & Inherent Limitations
### 3.1 Qualified Memory Bound
Rather than claiming constant $O(1)$ memory, the physical memory bound for
partitioned streams accounts for Arrow buffer sharing and asynchronous
buffering:
$$\text{Peak Working Memory} \le O(|P_{\max}|) + O(B_{\text{input}}) +
O(B_{\text{pending}}) + O(\text{eval\_temp})$$
Where:
- **$O(|P_{\max}|)$:** The raw row data of the largest single partition in the
stream.
- **$O(B_{\text{input}})$ (Buffer Pinning Slack):** Arrow zero-copy slicing
(`RecordBatch::slice`) retains `Arc` references to parent memory buffers. If an
open partition spans across multiple batches, the head slice pins the first
batch (holding unneeded prefix rows) and the tail slice pins the last batch
(holding unneeded suffix rows). This unneeded pinned data can add up to
approximately $2 \times B_{\text{input}}$ of buffer slack.
- **$O(B_{\text{pending}})$:** Output batches awaiting downstream polling,
bounded by $B_{\text{input}} \times \text{output\_width}$ due to backpressure.
- **$O(\text{eval\_temp})$:** During partition finalization,
`finalize_active_partition()` simultaneously holds the input slices, the
concatenated partition batch, the newly computed window columns, and the output
batch. We estimate this creates an unverified transient memory peak of roughly
$3\times\text{ to }4\times$ the active partition's physical row data before
intermediate buffers are freed.
### 3.2 The Single-Large-Partition Limitation
Partition-at-a-time evaluation does **not** solve memory accumulation when an
individual partition spans the entire input ($|P| = N$) or exceeds available
physical RAM. In that case, `active_partition_batches` accumulates all batches,
degenerating to $O(N)$ memory with measured near-parity (observed $0.99\times$
speedup, $<0.1\%$ peak RSS difference). Mitigating single-large-partition
memory exhaustion requires out-of-core disk spilling or intra-partition sliding
window operators (`BoundedWindowAggExec`).
---
## 4. Current Prototype Limitations: MemoryPool & OOM Risk
A critical architectural limitation of the current prototype is that **it does
not integrate with DataFusion's `MemoryPool`**:
- `WindowAggStream` currently allocates active partition buffers, output
queues, and evaluation temporaries directly via the system allocator without
calling `MemoryReservation::try_grow()`.
- **Operational Risk:** DataFusion's query engine currently has no visibility
into the operator's heap usage. If an individual partition is excessively large
(or data is unpartitioned), memory will expand unthrottled, risking an
unmanaged process termination (OS OOM kill) rather than returning a clean
`DataFusionError::ResourcesExhausted` error or triggering spills.
- We propose scoping pure in-memory streaming as Phase 1, with dynamic
`MemoryReservation` integration as Phase 2.
---
## 5. Controlled Benchmark Observations (and Caveats)
To evaluate the mechanical impact of partition streaming, controlled benchmarks
were run under isolated process invocations on Linux x86_64 (`rustc 1.99.0`
release build, 5 runs per workload):
- **Query:** `SELECT c1, sn, v, COUNT(sn) OVER (PARTITION BY c1 ORDER BY sn)
FROM t;`
- **Payload Schema:** `c1: Utf8` (100-byte string partition keys), `sn: Int64`
(ordering sequence), `v: Float64`.
| Workload | Rows / Partitions | Baseline Median Query Time | Prototype Median
Query Time | Observed Speedup | Baseline Peak RSS | Prototype Peak RSS |
Physical RAM Saved |
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
| **Small Partitions (1.0M rows)** | 1.0M / 6,666 (~150 rows/part) | 119.98 ms
| 30.19 ms | **3.97×** | 321.9 MB | 173.3 MB | 148.6 MB (46.2%) |
| **Small Partitions (2.5M rows)** | 2.5M / 16,666 (~150 rows/part) | 246.19 ms
| 77.76 ms | **3.17×** | 724.0 MB | 352.7 MB | 371.2 MB (51.3%) |
| **Large Partitions (2.0M rows)** | 2.0M / 4 (500k rows/part) | 157.85 ms |
134.44 ms | **1.17×** | 574.3 MB | 349.1 MB | 225.2 MB (39.2%) |
| **Single Partition (2.0M rows)** | 2.0M / 1 (2M rows/part) | 136.59 ms |
137.86 ms | **0.99×** | 543.7 MB | 543.6 MB | +0.1 MB (0.0%) |
| **Micro-Partitions (1.5M rows)** | 1.5M / 15,000 (100 rows/part) | 147.91 ms
| 56.93 ms | **2.60×** | 454.2 MB | 227.2 MB | 227.0 MB (50.0%) |
### Essential Measurement Caveats
1. **Sample Size & Variability:** On the 1M-row small-partition workload,
baseline query timings exhibited run-to-run variability between $98.01\text{
ms}$ and $133.04\text{ ms}$ (29.2% spread, mean $117.47 \pm 12.61\text{ ms}$),
whereas prototype timings ranged from $30.05\text{ ms}$ to $31.25\text{ ms}$
(mean $30.41 \pm 0.49\text{ ms}$). Pairwise speedup across the observed
extremes ranged from $3.14\times$ to $4.43\times$. Because sample sizes are
small ($N=5$), we report observed ranges rather than parametric confidence
intervals.
2. **Causal Caveat:** Baseline variability correlated with higher minor page
faults (~70k vs. ~32k) and higher system CPU time ($0.16\text{ s}$ vs.
$0.05\text{ s}$). We hypothesize that virtual memory reallocations during
monolithic batch concatenation contributed to this jitter, though direct
kernel-level tracing (e.g. `clear_page`) was not instrumented.
3. **No Universal Speedup Guarantee:** These speedups reflect synthetic
workloads where avoiding monolithic batch reallocations saved significant CPU
cycles. They are not universal guarantees for narrower payloads or
compute-heavy expressions.
4. **Single Partition Near-Parity:** On single-partition streams (2.0M rows),
the prototype exhibited performance and memory near-parity ($0.99\times$
speedup, $543.6\text{ MB}$ vs. $543.7\text{ MB}$ peak RSS—a 148 KB difference,
well within typical OS page allocator variance).
---
## 6. Correctness Status & Test Gaps
### What Has Been Formally Tested
Automated unit tests in the prototype verify:
- Cumulative frames: `COUNT(sn)` with `ROWS BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW`.
- Forward-looking frames: `SUM(sn)` with `ROWS BETWEEN CURRENT ROW AND
UNBOUNDED FOLLOWING`.
- Ranking functions: `RANK()` and `ROW_NUMBER()` computing partition-local
sequences on distinct inputs.
- Boundary edge conditions: Partitions spanning multiple batches, multiple
partitions per batch, empty inputs, NULL partition keys across batch
boundaries, and early stream cancellation.
### Outstanding Test Gaps (Prior to PR)
- **Unverified `LEAD`/`LAG` Boundary Assertions:** While
`test_partitioned_lead_lag_rank_following_frame` executes queries with `LEAD`
and `LAG` without panicking, the test asserts only on the emitted `rank` and
`sum` columns. It lacks explicit `assert_eq!` assertions verifying that `LEAD`
and `LAG` emit correct boundary values and `NULL`s. Proving navigation
semantics remains an outstanding test gap.
- **Ties in `RANK`:** Duplicate sort values within a streaming partition have
not yet been explicitly tested.
- **Complex Frames:** `GROUPS` units, sliding offset frames (`ROWS BETWEEN 2
PRECEDING AND 2 FOLLOWING`), and `RANGE` intervals.
- **Key Types:** Multi-column composite keys (`PARTITION BY col_a, col_b`),
DictionaryArray, and StringViewArray keys.
---
## 7. Phased Implementation Proposal
To keep review manageable, we suggest splitting this work into phases:
- **Phase 1 (This RFC):** Pure in-memory partition-at-a-time streaming in
`WindowAggStream`, direct batch emission via `pending_batches`, unpartitioned
fallback, and comprehensive edge-case unit tests. Zero public API changes.
- **Phase 2:** Dynamic `MemoryReservation` integration with DataFusion's
`MemoryPool` to prevent unmanaged OOMs, plus vectorized intra-batch evaluation
for micro-partitions.
- **Phase 3 (Future Idea for Design Review):** Explore potential out-of-core
spilling mechanisms if an individual partition $|P_i|$ exceeds available memory
allocation (e.g. leveraging DataFusion's disk management infrastructure),
subject to future architectural design review.
---
## 8. Specific Feedback Requested from Maintainers
We would appreciate guidance on the following questions:
1. **Output Batch Granularity vs. Coalescing:** When partitions are small (e.g.
100 rows), emitting completed partitions directly yields a higher frequency of
small batches. Should `WindowAggExec` rely on the physical optimizer inserting
`CoalesceBatchesExec` downstream, or should `WindowAggStream` support an
internal row threshold before emitting?
2. **Planner Routing with `BoundedWindowAggExec`:** DataFusion currently uses
`BoundedWindowAggExec` for certain sliding frames. What architectural criteria
or planner heuristics should guide the selection between `BoundedWindowAggExec`
(which slides within a partition) and partition-streaming `WindowAggExec`
(which streams across partitions) when both might apply?
3. **MemoryPool Reservation Granularity (Phase 2):** For dynamic memory
tracking in Phase 2, should `MemoryReservation` be updated per incoming batch,
or dynamically per partition transition?
4. **Sorted Stream Debug Assertions:** Should `WindowAggExec` include an
optional runtime debug assertion verifying that input batches strictly adhere
to declared partition sort order?
---
## Disclaimer
This is an early architectural design proposal to align on direction with
DataFusion maintainers. The prototype code is an experimental proof-of-concept
and is not yet a production-ready or merge-ready implementation.
GitHub link: https://github.com/apache/datafusion/discussions/26425
----
This is an automatically sent email for [email protected].
To unsubscribe, please send an email to:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
