comphead commented on code in PR #26167:
URL: https://github.com/apache/datafusion/pull/26167#discussion_r4265277245
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -268,51 +268,83 @@ impl OutputChannel {
/// into it. The last task to call [`Self::finalize`] is the one that
/// finalizes the coalescer and ships the residual batch.
///
+/// The rows the coalescer buffers are charged to the output partition's
+/// memory reservation until they are flushed into a completed batch, which
+/// [`OutputChannel::send`] then charges on its own.
+///
/// Cheap to [`Clone`]: both fields are [`Arc`]s.
#[derive(Clone)]
struct SharedCoalescer {
- inner: Arc<Mutex<LimitedBatchCoalescer>>,
+ inner: Arc<Mutex<BufferedCoalescer>>,
active_senders: Arc<AtomicUsize>,
}
+/// A [`LimitedBatchCoalescer`] and the number of bytes of its buffered rows
+/// currently charged to the output partition's memory reservation.
+struct BufferedCoalescer {
+ coalescer: LimitedBatchCoalescer,
+ reserved: usize,
+}
+
+impl BufferedCoalescer {
+ /// Set the charge for the buffered rows to the coalescer's current size.
+ /// Completed batches must be drained first, so their bytes are not
+ /// counted here as well as in [`OutputChannel::send`].
+ fn update_reservation(&mut self, reservation: &MemoryReservation) {
+ let buffered = self.coalescer.size();
+ if buffered > self.reserved {
+ reservation.grow(buffered - self.reserved);
+ } else {
+ reservation.shrink(self.reserved - buffered);
+ }
+ self.reserved = buffered;
+ }
Review Comment:
When `buffered == self.reserved` this takes the `else` branch and calls
`reservation.shrink(0)`. That is the common case for primitive columns, because
the first write reserves a full `batch_size` capacity and `size()` then stays
constant. `MemoryReservation::shrink` has no zero check (`free` and `resize`
do), so every push still reaches the pool: an atomic in `GreedyMemoryPool`, the
state mutex in `FairSpillPool`, and an extra `RwLock` read in
`TrackConsumersPool`. This runs per push per output partition with the
coalescer mutex held. Unmeasured.
`MemoryReservation::resize` already does this delta logic and is a no-op
when the size is unchanged. It cannot be used on the shared
`Arc<MemoryReservation>` because its size is the sum over all input tasks and
the receiver, which is why `reserved` exists. A dedicated consumer owned by
`BufferedCoalescer` would remove the need for it. Precedent is `{name}[Merge
{partition}]` at line 1882.
With that:
- `update_reservation` becomes
`self.reservation.resize(self.coalescer.size())`.
- The `reserved` field, the grow/shrink branch and the `reservation`
argument threaded through `coalesce`, `push_and_drain` and `finalize` go away.
- `finalize` becomes `free()`.
- The bytes are also released on drop in the hang-up path.
- The pool sees these bytes as unspillable, which is accurate, but it lowers
the share of every spiller in `FairSpillPool`.
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -4531,6 +4563,88 @@ mod tests {
Ok(())
}
+ /// Rows held in the output coalescer are charged to the memory pool until
+ /// they are flushed, and each byte is charged only once.
+ #[tokio::test]
+ async fn repartition_reserves_coalescer_buffered_rows() -> Result<()> {
Review Comment:
This covers one input, one output and a `UInt32` column. Not covered:
- Several input partitions pushing into one shared coalescer, which is the
case the mutex comment argues about.
- Hash partitioning into several outputs with a `Utf8View` column, which is
the scenario of #26166. Here `size()` includes the copied view buffers and
changes on most pushes.
- The hang-up path (`LIMIT` dropping the stream), asserting the pool returns
to 0 once the stream and tasks are dropped.
- A tight `GreedyMemoryPool` or `FairSpillPool` run that still returns all
rows by spilling, now that `grow` can exceed the limit.
The first two can be added by parameterizing this test over inputs and
column type. `repartition_with_partial_spilling` already has the tight-pool
setup for the last one.
--
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]
