diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 6a1a4acff8389..4eed41136746b 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -1514,11 +1514,12 @@ mod tests { }; use datafusion_physical_plan::ExecutionPlan; - for (budget, limit) in [ - (0, None), - (1, None), - (32 << 20, None), - (32 << 20, Some(123)), + for (budget, governed, limit) in [ + (0, false, None), + (1, false, None), + (32 << 20, false, None), + (32 << 20, true, None), + (32 << 20, false, Some(123)), ] { let pool: Arc = Arc::new(GreedyMemoryPool::new(512 << 20)); let (data, metadata, schema) = build_three_rg_file_data(); @@ -1544,7 +1545,7 @@ mod tests { let source = crate::source::ParquetSource::new(schema) .with_progressive_io(false) .with_row_group_prefetch(1 << 20, Arc::clone(&pool)) - .with_scan_read_ahead(4, budget, Arc::clone(&pool)) + .with_scan_read_ahead(4, budget, governed, Arc::clone(&pool)) .with_enable_page_index(false) .with_pushdown_filters(true) .with_predicate(predicate) diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index f42a6584db276..ca898b6545928 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -381,16 +381,23 @@ impl ParquetSource { /// Experimental shared file-range queue, enabled only for reorderable sibling /// streams. Prepares initial row-group bytes before workers claim each job. - /// A fixed cap bounds backlog bytes across the scan. + /// A fixed cap bounds backlog bytes; governed mode also reacts to pool headroom. /// This execution-local option is not serialized and does not enable pushdown. pub fn with_scan_read_ahead( mut self, max_jobs: usize, max_bytes: usize, + governed: bool, memory_pool: Arc, ) -> Self { self.scan_read_ahead = (max_jobs > 0 && max_bytes > 0).then(|| { - ReadAheadBudget::new(max_jobs, max_bytes, memory_pool, &self.metrics) + ReadAheadBudget::new( + max_jobs, + max_bytes, + governed, + memory_pool, + &self.metrics, + ) }); self } diff --git a/datafusion/datasource/src/file_stream/read_ahead.rs b/datafusion/datasource/src/file_stream/read_ahead.rs index 24a993b9ef03d..d6c84f64b7909 100644 --- a/datafusion/datasource/src/file_stream/read_ahead.rs +++ b/datafusion/datasource/src/file_stream/read_ahead.rs @@ -25,7 +25,9 @@ use std::task::{Context, Poll, Waker}; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result}; use datafusion_common_runtime::JoinSet; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; +use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; use datafusion_physical_plan::metrics::{ Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, }; @@ -44,6 +46,7 @@ const MIN_JOB_BYTES: usize = 16 * 1024 * 1024; pub struct ReadAheadBudget { max_jobs: usize, max_bytes: usize, + governed: bool, pool: Arc, used: Mutex, peak_bytes: Gauge, @@ -55,16 +58,19 @@ pub struct ReadAheadBudget { } impl ReadAheadBudget { - /// Construct a backlog bounded by job count and compressed bytes. + /// Construct a byte-bounded backlog. Governed mode additionally retains at + /// most one sixteenth of pool headroom, leaving room for downstream growth. pub fn new( max_jobs: usize, max_bytes: usize, + governed: bool, pool: Arc, metrics: &ExecutionPlanMetricsSet, ) -> Arc { Arc::new(Self { max_jobs, max_bytes, + governed, pool, used: Mutex::new(0), peak_bytes: MetricBuilder::new(metrics) @@ -81,6 +87,25 @@ impl ReadAheadBudget { }) } + fn limit(&self, used: usize) -> usize { + if self.governed + && let MemoryLimit::Finite(limit) = self.pool.memory_limit() + { + // ponytail: pool headroom is a coarse signal; use per-consumer pressure + // when the memory-pool API exposes it. + // Add our own reservation back when computing headroom so backlog + // does not mistake its own allocation for downstream pressure. + self.max_bytes.min( + limit + .saturating_sub(self.pool.reserved()) + .saturating_add(used) + / 16, + ) + } else { + self.max_bytes + } + } + fn reserve(self: &Arc) -> Option> { let reservation = Arc::new(ReadAheadReservation { budget: Arc::clone(self), @@ -117,7 +142,9 @@ impl ReadAheadReservation { self.budget.denied.add(1); return false; }; - if target > self.budget.max_bytes || self.reservation.try_resize(bytes).is_err() { + if target > self.budget.limit(*used) + || self.reservation.try_resize(bytes).is_err() + { self.budget.denied.add(1); return false; } @@ -261,6 +288,7 @@ mod tests { let budget = ReadAheadBudget::new( 2, 32 << 20, + false, Arc::clone(&pool), &ExecutionPlanMetricsSet::new(), ); @@ -276,6 +304,32 @@ mod tests { assert_eq!(*budget.used.lock(), 0); } + #[test] + fn budget_reacts_to_downstream_memory_and_releases_reservations() { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 << 20)); + let budget = ReadAheadBudget::new( + 32, + 512 << 20, + true, + Arc::clone(&pool), + &ExecutionPlanMetricsSet::new(), + ); + let mut reservations: Vec<_> = + (0..4).map(|_| budget.reserve().unwrap()).collect(); + assert_eq!(pool.reserved(), 64 << 20); + assert!(budget.reserve().is_none()); + let downstream = MemoryConsumer::new("aggregation").register(&pool); + downstream.try_grow(512 << 20).unwrap(); + assert!(budget.reserve().is_none()); + reservations.truncate(1); + let additional = budget.reserve().unwrap(); + assert!(budget.reserve().is_none()); + drop((additional, reservations, downstream)); + assert_eq!(pool.reserved(), 0); + assert_eq!(*budget.used.lock(), 0); + assert!(budget.reserve().is_some()); + } + #[derive(Debug)] struct GatedMorselizer(Mutex>>); @@ -332,6 +386,7 @@ mod tests { let budget = ReadAheadBudget::new( 2, 32 << 20, + false, Arc::clone(&pool), &ExecutionPlanMetricsSet::new(), );