[RFC/Discussion] Partition-at-a-Time Streaming Evaluation in WindowAggExec

[Subodh-17]

[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.

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

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 NULLs. 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.