andygrove commented on code in PR #2802:
URL:
https://github.com/apache/datafusion-ballista/pull/2802#discussion_r4505730718
##########
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:
Done in 7d51a949c. The inline imports moved up to the test module header.
##########
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:
Done in 7d51a949c. Spill and the in-memory encode now share
`write_ipc_stream`, `partition_counters` holds `EncodedStats`, and the options
are built once in `new()`.
##########
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:
Done in 7d51a949c. The writer is created up front now, so the `Option` and
the zero-row skip are gone.
##########
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:
Done in 7d51a949c. `spill_files` is gone. `has_spill_files` checks
`partition_counters` and `spill_path` computes an `Option<PathBuf>`.
##########
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:
Done in 7d51a949c. `spill()` is removed and the tests call `spill_batches`
through a small helper.
--
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]
