account Parquet read-ahead in the MemoryPool

Each read-ahead stream registers a `ParquetReadAhead[partition]`
consumer. The reservation follows the bytes the decoder holds plus the
bytes in flight. Bytes the decoder asks for are always reserved (grow).
Speculative read-ahead takes only what the pool can grant (try_grow);
ranges that do not fit stay pending.

`FileSource::create_morselizer_with_context` gives the source the scan's
`TaskContext`. The default delegates to `create_morselizer`.
`ParquetSource` uses it to get the memory pool.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: adriangb

datafusion/datasource-parquet/src/opener/mod.rs

@@ -29,7 +29,7 @@ use crate::metrics::{ByteProgress, RowFilterSkippedFullyMatchedMetric};
 use crate::page_filter::PagePruningAccessPlanFilter;
 use crate::push_decoder::{
     DecoderBuilderConfig, InitialDecoderState, PushDecoderStreamState, ReadAhead,
-    RgPlanEntry, RowFilterContext, RowGroupPruner,
+    ReadAheadMemory, RgPlanEntry, RowFilterContext, RowGroupPruner,
 };
 use crate::row_group_filter::{RowGroupAccessPlanFilter, row_group_in_range};
 use crate::{
@@ -58,6 +58,7 @@ use datafusion_common::{
     ColumnStatistics, HashSet, Result, ScalarValue, Statistics, exec_err, internal_err,
 };
 use datafusion_datasource::{PartitionedFile, TableSchema};
+use datafusion_execution::memory_pool::MemoryPool;
 use datafusion_physical_expr::expressions::{Column, DynamicFilterTracking, Literal};
 use datafusion_physical_expr::simplifier::PhysicalExprSimplifier;
 use datafusion_physical_expr::utils::collect_columns;
@@ -309,6 +310,8 @@ pub(super) struct ParquetMorselizer {
     /// Per-scan virtual-column state (validation already performed). `None`
     /// when no virtual columns are requested — the common path.
     pub(crate) virtual_state: Option<Arc<VirtualColumnsState>>,
+    /// Pool that accounts read-ahead buffers. `None` uses an unbounded pool.
+    pub(crate) memory_pool: Option<Arc<dyn MemoryPool>>,
 }
 
 impl fmt::Debug for ParquetMorselizer {
@@ -463,6 +466,7 @@ struct PreparedParquetOpen {
     /// the logical-with-virtual schema. `None` when no virtual columns were
     /// requested.
     virtual_state: Option<Arc<VirtualColumnsState>>,
+    memory_pool: Option<Arc<dyn MemoryPool>>,
     reorder_predicates: bool,
     pushdown_filters: bool,
     force_filter_selections: bool,
@@ -997,6 +1001,7 @@ impl ParquetMorselizer {
             projection,
             predicate,
             virtual_state: self.virtual_state.as_ref().map(Arc::clone),
+            memory_pool: self.memory_pool.clone(),
             reorder_predicates: self.reorder_filters,
             pushdown_filters: self.pushdown_filters,
             force_filter_selections: self.force_filter_selections,
@@ -1898,8 +1903,11 @@ impl RowGroupsPrunedParquetOpen {
             .file_metrics
             .row_groups_pruned_dynamic_filter
             .clone();
-        let read_ahead =
-            read_ahead_bytes.map(|window| ReadAhead::new(window, decoder.scan_plan()));
+        let read_ahead = read_ahead_bytes.map(|window| {
+            let memory =
+                ReadAheadMemory::new(prepared.memory_pool, prepared.partition_index);
+            ReadAhead::new(window, decoder.scan_plan(), memory)
+        });
         let stream = PushDecoderStreamState {
             decoder: Some(decoder),
             active_reader: None,
@@ -2219,6 +2227,7 @@ mod test {
         max_predicate_cache_size: Option<usize>,
         max_in_list_size: usize,
         read_ahead_bytes: Option<u64>,
+        memory_pool: Option<Arc<dyn MemoryPool>>,
         reverse_row_groups: bool,
         preserve_order: bool,
     }
@@ -2443,6 +2452,7 @@ mod test {
                 max_predicate_cache_size: None,
                 max_in_list_size: MAX_IN_LIST_SIZE,
                 read_ahead_bytes: None,
+                memory_pool: None,
                 reverse_row_groups: false,
                 preserve_order: false,
             }
@@ -2531,6 +2541,18 @@ mod test {
             self
         }
 
+        /// Set the read-ahead window (`read_ahead_bytes`).
+        fn with_read_ahead_bytes(mut self, bytes: u64) -> Self {
+            self.read_ahead_bytes = Some(bytes);
+            self
+        }
+
+        /// Set the pool that accounts read-ahead buffers.
+        fn with_memory_pool(mut self, pool: Arc<dyn MemoryPool>) -> Self {
+            self.memory_pool = Some(pool);
+            self
+        }
+
         /// Set a row limit.
         fn with_limit(mut self, limit: usize) -> Self {
             self.limit = Some(limit);
@@ -2632,6 +2654,7 @@ mod test {
                 reverse_row_groups: self.reverse_row_groups,
                 sort_order_for_reorder: None,
                 virtual_state,
+                memory_pool: self.memory_pool,
             })
         }
     }
@@ -5743,4 +5766,233 @@ mod test {
             assert_eq!(rows, 5);
         }
     }
+
+    /// Read-ahead (`read_ahead_bytes`) and its accounting in the memory pool.
+    mod read_ahead {
+        use super::*;
+        use std::ops::Range;
+        use std::sync::Mutex;
+
+        use bytes::Bytes;
+
+        use arrow::array::Int64Array;
+        use arrow::compute::concat_batches;
+        use datafusion_execution::memory_pool::{
+            GreedyMemoryPool, MemoryPool, PeakRecordingPool,
+        };
+        use futures::future::BoxFuture;
+        use parquet::arrow::arrow_reader::ArrowReaderOptions;
+        use parquet::arrow::async_reader::AsyncFileReader;
+
+        const WINDOW: u64 = 1024 * 1024;
+        const ROWS: i64 = 400_000;
+
+        /// Records the bytes of every data fetch of the readers it creates.
+        #[derive(Debug)]
+        struct RecordingReaderFactory {
+            inner: DefaultParquetFileReaderFactory,
+            fetches: Arc<Mutex<Vec<u64>>>,
+        }
+
+        struct RecordingReader {
+            inner: Box<dyn AsyncFileReader + Send>,
+            fetches: Arc<Mutex<Vec<u64>>>,
+        }
+
+        impl ParquetFileReaderFactory for RecordingReaderFactory {
+            fn create_reader(
+                &self,
+                partition_index: usize,
+                partitioned_file: PartitionedFile,
+                metadata_size_hint: Option<usize>,
+                metrics: &ExecutionPlanMetricsSet,
+            ) -> Result<Box<dyn AsyncFileReader + Send>> {
+                Ok(Box::new(RecordingReader {
+                    inner: self.inner.create_reader(
+                        partition_index,
+                        partitioned_file,
+                        metadata_size_hint,
+                        metrics,
+                    )?,
+                    fetches: Arc::clone(&self.fetches),
+                }))
+            }
+        }
+
+        impl AsyncFileReader for RecordingReader {
+            fn get_bytes(
+                &mut self,
+                range: Range<u64>,
+            ) -> BoxFuture<'_, parquet::errors::Result<Bytes>> {
+                self.fetches.lock().unwrap().push(range.end - range.start);
+                self.inner.get_bytes(range)
+            }
+
+            fn get_byte_ranges(
+                &mut self,
+                ranges: Vec<Range<u64>>,
+            ) -> BoxFuture<'_, parquet::errors::Result<Vec<Bytes>>> {
+                let bytes = ranges.iter().map(|r| r.end - r.start).sum();
+                self.fetches.lock().unwrap().push(bytes);
+                self.inner.get_byte_ranges(ranges)
+            }
+
+            fn get_metadata<'a>(
+                &'a mut self,
+                options: Option<&'a ArrowReaderOptions>,
+            ) -> BoxFuture<'a, parquet::errors::Result<Arc<ParquetMetaData>>>
+            {
+                self.inner.get_metadata(options)
+            }
+        }
+
+        /// A file of several MB: `a` is `0..ROWS`, `b` does not compress.
+        /// Ten row groups of several pages each.
+        async fn write_file(store: &Arc<dyn ObjectStore>) -> (SchemaRef, usize) {
+            let schema = Arc::new(Schema::new(vec![
+                Field::new("a", DataType::Int64, false),
+                Field::new("b", DataType::Int64, false),
+            ]));
+            let a: Vec<i64> = (0..ROWS).collect();
+            let b: Vec<i64> = a
+                .iter()
+                .map(|v| v.wrapping_mul(0x9E37_79B9_7F4A_7C15_u64 as i64))
+                .collect();
+            let batch = RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![Arc::new(Int64Array::from(a)), Arc::new(Int64Array::from(b))],
+            )
+            .unwrap();
+            let props = WriterProperties::builder()
+                .set_dictionary_enabled(false)
+                .set_max_row_group_row_count(Some(ROWS as usize / 10))
+                .set_data_page_row_count_limit(4096)
+                .set_write_batch_size(4096)
+                .build();
+            let len = write_parquet_batches(
+                Arc::clone(store),
+                "read_ahead.parquet",
+                vec![batch],
+                Some(props),
+            )
+            .await;
+            (schema, len)
+        }
+
+        struct Scan {
+            batches: Vec<RecordBatch>,
+            /// Bytes of each data fetch.
+            fetches: Vec<u64>,
+            /// `pool.reserved()` after each batch.
+            reserved: Vec<usize>,
+        }
+
+        /// Scan the file, with read-ahead if `pool` is set.
+        async fn scan(
+            pool: Option<Arc<dyn MemoryPool>>,
+            predicate: Option<Arc<dyn PhysicalExpr>>,
+        ) -> Scan {
+            let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+            let (schema, len) = write_file(&store).await;
+            let fetches = Arc::new(Mutex::new(vec![]));
+            let factory = RecordingReaderFactory {
+                inner: DefaultParquetFileReaderFactory::new(Arc::clone(&store)),
+                fetches: Arc::clone(&fetches),
+            };
+            let mut builder = ParquetMorselizerBuilder::new()
+                .with_store(Arc::clone(&store))
+                .with_schema(Arc::clone(&schema))
+                .with_projection_indices(&[0, 1])
+                .with_parquet_file_reader_factory(Arc::new(factory));
+            if let Some(predicate) = predicate {
+                builder = builder
+                    .with_predicate(predicate)
+                    .with_pushdown_filters(true);
+            }
+            if let Some(pool) = &pool {
+                builder = builder
+                    .with_read_ahead_bytes(WINDOW)
+                    .with_memory_pool(Arc::clone(pool));
+            }
+            let morselizer = builder.build();
+            let file = PartitionedFile::new("read_ahead.parquet".to_string(), len as u64);
+            let mut stream = open_file(&morselizer, file).await.unwrap();
+            let mut batches = vec![];
+            let mut reserved = vec![];
+            while let Some(batch) = stream.next().await {
+                batches.push(batch.unwrap());
+                reserved.extend(pool.as_ref().map(|p| p.reserved()));
+            }
+            drop(stream);
+            drop(morselizer);
+            if let Some(pool) = &pool {
+                assert_eq!(pool.reserved(), 0, "the stream returns all it reserved");
+            }
+            let fetches = fetches.lock().unwrap().clone();
+            Scan {
+                batches,
+                fetches,
+                reserved,
+            }
+        }
+
+        fn concat(batches: &[RecordBatch]) -> RecordBatch {
+            concat_batches(&batches[0].schema(), batches).unwrap()
+        }
+
+        fn a_lt(value: i64) -> Arc<dyn PhysicalExpr> {
+            let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
+            logical2physical(&col("a").lt(lit(value)), &schema)
+        }
+
+        /// A pool far smaller than the file cannot stop the scan: the bytes
+        /// the decoder needs are always granted.
+        #[tokio::test]
+        async fn small_pool_scan_is_correct_and_releases_all() {
+            let expected = scan(None, None).await;
+            let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(64 * 1024));
+            let actual = scan(Some(pool), None).await;
+            assert_eq!(concat(&actual.batches), concat(&expected.batches));
+            assert_eq!(concat(&actual.batches).num_rows(), ROWS as usize);
+        }
+
+        /// With room in the pool, the reservation follows what read-ahead
+        /// holds and stays within the window plus one required fetch.
+        #[tokio::test]
+        async fn reservation_is_bounded_by_window() {
+            let recorder = Arc::new(PeakRecordingPool::new(Arc::new(
+                GreedyMemoryPool::new(1 << 30),
+            )));
+            let pool: Arc<dyn MemoryPool> = Arc::clone(&recorder) as _;
+            let actual = scan(Some(pool), None).await;
+            // A required fetch is at most the largest fetch.
+            let largest_fetch = *actual.fetches.iter().max().unwrap();
+            let peak = recorder.max_reserved() as u64;
+            assert!(actual.reserved.iter().any(|&r| r > 0));
+            assert!(
+                peak <= WINDOW + largest_fetch,
+                "peak {peak}, largest fetch {largest_fetch}"
+            );
+        }
+
+        /// A filter that rejects every row: without read-ahead the scan never
+        /// reads `b`. Read-ahead fetches `b` ahead (a conditional range),
+        /// unless the pool has no room for it.
+        #[tokio::test]
+        async fn conditional_ranges_are_read_ahead_within_the_pool() {
+            let bytes = |scan: &Scan| scan.fetches.iter().sum::<u64>();
+            let baseline = scan(None, Some(a_lt(-1))).await;
+            assert_eq!(
+                baseline.batches.iter().map(|b| b.num_rows()).sum::<usize>(),
+                0
+            );
+
+            let unbounded: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 30));
+            let ahead = scan(Some(unbounded), Some(a_lt(-1))).await;
+            let full: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1));
+            let required_only = scan(Some(full), Some(a_lt(-1))).await;
+            assert!(bytes(&ahead) > bytes(&baseline));
+            assert_eq!(bytes(&required_only), bytes(&baseline));
+        }
+    }
 }