Sample-based sort spill reservation to mitigate merge OOM

DataFusion 52.1.0 has a TOCTOU race in ExternalSorter where merge
reservations are freed and re-created empty, letting other partitions
steal the memory (apache/datafusion#20642). Until the upstream fix
lands, compute a data-aware sort_spill_reservation_bytes by sampling
actual Arrow row sizes from the input, estimating spill file count,
and reserving enough for the merge phase.

Author: coracuity

src/commands/transform.rs

@@ -5,7 +5,12 @@ use crate::{
     ColumnDictionaryConfig, ColumnEncodingConfig, DEFAULT_BLOOM_FILTER_FPP, DataFormat,
     DictionaryMode, ListOutputsFormat, ParquetCompression, ParquetEncoding, ParquetStatistics,
     ParquetWriterVersion, PartitionStrategy, SortSpec, TransformCommand, default_thread_budget,
-    io_strategies::{OutputFileInfo, output_strategy::SinkFactory, path_template::PathTemplate},
+    io_strategies::{
+        OutputFileInfo,
+        input_strategy::InputStrategy,
+        output_strategy::SinkFactory,
+        path_template::PathTemplate,
+    },
     operations::{query::QueryOperation, sort::SortOperation},
     pipeline::Pipeline,
     sinks::{
@@ -18,6 +23,7 @@ use crate::{
         arrow::ArrowDataSource, data_source::DataSource, parquet::ParquetDataSource,
         vortex::VortexDataSource,
     },
+    utils::memory::{estimate_sort_spill_reservation, sample_avg_row_bytes},
 };
 use anyhow::{Result, anyhow};
 use arrow::datatypes::SchemaRef;
@@ -178,70 +184,84 @@ pub async fn run(args: TransformCommand) -> Result<()> {
         (from_many, true)
     };
 
-    let setup_result: Result<()> = {
-        if !should_glob && input_paths.len() == 1 {
-            let input_path = &input_paths[0];
-            let detected_input_format = detect_format(input_path, input_format)?;
-
-            let source: Box<dyn DataSource> = match detected_input_format {
-                DataFormat::Arrow => Box::new(ArrowDataSource::new(input_path.clone())),
-                DataFormat::Parquet => Box::new(ParquetDataSource::new(input_path.clone())),
-                DataFormat::Vortex => Box::new(VortexDataSource::new(input_path.clone())),
-            };
-
-            pipeline = pipeline.with_input_strategy_with_single_source(source);
-            Ok(())
-        } else {
-            let mut expanded_paths = Vec::new();
-
-            for pattern in &input_paths {
-                for entry in glob(pattern)
-                    .map_err(|e| anyhow!("Error expanding glob pattern {}: {}", pattern, e))?
-                {
-                    expanded_paths.push(
-                        entry
-                            .map_err(|e| anyhow!("Error decoding file path: {}", e))?
-                            .to_string_lossy()
-                            .to_string(),
-                    );
-                }
+    // resolve input paths (glob-expand if needed), build sources, and create InputStrategy
+    let resolved_paths: Vec<String>;
+    let input_strategy = if !should_glob && input_paths.len() == 1 {
+        let input_path = &input_paths[0];
+        let source = make_source(input_path, input_format)?;
+        resolved_paths = vec![input_path.clone()];
+        InputStrategy::Single(source)
+    } else {
+        let mut expanded_paths = Vec::new();
+
+        for pattern in &input_paths {
+            for entry in glob(pattern)
+                .map_err(|e| anyhow!("Error expanding glob pattern {}: {}", pattern, e))?
+            {
+                expanded_paths.push(
+                    entry
+                        .map_err(|e| anyhow!("Error decoding file path: {}", e))?
+                        .to_string_lossy()
+                        .to_string(),
+                );
             }
+        }
 
-            expanded_paths.sort();
-            expanded_paths.dedup();
+        expanded_paths.sort();
+        expanded_paths.dedup();
 
-            if expanded_paths.is_empty() {
-                anyhow::bail!("No input files found matching patterns: {:?}", input_paths);
-            }
+        if expanded_paths.is_empty() {
+            anyhow::bail!("No input files found matching patterns: {:?}", input_paths);
+        }
 
-            let mut sources: Vec<Box<dyn DataSource>> = Vec::new();
-            let mut schema: Option<SchemaRef> = None;
-            for input_path in expanded_paths {
-                let detected_input_format = detect_format(&input_path, input_format)?;
-                let source: Box<dyn DataSource> = match detected_input_format {
-                    DataFormat::Arrow => Box::new(ArrowDataSource::new(input_path.clone())),
-                    DataFormat::Parquet => Box::new(ParquetDataSource::new(input_path.clone())),
-                    DataFormat::Vortex => Box::new(VortexDataSource::new(input_path.clone())),
-                };
-                if let Some(ref schema) = schema {
-                    let source_schema = source.schema()?;
-                    if *schema != source_schema {
-                        anyhow::bail!(
-                            "Schema mismatch for input file {} (does not match other file(s))",
-                            &input_path
-                        );
-                    }
-                } else {
-                    schema = Some(source.schema()?);
+        let mut sources: Vec<Box<dyn DataSource>> = Vec::new();
+        let mut schema: Option<SchemaRef> = None;
+        for input_path in &expanded_paths {
+            let source = make_source(input_path, input_format)?;
+            if let Some(ref schema) = schema {
+                let source_schema = source.schema()?;
+                if *schema != source_schema {
+                    anyhow::bail!(
+                        "Schema mismatch for input file {} (does not match other file(s))",
+                        input_path
+                    );
                 }
-                sources.push(source);
+            } else {
+                schema = Some(source.schema()?);
             }
-            pipeline = pipeline.with_input_strategy_with_multiple_sources(sources);
-            Ok(())
+            sources.push(source);
         }
+        resolved_paths = expanded_paths;
+        InputStrategy::Multiple(sources)
     };
 
-    setup_result?;
+    // sample rows to estimate sort spill reservation before handing strategy to pipeline
+    if has_sort {
+        let avg_row_bytes = sample_avg_row_bytes(&input_strategy, 100_000).await?;
+
+        if avg_row_bytes > 0 {
+            let total_input_bytes: u64 = resolved_paths
+                .iter()
+                .filter_map(|p| std::fs::metadata(p).ok())
+                .map(|m| m.len())
+                .sum();
+
+            let memory_limit = effective_memory_limit.unwrap_or(total_budget * 60 / 100);
+            let partitions = effective_target_partitions.unwrap_or(three_quarter_cpus);
+            let memory_per_partition = memory_limit / partitions.max(1);
+
+            let reservation = estimate_sort_spill_reservation(
+                avg_row_bytes,
+                total_input_bytes,
+                memory_per_partition,
+                8192, // DataFusion default batch size
+            );
+
+            pipeline = pipeline.with_sort_spill_reservation_bytes(Some(reservation));
+        }
+    }
+
+    pipeline = pipeline.with_input_strategy(input_strategy);
 
     let list_outputs_format = list_outputs;
 
@@ -483,6 +503,18 @@ fn to_title_case(s: &str) -> String {
         .join(" ")
 }
 
+fn make_source(
+    path: &str,
+    input_format: Option<DataFormat>,
+) -> Result<Box<dyn DataSource>> {
+    let format = detect_format(path, input_format)?;
+    Ok(match format {
+        DataFormat::Arrow => Box::new(ArrowDataSource::new(path.to_string())),
+        DataFormat::Parquet => Box::new(ParquetDataSource::new(path.to_string())),
+        DataFormat::Vortex => Box::new(VortexDataSource::new(path.to_string())),
+    })
+}
+
 fn detect_format(path: &str, explicit_format: Option<DataFormat>) -> Result<DataFormat> {
     if let Some(format) = explicit_format {
         return Ok(format);

src/pipeline/mod.rs

@@ -28,6 +28,7 @@ pub struct PipelineConfig {
     pub target_partitions: Option<usize>,
     pub spill_path: Option<Utf8PathBuf>,
     pub spill_compression: SpillCompression,
+    pub sort_spill_reservation_bytes: Option<usize>,
 }
 
 #[derive(Default)]
@@ -45,6 +46,11 @@ impl Pipeline {
         Self::default()
     }
 
+    pub fn with_input_strategy(mut self, input_strategy: InputStrategy) -> Self {
+        self.input_strategy = Some(input_strategy);
+        self
+    }
+
     pub fn with_input_strategy_with_single_source(mut self, source: Box<dyn DataSource>) -> Self {
         self.input_strategy = Some(InputStrategy::Single(source));
 
@@ -163,6 +169,14 @@ impl Pipeline {
         self
     }
 
+    pub fn with_sort_spill_reservation_bytes(
+        mut self,
+        sort_spill_reservation_bytes: Option<usize>,
+    ) -> Self {
+        self.config.sort_spill_reservation_bytes = sort_spill_reservation_bytes;
+        self
+    }
+
     pub async fn execute(&mut self) -> Result<Vec<OutputFileInfo>> {
         let mut ctx = self.build_session_context()?;
         self.execute_with_session_context(&mut ctx).await
@@ -214,6 +228,10 @@ impl Pipeline {
         cfg.options_mut().sql_parser.dialect = self.config.query_dialect.into();
         cfg.options_mut().execution.spill_compression = self.config.spill_compression.into();
 
+        if let Some(reservation) = self.config.sort_spill_reservation_bytes {
+            cfg.options_mut().execution.sort_spill_reservation_bytes = reservation;
+        }
+
         if let Some(target_partitions) = self.config.target_partitions {
             cfg = cfg.with_target_partitions(target_partitions);
         }

src/utils/memory.rs

@@ -1,11 +1,14 @@
-//! Container-aware memory detection.
+//! Container-aware memory detection and sort memory estimation.
 //!
 //! Uses cgroup limits when running in a container (Linux), falls back to system
 //! memory on macOS or when not containerized.
 
 use arrow::datatypes::{DataType, Schema};
+use futures::StreamExt;
 use sysinfo::System;
 
+use crate::io_strategies::input_strategy::InputStrategy;
+
 /// Returns total memory in bytes, respecting container cgroup limits.
 ///
 /// In containerized environments (Docker, Kubernetes), this returns the cgroup memory
@@ -165,6 +168,67 @@ fn parse_memory_stat_net_used(content: &str) -> Option<u64> {
     Some(total.saturating_sub(slab_reclaimable))
 }
 
+/// Sample actual rows from an `InputStrategy` to measure average in-memory Arrow row size.
+///
+/// Creates a throwaway `SessionContext` and streams batches until `target_rows` rows are
+/// read (or EOF). Returns the average bytes per row based on
+/// `RecordBatch::get_array_memory_size()`, which reflects the real in-memory footprint
+/// including variable-width columns like strings and lists.
+pub async fn sample_avg_row_bytes(
+    input_strategy: &InputStrategy,
+    target_rows: usize,
+) -> anyhow::Result<usize> {
+    let mut ctx = datafusion::prelude::SessionContext::new();
+    let mut stream = input_strategy.as_stream(&mut ctx).await?;
+
+    let mut total_bytes: usize = 0;
+    let mut total_rows: usize = 0;
+    while let Some(batch) = stream.next().await {
+        let batch = batch?;
+        total_bytes += batch.get_array_memory_size();
+        total_rows += batch.num_rows();
+        if total_rows >= target_rows {
+            break;
+        }
+    }
+
+    if total_rows == 0 {
+        return Ok(0);
+    }
+    Ok(total_bytes / total_rows)
+}
+
+const MIN_SORT_SPILL_RESERVATION: usize = 10 * 1024 * 1024; // 10MB (DataFusion default)
+
+/// Estimate `sort_spill_reservation_bytes` from sampled row size and input characteristics.
+///
+/// The merge phase holds one batch per spill file in memory simultaneously.
+/// We estimate spill file count from total input size vs per-partition memory budget,
+/// then multiply by `batch_size * avg_row_bytes` to get the reservation.
+///
+/// Clamped to [10MB, memory_per_partition / 2] so there's always room for the sorter
+/// to accumulate batches before spilling.
+pub fn estimate_sort_spill_reservation(
+    avg_row_bytes: usize,
+    total_input_bytes: u64,
+    memory_per_partition: usize,
+    batch_size: usize,
+) -> usize {
+    if avg_row_bytes == 0 || memory_per_partition == 0 {
+        return MIN_SORT_SPILL_RESERVATION;
+    }
+
+    let estimated_spill_files = (total_input_bytes as usize)
+        .checked_div(memory_per_partition)
+        .unwrap_or(1)
+        .max(1);
+    let merge_batch_bytes = batch_size.saturating_mul(avg_row_bytes);
+    let reservation = estimated_spill_files.saturating_mul(merge_batch_bytes);
+    let max_reservation = memory_per_partition / 2;
+
+    reservation.clamp(MIN_SORT_SPILL_RESERVATION, max_reservation)
+}
+
 /// Returns the exact byte size for fixed-width Arrow types, or `None` for variable-width types.
 #[allow(clippy::cast_sign_loss)]
 pub fn estimate_fixed_type_bytes(dt: &DataType) -> Option<usize> {
@@ -384,4 +448,56 @@ kernel 500";
         let schema = Schema::empty();
         assert_eq!(estimate_row_bytes(&schema), 0);
     }
+
+    #[test]
+    fn test_sort_spill_reservation_hits_floor() {
+        // small input that would compute below 10MB
+        let reservation = estimate_sort_spill_reservation(
+            100,            // 100 bytes/row
+            1_000_000,      // 1MB input
+            100_000_000,    // 100MB per partition
+            8192,           // batch size
+        );
+        assert_eq!(reservation, 10 * 1024 * 1024); // 10MB floor
+    }
+
+    #[test]
+    fn test_sort_spill_reservation_scales_up() {
+        // large input: 10GB input, 500MB per partition = ~20 spill files
+        // 20 * 8192 * 200 = ~32MB
+        let reservation = estimate_sort_spill_reservation(
+            200,                // 200 bytes/row
+            10_000_000_000,     // 10GB input
+            500_000_000,        // 500MB per partition
+            8192,               // batch size
+        );
+        assert_eq!(reservation, 20 * 8192 * 200);
+        assert!(reservation > 10 * 1024 * 1024); // above floor
+        assert!(reservation < 250_000_000);       // below ceiling (250MB)
+    }
+
+    #[test]
+    fn test_sort_spill_reservation_hits_ceiling() {
+        // huge input with wide rows: would exceed half of per-partition budget
+        let memory_per_partition = 100_000_000; // 100MB
+        let reservation = estimate_sort_spill_reservation(
+            2000,               // 2KB/row (wide rows)
+            100_000_000_000,    // 100GB input
+            memory_per_partition,
+            8192,
+        );
+        assert_eq!(reservation, memory_per_partition / 2); // ceiling
+    }
+
+    #[test]
+    fn test_sort_spill_reservation_zero_row_bytes() {
+        let reservation = estimate_sort_spill_reservation(0, 1_000_000, 100_000_000, 8192);
+        assert_eq!(reservation, 10 * 1024 * 1024); // falls back to floor
+    }
+
+    #[test]
+    fn test_sort_spill_reservation_zero_memory_per_partition() {
+        let reservation = estimate_sort_spill_reservation(100, 1_000_000, 0, 8192);
+        assert_eq!(reservation, 10 * 1024 * 1024); // falls back to floor
+    }
 }