NoahKusaba commented on code in PR #2802:
URL:
https://github.com/apache/datafusion-ballista/pull/2802#discussion_r4504475929
##########
ballista/core/src/execution_plans/sort_shuffle/spill.rs:
##########
@@ -22,36 +22,41 @@
//! because the per-task buffer budget was reached.
//! At finalization, the spill bytes are concatenated verbatim into the
//! consolidated output file alongside the in-memory remainder.
+//!
+//! Each spill event appends one complete Arrow IPC stream to each spilled
+//! partition's file and closes its writer, so a partition's spill file holds
+//! one stream per event and no writer outlives the event. Keeping a writer
+//! open per partition would hold a file descriptor and a compression context
+//! (about 0.3-0.7 MB each for zstd, outside the memory pool) for every output
+//! partition until the task finishes.
use crate::JobId;
use crate::error::{BallistaError, Result};
use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::arrow::ipc::reader::StreamReader;
use datafusion::arrow::ipc::writer::StreamWriter;
use datafusion::arrow::ipc::{CompressionType, writer::IpcWriteOptions};
use datafusion::arrow::record_batch::RecordBatch;
use log::debug;
use std::collections::HashMap;
-use std::fs::File;
+use std::fs::OpenOptions;
use std::io::BufWriter;
use std::path::{Path, PathBuf};
/// Manages spill files for sort-based shuffle.
///
/// When partition buffers exceed memory limits, they are spilled to disk
-/// as Arrow IPC streams. Each output partition has at most one spill file
-/// that is appended to across multiple spill calls. During finalization,
-/// the spill file bytes are concatenated directly into the consolidated
-/// output file (no decode/re-encode round-trip).
+/// as Arrow IPC streams. Each output partition has at most one spill file,
+/// and each spill call appends one complete stream to it. During
+/// finalization, the spill file bytes are concatenated directly into the
+/// consolidated output file (no decode/re-encode round-trip); readers
+/// already cross concatenated streams.
pub struct SpillManager {
/// Base directory for spill files
spill_dir: PathBuf,
/// Schema shared by all spill writers
schema: SchemaRef,
/// Spill file path per output partition: partition_id -> spill_file_path
spill_files: HashMap<usize, PathBuf>,
Review Comment:
`spill_files` and `partition_counters` now track the same thing: each gets
an entry exactly when a partition has spilled at least once. The stored path is
also just `spill_file_path(partition_id)`. And now that this runs per spill, it
rebuilds the `PathBuf` on every call (line 125) and re-inserts it (line 171),
not only on the first spill.
Could we drop `spill_files`? `has_spill_files` could check
`partition_counters.contains_key`, and `spill_path` could compute the path,
returning `PathBuf`, or `Option<PathBuf>` when the partition has spilled.
--
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]
