validate whether batch read from spill exceeds max_record_batch_mem

Author: ding-young

datafusion/core/tests/sql/runtime_config.rs

@@ -193,7 +193,7 @@ async fn test_max_temp_directory_size_enforcement() {
         .unwrap();
 
     let result = ctx.sql(query).await.unwrap().collect().await;
-
+    println!("result is {result:?}");
     assert!(
         result.is_ok(),
         "Should not fail due to max temp directory size limit"

datafusion/physical-plan/benches/spill_io.rs

@@ -115,8 +115,9 @@ fn bench_spill_io(c: &mut Criterion) {
                 // - Wait for the consumer to finish processing
                 |spill_file| {
                     rt.block_on(async {
-                        let stream =
-                            spill_manager.read_spill_as_stream(spill_file).unwrap();
+                        let stream = spill_manager
+                            .read_spill_as_stream(spill_file, None)
+                            .unwrap();
                         let _ = collect(stream).await.unwrap();
                     })
                 },
@@ -519,8 +520,9 @@ fn benchmark_spill_batches_for_all_codec(
                             )
                             .unwrap()
                             .unwrap();
-                        let stream =
-                            spill_manager.read_spill_as_stream(spill_file).unwrap();
+                        let stream = spill_manager
+                            .read_spill_as_stream(spill_file, None)
+                            .unwrap();
                         let _ = collect(stream).await.unwrap();
                     })
                 },
@@ -553,7 +555,9 @@ fn benchmark_spill_batches_for_all_codec(
         let rt = Runtime::new().unwrap();
         let start = Instant::now();
         rt.block_on(async {
-            let stream = spill_manager.read_spill_as_stream(spill_file).unwrap();
+            let stream = spill_manager
+                .read_spill_as_stream(spill_file, None)
+                .unwrap();
             let _ = collect(stream).await.unwrap();
         });
         let read_time = start.elapsed();

datafusion/physical-plan/src/sorts/merge.rs

@@ -195,9 +195,7 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
 
         match futures::ready!(self.streams.poll_next(cx, idx)) {
             None => Poll::Ready(Ok(())),
-            Some(Err(e)) => {
-                Poll::Ready(Err(e))
-            }
+            Some(Err(e)) => Poll::Ready(Err(e)),
             Some(Ok((cursor, batch))) => {
                 self.cursors[idx] = Some(Cursor::new(cursor));
                 Poll::Ready(self.in_progress.push_batch(idx, batch))

datafusion/physical-plan/src/sorts/multi_level_merge.rs

@@ -237,7 +237,8 @@ impl MultiLevelMergeBuilder {
                 let spill_file = self.sorted_spill_files.remove(0);
 
                 // Not reserving any memory for this disk as we are not holding it in memory
-                self.spill_manager.read_spill_as_stream(spill_file.file)
+                self.spill_manager
+                    .read_spill_as_stream(spill_file.file, None)
             }
 
             // Only in memory streams, so merge them all in a single pass
@@ -248,15 +249,11 @@ impl MultiLevelMergeBuilder {
                     // If we have no sorted spill files left, this is the last run
                     true,
                     true,
-                    0, // TODO(ding-young)
                 )
             }
 
             // Need to merge multiple streams
-            (spill, inmem) => {
-                println!(
-                    "[merge_sorted_runs_within_mem_limit] spill{spill}, inmem{inmem}"
-                );
+            (_, _) => {
                 let mut memory_reservation = self.reservation.new_empty();
 
                 // Don't account for existing streams memory
@@ -278,20 +275,17 @@ impl MultiLevelMergeBuilder {
                         .spill_manager
                         .clone()
                         .with_batch_read_buffer_capacity(buffer_size)
-                        .read_spill_as_stream(spill.file)?;
+                        .read_spill_as_stream(
+                            spill.file,
+                            Some(spill.max_record_batch_memory),
+                        )?;
                     sorted_streams.push(stream);
                 }
-                println!(
-                    "[merge_sorted_runs_within_mem_limit] memory_reservation:{}",
-                    memory_reservation.size()
-                );
-                // TODO(ding-young) pass memory limit but what about in mem size?
                 let merge_sort_stream = self.create_new_merge_sort(
                     sorted_streams,
                     // If we have no sorted spill files left, this is the last run
                     self.sorted_spill_files.is_empty(),
                     is_only_merging_memory_streams,
-                    memory_reservation.size(),
                 )?;
 
                 // If we're only merging memory streams, we don't need to attach the memory reservation
@@ -317,7 +311,6 @@ impl MultiLevelMergeBuilder {
         streams: Vec<SendableRecordBatchStream>,
         is_output: bool,
         all_in_memory: bool,
-        expected_max_memory_usage: usize,
     ) -> Result<SendableRecordBatchStream> {
         let mut builder = StreamingMergeBuilder::new()
             .with_schema(Arc::clone(&self.schema))
@@ -338,8 +331,7 @@ impl MultiLevelMergeBuilder {
             // (reserving memory for the biggest batch in each stream)
             // TODO - avoid this hack as this can be broken easily when `SortPreservingMergeStream`
             //        changes the implementation to use more/less memory
-            println!("create new merge sort with bypass mempool");
-            builder = builder.with_bypass_mempool(expected_max_memory_usage);
+            builder = builder.with_bypass_mempool();
         } else {
             // If we are only merging in-memory streams, we need to use the memory reservation
             // because we don't know the maximum size of the batches in the streams

datafusion/physical-plan/src/sorts/streaming_merge.rs

@@ -162,10 +162,8 @@ impl<'a> StreamingMergeBuilder<'a> {
     /// Bypass the mempool and avoid using the memory reservation.
     ///
     /// This is not marked as `pub` because it is not recommended to use this method
-    pub(super) fn with_bypass_mempool(self, memory_limit: usize) -> Self {
-        // TODO(ding-young) Bypass main memory pool and use separate GreedyMemoryPool for sanity check
-        // let mem_pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default());
-        let mem_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(memory_limit));
+    pub(super) fn with_bypass_mempool(self) -> Self {
+        let mem_pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default());
 
         self.with_reservation(
             MemoryConsumer::new("merge stream mock memory").register(&mem_pool),

datafusion/physical-plan/src/spill/mod.rs

@@ -54,8 +54,13 @@ use futures::{FutureExt as _, Stream};
 struct SpillReaderStream {
     schema: SchemaRef,
     state: SpillReaderStreamState,
+    /// how much memory the largest memory batch is taking
+    pub max_record_batch_memory: Option<usize>,
 }
 
+// Small margin allowed to accommodate slight memory accounting variation
+const MEMORY_MARGIN: usize = 4096;
+
 /// When we poll for the next batch, we will get back both the batch and the reader,
 /// so we can call `next` again.
 type NextRecordBatchResult = Result<(StreamReader<BufReader<File>>, Option<RecordBatch>)>;
@@ -76,10 +81,15 @@ enum SpillReaderStreamState {
 }
 
 impl SpillReaderStream {
-    fn new(schema: SchemaRef, spill_file: RefCountedTempFile) -> Self {
+    fn new(
+        schema: SchemaRef,
+        spill_file: RefCountedTempFile,
+        max_record_batch_memory: Option<usize>,
+    ) -> Self {
         Self {
             schema,
             state: SpillReaderStreamState::Uninitialized(spill_file),
+            max_record_batch_memory,
         }
     }
 
@@ -125,6 +135,27 @@ impl SpillReaderStream {
                     Ok((reader, batch)) => {
                         match batch {
                             Some(batch) => {
+                                if let Some(max_record_batch_memory) =
+                                    self.max_record_batch_memory
+                                {
+                                    let actual_size =
+                                        get_record_batch_memory_size(&batch);
+                                    if actual_size
+                                        > max_record_batch_memory + MEMORY_MARGIN
+                                    {
+                                        return Poll::Ready(Some(Err(
+                                            DataFusionError::ResourcesExhausted(
+                                                format!(
+                                                    "Record batch memory usage ({actual_size} bytes) exceeds the expected limit ({max_record_batch_memory} bytes)\n
+                                                    by more than the allowed tolerance ({MEMORY_MARGIN} bytes).\n
+                                                    This likely indicates a bug in memory accounting during spilling.\n
+                                                    Please report this issue",
+                                                )
+                                                .to_owned(),
+                                            ),
+                                        )));
+                                    }
+                                }
                                 self.state = SpillReaderStreamState::Waiting(reader);
 
                                 Poll::Ready(Some(Ok(batch)))
@@ -417,7 +448,7 @@ mod tests {
         let spilled_rows = spill_manager.metrics.spilled_rows.value();
         assert_eq!(spilled_rows, num_rows);
 
-        let stream = spill_manager.read_spill_as_stream(spill_file)?;
+        let stream = spill_manager.read_spill_as_stream(spill_file, None)?;
         assert_eq!(stream.schema(), schema);
 
         let batches = collect(stream).await?;
@@ -481,7 +512,7 @@ mod tests {
         let spilled_rows = spill_manager.metrics.spilled_rows.value();
         assert_eq!(spilled_rows, num_rows);
 
-        let stream = spill_manager.read_spill_as_stream(spill_file)?;
+        let stream = spill_manager.read_spill_as_stream(spill_file, None)?;
         assert_eq!(stream.schema(), dict_schema);
         let batches = collect(stream).await?;
         assert_eq!(batches.len(), 2);
@@ -512,7 +543,7 @@ mod tests {
         assert!(spill_file.path().exists());
         assert!(max_batch_mem > 0);
 
-        let stream = spill_manager.read_spill_as_stream(spill_file)?;
+        let stream = spill_manager.read_spill_as_stream(spill_file, None)?;
         assert_eq!(stream.schema(), schema);
 
         let batches = collect(stream).await?;
@@ -547,7 +578,7 @@ mod tests {
         let spilled_rows = spill_manager.metrics.spilled_rows.value();
         assert_eq!(spilled_rows, num_rows);
 
-        let stream = spill_manager.read_spill_as_stream(spill_file)?;
+        let stream = spill_manager.read_spill_as_stream(spill_file, None)?;
         assert_eq!(stream.schema(), schema);
 
         let batches = collect(stream).await?;
@@ -931,8 +962,10 @@ mod tests {
                     .spill_record_batch_and_finish(&batches, "Test2")?
                     .unwrap();
 
-                let mut stream_1 = spill_manager.read_spill_as_stream(spill_file_1)?;
-                let mut stream_2 = spill_manager.read_spill_as_stream(spill_file_2)?;
+                let mut stream_1 =
+                    spill_manager.read_spill_as_stream(spill_file_1, None)?;
+                let mut stream_2 =
+                    spill_manager.read_spill_as_stream(spill_file_2, None)?;
                 stream_1.next().await;
                 stream_2.next().await;
 

datafusion/physical-plan/src/spill/spill_manager.rs

@@ -174,10 +174,12 @@ impl SpillManager {
     pub fn read_spill_as_stream(
         &self,
         spill_file_path: RefCountedTempFile,
+        max_record_batch_memory: Option<usize>,
     ) -> Result<SendableRecordBatchStream> {
         let stream = Box::pin(cooperative(SpillReaderStream::new(
             Arc::clone(&self.schema),
             spill_file_path,
+            max_record_batch_memory,
         )));
 
         Ok(spawn_buffered(stream, self.batch_read_buffer_capacity))