diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 2deb8fde2da6..8f690edc5426 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -468,6 +468,31 @@ async fn test_stringview_external_sort() { let _ = df.collect().await.expect("Query execution failed"); } +/// This test case is for a previously detected bug: +/// When `ExternalSorter` has read all input batches +/// - It has spilled many sorted runs to disk +/// - Its in-memory buffer for batches is almost full +/// The previous implementation will try to merge the spills and in-memory batches +/// together, without spilling the in-memory batches first, causing OOM. +#[tokio::test] +async fn test_in_mem_buffer_almost_full() { + let config = SessionConfig::new() + .with_sort_spill_reservation_bytes(3000000) + .with_target_partitions(1); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(10 * 1024 * 1024))) + .build_arc() + .unwrap(); + + let ctx = SessionContext::new_with_config_rt(config, runtime); + + let query = "select * from generate_series(1,9000000) as t1(v1) order by v1;"; + let df = ctx.sql(query).await.unwrap(); + + // Check not fail + let _ = df.collect().await.unwrap(); +} + /// Run the query with the specified memory limit, /// and verifies the expected errors are returned #[derive(Clone, Debug)] diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index a6c5dbec74dc..55ba77096a7d 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -307,7 +307,7 @@ impl ExternalSorter { let size = get_reserved_byte_for_record_batch(&input); if self.reservation.try_grow(size).is_err() { - self.sort_or_spill_in_mem_batches().await?; + self.sort_or_spill_in_mem_batches(false).await?; // We've already freed more than half of reserved memory, // so we can grow the reservation again. There's nothing we can do // if this try_grow fails. @@ -332,7 +332,7 @@ impl ExternalSorter { /// /// 2. A combined streaming merge incorporating both in-memory /// batches and data from spill files on disk. - fn sort(&mut self) -> Result { + async fn sort(&mut self) -> Result { // Release the memory reserved for merge back to the pool so // there is some left when `in_mem_sort_stream` requests an // allocation. @@ -340,10 +340,12 @@ impl ExternalSorter { if self.spilled_before() { let mut streams = vec![]; + + // Sort `in_mem_batches` and spill it first. If there are many + // `in_mem_batches` and the memory limit is almost reached, merging + // them with the spilled files at the same time might cause OOM. if !self.in_mem_batches.is_empty() { - let in_mem_stream = - self.in_mem_sort_stream(self.metrics.baseline.intermediate())?; - streams.push(in_mem_stream); + self.sort_or_spill_in_mem_batches(true).await?; } for spill in self.spills.drain(..) { @@ -488,11 +490,17 @@ impl ExternalSorter { /// the memory usage has dropped by a factor of 2, then we don't have /// to spill. Otherwise, we spill to free up memory for inserting /// more batches. - /// /// The factor of 2 aims to avoid a degenerate case where the /// memory required for `fetch` is just under the memory available, - // causing repeated re-sorting of data - async fn sort_or_spill_in_mem_batches(&mut self) -> Result<()> { + /// causing repeated re-sorting of data + /// + /// # Arguments + /// + /// * `force_spill` - If true, the method will spill the in-memory batches + /// even if the memory usage has not dropped by a factor of 2. Otherwise it will + /// only spill when the memory usage has dropped by the pre-defined factor. + /// + async fn sort_or_spill_in_mem_batches(&mut self, force_spill: bool) -> Result<()> { // Release the memory reserved for merge back to the pool so // there is some left when `in_mem_sort_stream` requests an // allocation. At the end of this function, memory will be @@ -529,7 +537,7 @@ impl ExternalSorter { // Sorting may free up some memory especially when fetch is `Some`. If we have // not freed more than 50% of the memory, then we have to spill to free up more // memory for inserting more batches. - if self.reservation.size() > before / 2 { + if (self.reservation.size() > before / 2) || force_spill { // We have not freed more than 50% of the memory, so we have to spill to // free up more memory self.spill().await?; @@ -1114,7 +1122,7 @@ impl ExecutionPlan for SortExec { let batch = batch?; sorter.insert_batch(batch).await?; } - sorter.sort() + sorter.sort().await }) .try_flatten(), ))) @@ -1409,7 +1417,7 @@ mod tests { // bytes. We leave a little wiggle room for the actual numbers. assert!((3..=10).contains(&spill_count)); assert!((9000..=10000).contains(&spilled_rows)); - assert!((36000..=40000).contains(&spilled_bytes)); + assert!((38000..=42000).contains(&spilled_bytes)); let columns = result[0].columns(); @@ -1482,7 +1490,7 @@ mod tests { // `number_of_batches / (sort_spill_reservation_bytes / batch_size)` assert!((12..=18).contains(&spill_count)); assert!((15000..=20000).contains(&spilled_rows)); - assert!((700000..=900000).contains(&spilled_bytes)); + assert!((900000..=1000000).contains(&spilled_bytes)); // Verify that the result is sorted let concated_result = concat_batches(&schema, &result)?;