NoahKusaba commented on code in PR #2802:
URL:
https://github.com/apache/datafusion-ballista/pull/2802#discussion_r4504475839
##########
ballista/core/src/execution_plans/sort_shuffle/spill.rs:
##########
@@ -98,97 +103,98 @@ impl SpillManager {
spill_dir,
schema,
spill_files: HashMap::new(),
- active_writers: HashMap::new(),
compression,
total_spilled_batches: 0,
total_bytes_spilled: 0,
partition_counters: HashMap::new(),
})
}
- /// Spills a single `batch` for `partition_id` to disk. The first call for
- /// a given `partition_id` creates the spill file; subsequent calls append.
+ /// Spills `batches` for `partition_id` to disk as one complete IPC
+ /// stream appended to the partition's spill file, which the first call
+ /// for a given `partition_id` creates. The writer is finished and dropped
+ /// before returning.
///
- /// Returns the number of bytes written (estimated from the batch's array
- /// memory size).
- pub fn spill(&mut self, partition_id: usize, batch: &RecordBatch) ->
Result<u64> {
- if batch.num_rows() == 0 {
- return Ok(0);
+ /// Returns `(batches, bytes)` written, with bytes estimated from each
+ /// batch's array memory size.
+ pub fn spill_batches(
Review Comment:
This loop is now essentially `encode_buffered_partitions` (`writer.rs:121`):
write one partition's batches as a single IPC stream and count batches, rows
and bytes. The only difference is the sink: a file opened for appending here, a
`Vec<u8>` there. Could both call one shared helper?
```rust
fn write_stream<W: Write>(
w: W,
schema: &SchemaRef,
opts: &IpcWriteOptions,
batches: &mut dyn Iterator<Item = Result<RecordBatch>>,
) -> Result<(W, EncodedStats)>
```
Spilling would then be: open the file in append mode, call `write_stream`
with a `BufWriter`, flush, and add the stats to the partition's totals.
`partition_counters` could hold `EncodedStats` instead of a `(u64, u64, u64)`
tuple, since it stores the same three counts.
Also, `SpillManager` could build its `IpcWriteOptions` once in `new()`
instead of calling `try_with_compression` on every spill (line 144).
##########
ballista/core/src/execution_plans/sort_shuffle/spill.rs:
##########
@@ -98,97 +103,98 @@ impl SpillManager {
spill_dir,
schema,
spill_files: HashMap::new(),
- active_writers: HashMap::new(),
compression,
total_spilled_batches: 0,
total_bytes_spilled: 0,
partition_counters: HashMap::new(),
})
}
- /// Spills a single `batch` for `partition_id` to disk. The first call for
- /// a given `partition_id` creates the spill file; subsequent calls append.
+ /// Spills `batches` for `partition_id` to disk as one complete IPC
+ /// stream appended to the partition's spill file, which the first call
+ /// for a given `partition_id` creates. The writer is finished and dropped
+ /// before returning.
///
- /// Returns the number of bytes written (estimated from the batch's array
- /// memory size).
- pub fn spill(&mut self, partition_id: usize, batch: &RecordBatch) ->
Result<u64> {
- if batch.num_rows() == 0 {
- return Ok(0);
+ /// Returns `(batches, bytes)` written, with bytes estimated from each
+ /// batch's array memory size.
+ pub fn spill_batches(
+ &mut self,
+ partition_id: usize,
+ batches: &mut dyn Iterator<Item =
datafusion::error::Result<RecordBatch>>,
+ ) -> Result<(u64, u64)> {
+ let spill_path = self.spill_file_path(partition_id);
+ let mut writer: Option<StreamWriter<BufWriter<std::fs::File>>> = None;
Review Comment:
nit - I don't think the lazy open or the `num_rows() == 0` skip can be
reached from `spill_all_partitions`. `BufferedBatches::drain` only calls `f`
for partitions with at least one row, and `PartitionedBatchIterator` never
returns an empty batch. The only caller that could pass an empty batch is the
`spill()` wrapper.
If that wrapper returns early on an empty batch, the writer can be created
before the loop, and the `Option` and `match` go away. Both functions get
shorter.
##########
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,
line 171 rebuilds the `PathBuf` and re-inserts it on every spill, not only the
first.
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.
##########
ballista/core/src/execution_plans/sort_shuffle/spill.rs:
##########
@@ -273,6 +278,30 @@ mod tests {
RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap()
}
+ /// Reads every batch of a partition's spill file, across the one stream
+ /// each spill event appended.
+ fn read_spill(
+ manager: &SpillManager,
+ partition_id: usize,
+ ) -> Result<Vec<RecordBatch>> {
+ use datafusion::arrow::ipc::reader::StreamReader;
+ use std::io::{Seek, SeekFrom};
Review Comment:
nit - I've seen in other repo's the precedence of keeping imports at the
file header and not in-lining them in functions.
##########
ballista/core/src/execution_plans/sort_shuffle/spill.rs:
##########
@@ -98,97 +103,98 @@ impl SpillManager {
spill_dir,
schema,
spill_files: HashMap::new(),
- active_writers: HashMap::new(),
compression,
total_spilled_batches: 0,
total_bytes_spilled: 0,
partition_counters: HashMap::new(),
})
}
- /// Spills a single `batch` for `partition_id` to disk. The first call for
- /// a given `partition_id` creates the spill file; subsequent calls append.
+ /// Spills `batches` for `partition_id` to disk as one complete IPC
+ /// stream appended to the partition's spill file, which the first call
+ /// for a given `partition_id` creates. The writer is finished and dropped
+ /// before returning.
///
- /// Returns the number of bytes written (estimated from the batch's array
- /// memory size).
- pub fn spill(&mut self, partition_id: usize, batch: &RecordBatch) ->
Result<u64> {
- if batch.num_rows() == 0 {
- return Ok(0);
+ /// Returns `(batches, bytes)` written, with bytes estimated from each
+ /// batch's array memory size.
+ pub fn spill_batches(
+ &mut self,
+ partition_id: usize,
+ batches: &mut dyn Iterator<Item =
datafusion::error::Result<RecordBatch>>,
+ ) -> Result<(u64, u64)> {
+ let spill_path = self.spill_file_path(partition_id);
+ let mut writer: Option<StreamWriter<BufWriter<std::fs::File>>> = None;
+ let mut event_batches: u64 = 0;
+ let mut event_bytes: u64 = 0;
+ for batch in batches {
+ let batch = batch?;
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let writer = match &mut writer {
+ Some(writer) => writer,
+ None => {
+ debug!("Spilling partition {} to {:?}", partition_id,
spill_path);
+ let file = OpenOptions::new()
+ .create(true)
+ .append(true)
+ .open(&spill_path)
+ .map_err(BallistaError::IoError)?;
+ let options = IpcWriteOptions::default()
+ .try_with_compression(self.compression)?;
+ writer.insert(StreamWriter::try_new_with_options(
+ BufWriter::new(file),
+ &self.schema,
+ options,
+ )?)
+ }
+ };
+ let bytes = batch.get_array_memory_size() as u64;
+ writer.write(&batch)?;
+
+ let entry = self
+ .partition_counters
+ .entry(partition_id)
+ .or_insert((0, 0, 0));
+ entry.0 += 1;
+ entry.1 += batch.num_rows() as u64;
+ entry.2 += bytes;
+ event_batches += 1;
+ event_bytes += bytes;
}
-
- if !self.active_writers.contains_key(&partition_id) {
- let spill_path = self.spill_file_path(partition_id);
- debug!(
- "Creating spill file for partition {} at {:?}",
- partition_id, spill_path
- );
- let file =
File::create(&spill_path).map_err(BallistaError::IoError)?;
- let buffered = BufWriter::new(file);
-
- let options =
-
IpcWriteOptions::default().try_with_compression(self.compression)?;
-
- let writer =
- StreamWriter::try_new_with_options(buffered, &self.schema,
options)?;
-
- self.active_writers.insert(partition_id, writer);
+ if let Some(mut writer) = writer {
+ writer.finish()?;
+ writer
+ .into_inner()?
+ .into_inner()
+ .map_err(|e| BallistaError::IoError(e.into_error()))?;
self.spill_files.insert(partition_id, spill_path);
}
- let writer = self.active_writers.get_mut(&partition_id).unwrap();
- let bytes_written = batch.get_array_memory_size() as u64;
- writer.write(batch)?;
-
- let entry = self
- .partition_counters
- .entry(partition_id)
- .or_insert((0, 0, 0));
- entry.0 += 1;
- entry.1 += batch.num_rows() as u64;
- entry.2 += bytes_written;
+ self.total_spilled_batches += event_batches;
+ self.total_bytes_spilled += event_bytes;
- self.total_spilled_batches += 1;
- self.total_bytes_spilled += bytes_written;
+ Ok((event_batches, event_bytes))
+ }
- Ok(bytes_written)
+ /// Spills a single `batch` for `partition_id` to disk; see
+ /// [`Self::spill_batches`].
+ ///
+ /// Returns the number of bytes written (estimated from the batch's array
+ /// memory size).
+ pub fn spill(&mut self, partition_id: usize, batch: &RecordBatch) ->
Result<u64> {
Review Comment:
nit - Only the tests call `spill` now. Since this PR already has
`api-change`, how about removing it and having the tests call
`spill_batches(id, &mut std::iter::once(Ok(batch)))`? The alternative is
`#[cfg(test)]`. Either way, the public API doesn't keep a single-batch wrapper
that clones its input on every call.
--
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]
