GitHub user Subodh-17 added a comment to the discussion: [RFC/Discussion]
Partition-at-a-Time Streaming Evaluation in WindowAggExec
## Update: Prototype Milestone — Partition-at-a-Time Streaming & Memory
Accounting Implemented
### Context & Motivation
In Apache DataFusion, `WindowAggExec` currently buffers all incoming
`RecordBatch`es across the entire query input before evaluating window
functions or emitting a single row. This imposes an $O(N)$ memory footprint and
blocks pipeline streaming, even when the input stream is already sorted by the
window expression's `PARTITION BY` keys. While `BoundedWindowAggExec` provides
row-level streaming for sliding frames (`ROWS BETWEEN ... PRECEDING AND CURRENT
ROW`), queries with non-sliding frames or general window functions (e.g.,
`RANK()`, `DENSE_RANK()`, `ROW_NUMBER()`, `LEAD()`, `LAG()`, or `OVER
(PARTITION BY ... ORDER BY ... ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED
FOLLOWING)`) currently force full-dataset buffering.
Following up on the initial RFC in discussion #26425, we have implemented and
validated a working prototype that adds **partition-at-a-time streaming** with
integrated **`MemoryReservation` accounting** to `WindowAggExec`.
When the input stream's physical plan properties declare an output ordering
that satisfies the window `PARTITION BY` keys, the operator can activate its
partition-streaming path. On this path, the operator accumulates incoming rows
only for the active partition, evaluates window expressions, and queues
completed partitions for progressive pipelined emission downstream. This bounds
active partition accumulation to $O(|P_{\max}|)$ rows (where $|P_{\max}|$ is
the row count of the largest partition) rather than buffering the entire $O(N)$
input dataset before emitting. Logically, completed output batches awaiting
emission in `pending_batches` add to retained operator memory across poll
cycles, while physical resident memory remains subject to Arrow zero-copy
buffer pinning (as detailed below). (If input ordering guarantees are absent or
the query is unpartitioned, the operator falls back to bulk whole-dataset
buffering).
---
### Implemented Prototype Behavior
1. **Partition-at-a-Time Execution:**
- Detects partition boundaries on incoming batches using DataFusion's
evaluate_partition_ranges helper, which uses lexicographical partition-range
detection.
- Emits completed interior partitions immediately via a FIFO queue
(`pending_batches`).
- Slices trailing partitions at batch boundaries, retaining only the
trailing slice in `active_partition_batches` to accumulate with subsequent
incoming batches.
- Automatically falls back to full-dataset bulk aggregation when the input
is unpartitioned or lacks the necessary partition-key ordering.
- Fuses stream error handling so any evaluation failure cleanly halts the
stream, frees memory reservations, and returns `Poll::Ready(None)` on
subsequent polls.
2. **Logical Memory Accounting Model:**
- Sliced partition ranges from an incoming batch $B$ are allocated
proportional shares of $B$'s physical allocation based on row count
($\sum_{i=0}^{M-1} r\_mem_i \equiv \text{batch\_size}$).
- Within this operator, this avoids significant $M \times$ logical
reservation overcounting that occurs if `get_record_batch_memory_size()` is
called naively on zero-copy Arrow slices (downstream operators' independent
measurements are noted below).
- Batches waiting in `pending_batches` remain reserved in the operator's
`MemoryReservation` across async poll cycles.
- When partitions are emitted downstream, the reservation shrinks by the
exact byte amount attributed to that partition.
- Invariant: Every byte allocated via `reservation.try_grow()` has an exact,
byte-for-byte matching `shrink()`. On stream exhaustion, error, or early drop,
`reservation.size()` returns to exactly `0`.
---
### Proposed Future Work (Not Implemented in this Prototype)
To maintain a clean boundary between current work and future enhancements:
1. **Physical Planner Optimizer Rules:** No planner rule integration (e.g. in
`EnforceDistribution` or physical optimizer passes) is included in this
prototype. Execution modes are constructed and evaluated at the operator level.
2. **Out-of-Core Spilling (Phase 3):** Disk spilling is not implemented. When
an individual partition $|P_i|$ exceeds the available pool capacity, the
operator halts safely with `ResourcesExhausted` rather than spilling to disk.
3. **Slice Compaction / Copying:** Slices are zero-copy; optional deep-copy
compaction of trailing slices to free parent buffers is deferred for maintainer
input.
4. **Micro-Partition Vectorization:** Vectorized batching across
micro-partitions before emission is left for downstream coalescing.
---
### Validation & Test Coverage (Command-Verified Output)
All checks and tests pass cleanly under WSL Linux x86_64 (`rustc 1.99.0`):
- **Unit Tests:** `cargo test -p datafusion-physical-plan --lib --
windows::window_agg_exec::tests` passes **28/28 tests** (0.17s).
- **Physical Plan Window Suite:** `cargo test -p datafusion-physical-plan --lib
-- windows` passes **46/46 tests** (4.43s).
- **Linter:** `cargo clippy -p datafusion-physical-plan --lib -- -D warnings`
completes with **0 warnings**.
- **Formatting & Hygiene:** `cargo fmt --check` and `git diff --check` complete
with exit code **0**.
Key scenarios validated by automated tests:
- Multi-partition single-batch slicing ($M = 3$) with verified non-zero
intermediate reservations and clean return to 0.
- Partitions spanning across multiple incoming batches.
- Boundary detection at exact batch edges, interspersed empty batches, and rank
ties.
- Graceful stream termination and error fusing with zero memory leaks upon
memory-pool exhaustion (`ResourcesExhausted`).
- Safe RAII reservation cleanup on early stream drops and upstream source
errors.
---
### Disclosed Architectural Limitations
We explicitly distinguish **logical `MemoryPool` reservation accounting** from
**physical-memory guarantees**:
1. **Logical Accounting $\neq$ Hard Physical Bound:** Proportional row
attribution guarantees exact arithmetic balancing within DataFusion's
`MemoryPool`, but does not bound actual resident process RAM.
2. **Zero-Copy Arrow Buffer Pinning:** Zero-copy slicing (`batch.slice()`)
creates slice array headers pointing to shared `Arc<Buffer>` allocations. If an
incoming 64 MB batch contains a trailing 1-row partition, that 1-row slice
(logically accounting for ~6.7 KB) pins the underlying 64 MB physical buffer in
RAM until that partition boundary is finalized.
3. **Transient Concatenation Spikes:** When a partition spans multiple batches,
`concat_batches()` allocates a contiguous `RecordBatch` before dropping input
batches, temporarily consuming ~$2\times$ memory during the concatenation step
without a separate pre-reservation.
4. **Kernel Pre-Allocation:** Arrow window compute kernels allocate output
arrays before `reservation.try_grow(window_cols_size)` is called.
5. **Downstream Multiplied Reservation:** If downstream operators call
`get_record_batch_memory_size()` on output slices, they will independently
measure the full parent buffer capacity for each slice.
6. **No Disk Spilling:** A single partition exceeding available memory pool
budget will cleanly terminate with `ResourcesExhausted`; it cannot spill to
disk.
---
### Questions for Maintainers
We would appreciate maintainer guidance on two specific architectural questions
before preparing a formal pull request:
1. **Operator Topology: Unified `WindowAggExec` vs. Dedicated Physical
Operator?**
In the current prototype, streaming partition execution is embedded directly
within `WindowAggExec` via an internal mode branch, falling back to bulk
evaluation when input ordering is insufficient.
*Question:* Do maintainers prefer keeping partition streaming encapsulated
within `WindowAggExec` (mirroring how `SortPreservingMergeExec` handles
streaming merges), or should it be split into a distinct execution plan node
(e.g., `StreamingWindowAggExec` or `PartitionedWindowAggExec`) selected by a
physical optimizer rule?
2. **Memory-Accounting Contract for Arrow Slices and Temporaries:**
The prototype uses proportional row attribution ($\sum r\_mem_i =
\text{batch\_size}$) to provide an exact logical balance within `MemoryPool`,
while retaining zero-copy Arrow slices.
*Question:* Is proportional logical attribution acceptable to maintainers
for partition streaming, or should trailing slices that span batch boundaries
be copied/compacted to physically release parent `Arc<Buffer>` allocations at
the cost of memcpy overhead? Additionally, should temporary memory during
`concat_batches` and window-kernel execution be pre-reserved through a scratch
consumer, or is transient unreserved allocation standard practice in DataFusion
for intra-step kernel calls?
---
### Disclaimer
This is a local design proposal to align on direction with DataFusion
maintainers. The prototype demonstrates functional partition-at-a-time
evaluation and exact `MemoryPool` accounting, but does not claim production
readiness, physical memory bounds, or planner integration.
GitHub link:
https://github.com/apache/datafusion/discussions/26425#discussioncomment-19021725
----
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]
