jayzhan211 commented on code in PR #23522:
URL: https://github.com/apache/datafusion/pull/23522#discussion_r3596209770
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -92,69 +91,112 @@ impl SpillPoolShared {
}
}
-/// Writer for a spill pool. Provides coordinated write access with FIFO
semantics.
+/// Writer for a spill pool that can be cloned to produce additional writers.
///
-/// Created by [`channel`]. See that function for architecture diagrams and
usage examples.
-///
-/// The writer is `Clone`, allowing multiple writers to coordinate on the same
pool.
-/// All clones share the same current write file and coordinate file rotation.
-/// The writer automatically manages file rotation based on the
`max_file_size_bytes`
-/// configured in [`channel`]. When the last writer clone is dropped, it
finalizes the
-/// current file so readers can access all written data.
+/// Created by [`mspc_channel`]. See that function for architecture diagrams
and usage
+/// examples.
pub struct SpillPoolWriter {
- /// Maximum size in bytes before rotating to a new file.
- /// Typically set from configuration
`datafusion.execution.max_spill_file_size_bytes`.
- max_file_size_bytes: usize,
- /// Shared state with readers (includes current_write_file for
coordination)
- shared: Arc<Mutex<SpillPoolShared>>,
+ /// The underlying shared writer. Kept private and never cloned, so this
pool always has
+ /// exactly one writer.
+ inner: SpillPoolSink,
+}
+
+impl SpillPoolWriter {
+ /// Spills a batch to the pool, rotating files when necessary.
+ ///
+ /// See [`mpsc_channel`] for the rotation semantics.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if disk I/O fails or disk quota is exceeded.
+ pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> {
+ self.inner.push_batch(batch)
+ }
+}
+
+impl SpillPoolWriter {
+ /// Returns a new sink that can be used to spill batches to the pool.
+ ///
+ /// As an alternative to this function, it is also possible to clone the
writer. The benefit
+ /// of this method is that the output type matches the type used by
[`spsc_channel`]. This
+ /// enables cost-free abstraction for producers over SPSC and MPSC
channels.
+ pub fn new_sink(&self) -> SpillPoolSink {
+ // Increment `remaining_writer_count`. The corresponding decrement is
done in the `Drop`
+ // implementation of `SpillPoolWriter`.
+ self.inner.shared.lock().remaining_writer_count += 1;
+ SpillPoolSink {
+ max_file_size_bytes: self.inner.max_file_size_bytes,
+ shared: Arc::clone(&self.inner.shared),
+ }
+ }
}
impl Clone for SpillPoolWriter {
fn clone(&self) -> Self {
- // Increment the active writer count so that `writer_dropped` is only
- // set to true when the *last* clone is dropped.
- self.shared.lock().active_writer_count += 1;
Self {
- max_file_size_bytes: self.max_file_size_bytes,
- shared: Arc::clone(&self.shared),
+ inner: self.new_sink(),
}
}
}
-impl SpillPoolWriter {
+impl Drop for SpillPoolSink {
+ fn drop(&mut self) {
+ let mut shared = self.shared.lock();
+
+ shared.remaining_writer_count -= 1;
+ let is_last_writer = shared.remaining_writer_count == 0;
+
+ if !is_last_writer {
+ // Other writer clones are still active; do not finalize or
+ // signal EOF to readers.
+ return;
+ }
+
+ // Finalize any spill files that were not finished yet
+ if !shared.open_write_files.is_empty() {
+ let files = mem::take(&mut shared.open_write_files);
+ drop(shared);
+
+ for file in files {
+ let mut file_shared = file.lock();
+
+ // Finish the current writer if it exists
+ if let Some(mut writer) = file_shared.writer.take() {
Review Comment:
👍🏻
--
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]
