Repository navigation
Replies: 2 comments
Update: Prototype Milestone — Partition-at-a-Time Streaming & Memory Accounting ImplementedContext & MotivationIn Apache DataFusion, 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 When the input stream's physical plan properties declare an output ordering that satisfies the window Implemented Prototype Behavior
Proposed Future Work (Not Implemented in this Prototype)To maintain a clean boundary between current work and future enhancements:
Validation & Test Coverage (Command-Verified Output)All checks and tests pass cleanly under WSL Linux x86_64 (
Key scenarios validated by automated tests:
Disclosed Architectural LimitationsWe explicitly distinguish logical
Questions for MaintainersWe would appreciate maintainer guidance on two specific architectural questions before preparing a formal pull request:
DisclaimerThis is a local design proposal to align on direction with DataFusion maintainers. The prototype demonstrates functional partition-at-a-time evaluation and exact |
Design Consultation: Partition-by-Partition Streaming for WindowAggExecHello DataFusion maintainers and community, We have been investigating memory efficiency and time-to-first-batch latency in 1. Context & MotivationIn DataFusion, queries with bounded window frames can slide incrementally via When queries contain many distinct partitions (e.g., millions of user sessions or sensor IDs) and the input stream is ordered by those partition keys, buffering the entire dataset introduces unnecessary peak memory accumulation and delays the emission of completed partition results. We have developed an isolated prototype of partition-by-partition streaming directly within 2. Mechanism & Conservative EligibilityThe prototype introduces incremental partition boundary detection:
3. Verification & Semantic EquivalenceThe prototype is covered by 30 focused unit and regression tests in
4. Disclosed Architectural BoundariesTo ensure clear alignment with DataFusion's design principles, we want to highlight three current engine boundaries:
5. Questions for Maintainers:
We appreciate your feedback and architectural guidance! |
Uh oh!
There was an error while loading. Please reload this page.
[RFC/Discussion] Partition-at-a-Time Streaming Evaluation in WindowAggExec
datafusion-physical-plan(WindowAggExec,WindowAggStream)1. Problem Statement
In Apache DataFusion, window functions that require full partition awareness (such as
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING,LEAD,LAG, andRANK) are executed byWindowAggExec.Currently,
WindowAggStreamacts as a pipeline-breaking operator:RecordBatchintoVec<RecordBatch>until the upstream stream yieldsNone(End-of-Stream / EOS).concat_batchesacross all accumulated batches, materializing a monolithic batch ofOperational Drawbacks
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$P_k$ arrive contiguously in the stream.
WindowAggExecto be pre-sorted byPARTITION BYkeys (enforced viarequired_input_ordering()). Consequently, all rows belonging to any partitionOnce 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
WindowAggStreamto evaluate and emit completed partitions incrementally:active_partition_batches: Vec<RecordBatch>andactive_partition_key: Option<Vec<ScalarValue>>for the currently open partition.process_batch): Inspect partition keys in incoming batches usingevaluate_partition_ranges:batch.slice(start, len)without cross-batch concatenation.pending_batches: VecDeque<RecordBatch>. The stream drainspending_batchesbefore polling upstream input, ensuring upstream ingestion pauses when downstream consumption is slow.PARTITION BYclause (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 inactive_partition_batchesuntil EOF).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:
Where:
RecordBatch::slice) retainsArcreferences 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 approximatelyfinalize_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 roughly3.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, $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 (
active_partition_batchesaccumulates all batches, degenerating toBoundedWindowAggExec).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:WindowAggStreamcurrently allocates active partition buffers, output queues, and evaluation temporaries directly via the system allocator without callingMemoryReservation::try_grow().DataFusionError::ResourcesExhaustederror or triggering spills.MemoryReservationintegration 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.0release build, 5 runs per workload):SELECT c1, sn, v, COUNT(sn) OVER (PARTITION BY c1 ORDER BY sn) FROM t;c1: Utf8(100-byte string partition keys),sn: Int64(ordering sequence),v: Float64.Essential Measurement Caveats
clear_page) was not instrumented.6. Correctness Status & Test Gaps
What Has Been Formally Tested
Automated unit tests in the prototype verify:
COUNT(sn)withROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW.SUM(sn)withROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING.RANK()andROW_NUMBER()computing partition-local sequences on distinct inputs.Outstanding Test Gaps (Prior to PR)
LEAD/LAGBoundary Assertions: Whiletest_partitioned_lead_lag_rank_following_frameexecutes queries withLEADandLAGwithout panicking, the test asserts only on the emittedrankandsumcolumns. It lacks explicitassert_eq!assertions verifying thatLEADandLAGemit correct boundary values andNULLs. Proving navigation semantics remains an outstanding test gap.RANK: Duplicate sort values within a streaming partition have not yet been explicitly tested.GROUPSunits, sliding offset frames (ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING), andRANGEintervals.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:
WindowAggStream, direct batch emission viapending_batches, unpartitioned fallback, and comprehensive edge-case unit tests. Zero public API changes.MemoryReservationintegration with DataFusion'sMemoryPoolto prevent unmanaged OOMs, plus vectorized intra-batch evaluation for micro-partitions.8. Specific Feedback Requested from Maintainers
We would appreciate guidance on the following questions:
WindowAggExecrely on the physical optimizer insertingCoalesceBatchesExecdownstream, or shouldWindowAggStreamsupport an internal row threshold before emitting?BoundedWindowAggExec: DataFusion currently usesBoundedWindowAggExecfor certain sliding frames. What architectural criteria or planner heuristics should guide the selection betweenBoundedWindowAggExec(which slides within a partition) and partition-streamingWindowAggExec(which streams across partitions) when both might apply?MemoryReservationbe updated per incoming batch, or dynamically per partition transition?WindowAggExecinclude 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.
All reactions