avantgardnerio opened a new pull request, #2095:
URL: https://github.com/apache/datafusion-ballista/pull/2095
## Summary
Adds `BufferExec`: a passthrough physical operator that buffers upstream
batches with a policy-configurable release trigger. Currently supports one mode:
- **`BufferMode::Dam`** — buffer batches while a `MemoryReservation` grows
against DataFusion's `MemoryPool`. Break the dam when `MemoryPool::try_grow`
refuses further allocation (any pressure — this operator, or contention from
elsewhere in the query). After the dam breaks, buffered batches drain in order
and the remainder of the input passes through unchanged. No user-tunable byte
threshold; the pool decides pressure, mirroring the mechanism `SortExec` uses
to trigger spill.
Growth path is documented on the mode enum: `PassThrough`, `BufferAll`,
`SpillOnPressure`, `BufferAndSpill`, `StageBoundary`. Only `Dam` is implemented
today; the other names are fixed so future variants slot in without renaming
callers. Ultimately this operator should land upstream in DataFusion so every
blocking op (`SortExec`, `HashAggregateExec`, joins, ...) can share the same
materialise/spill primitive instead of each rolling its own.
## Why this exists — pairing with `RuntimeStatsExec` (#2094)
Sketch-based routing has a chicken-and-egg problem: a range-repartition
operator wants to pick cut points from an upstream `RuntimeStatsExec` T-Digest,
but a streaming pipeline pulls one batch through the tap and hands it straight
to the router — the router runs with a one-batch sample.
`BufferMode::Dam` inserted between `RuntimeStatsExec` and the
sketch-consuming router closes the gap **without a plan-time size guess**:
1. Batches accumulate above the router, feeding the tap's T-Digest as they
arrive.
2. Growth continues until `MemoryPool::try_grow` refuses — the pool decides
\"enough sample\".
3. The dam breaks; the sketch now reflects representative data.
4. The router picks cuts once, then buffered batches (plus the input
remainder) drain through in order.
The pool — not a magic byte count in code — decides when the sample is large
enough. Same mechanism `SortExec` uses to decide spill timing; no new tuning
knob.
## Scope of this PR
Just the operator, its proto/serde plumbing, and tests. Nothing inserts
`BufferExec` into a plan yet — no rewrite rule, no scheduler hook. It lands
alongside the parallel-window detection rule and the range-repartition
operators (`OrderedRangeRepartitionExec` / `UnorderedRangeRepartitionExec`)
that will consume the sketch.
Included:
- Proto: `BufferExecNode { BufferMode mode }` on
`BallistaPhysicalPlanNode.oneof` (tag 6) and `enum BufferMode { BUFFER_MODE_DAM
= 0; }`.
- Codec encode/decode registration on `BallistaPhysicalExtensionCodec`.
Unknown mode discriminants surface as a decode error rather than a panic.
- No new external deps.
## Tests (4)
- `dam_drains_full_buffer_when_pool_has_headroom` — input fits within the
pool, everything buffers and drains at end-of-stream.
- `dam_breaks_on_pool_pressure_and_passes_remainder` — pool tight enough
that `try_grow` refuses partway through the stream; buffered rows plus the
remainder all pass through, none lost or duplicated.
- `dam_propagates_oom_when_first_batch_wont_fit` — pool so tight even the
first batch's reservation fails; operator surfaces OOM (mirrors `SortExec`
behavior with nothing already buffered).
- `test_buffer_exec_dam_roundtrip` — plan-node serde roundtrip through the
codec.
All three streaming tests use `GreedyMemoryPool` with a per-test byte cap so
the dam-break branch fires without hardcoded thresholds in the operator.
## Test plan
- [x] `cargo test -p ballista-core` (121 lib + 2 doc tests pass; 4 new).
- [x] `cargo check --workspace --all-targets --locked` clean (no lock churn
— no new external deps).
- [x] `cargo clippy --all-targets` clean.
- [x] `cargo fmt --all`.
## Relation to #2094
Independent PR. If #2094 (RuntimeStatsExec) lands first, this rebases to tag
7 for `BufferExecNode`; if this lands first, RuntimeStatsExec rebases the same
way. No wire compatibility concerns since neither has shipped.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
