diff --git a/benchmarks/src/clickbench.rs b/benchmarks/src/clickbench.rs index b6d118129a485..4d6b5810a14f7 100644 --- a/benchmarks/src/clickbench.rs +++ b/benchmarks/src/clickbench.rs @@ -18,16 +18,21 @@ use std::fs; use std::io::ErrorKind; use std::path::{Path, PathBuf}; +use std::sync::Arc; use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats}; use clap::Args; +use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::{ExplainFormat, ExplainOption}; +use datafusion::physical_plan::collect; use datafusion::{ error::{DataFusionError, Result}, prelude::SessionContext, }; use datafusion_common::exec_datafusion_err; use datafusion_common::instant::Instant; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; /// SQL to create the hits view with proper EventDate casting. /// @@ -62,6 +67,11 @@ pub struct RunOpt { #[arg(long = "pushdown")] pushdown: bool, + /// Maximum compressed bytes to prefetch for the next Parquet row group. + /// Zero disables prefetch. Uses the query's memory pool. + #[arg(long, default_value_t = 0)] + prefetch_bytes: usize, + /// Common options #[command(flatten)] common: CommonOpt, @@ -254,8 +264,46 @@ impl RunOpt { let mut query_results = vec![]; for i in 0..self.iterations() { let start = Instant::now(); - let results = ctx.sql(sql).await?.collect().await?; + let dataframe = ctx.sql(sql).await?; + let task_ctx = Arc::new(dataframe.task_ctx()); + let mut plan = dataframe.create_physical_plan().await?; + if self.prefetch_bytes > 0 { + plan = plan + .transform_up(|plan| { + if let Some(exec) = plan.downcast_ref::() + && let Some((config, source)) = + exec.downcast_to_file_source::() + { + let source = source.clone().with_row_group_prefetch( + self.prefetch_bytes, + Arc::clone(&ctx.runtime_env().memory_pool), + ); + let config = FileScanConfigBuilder::from(config.clone()) + .with_source(Arc::new(source)) + .build(); + return Ok(Transformed::yes(Arc::new( + exec.clone().with_data_source(Arc::new(config)), + ))); + } + Ok(Transformed::no(plan)) + })? + .data; + } + let results = collect(Arc::clone(&plan), task_ctx).await?; let elapsed = start.elapsed(); + if self.prefetch_bytes > 0 { + let mut prefetched = 0; + plan.apply(|plan| { + if let Some(value) = plan + .metrics() + .and_then(|metrics| metrics.sum_by_name("prefetch_row_groups")) + { + prefetched += value.as_usize(); + } + Ok(TreeNodeRecursion::Continue) + })?; + println!("Prefetched row groups: {prefetched}"); + } let ms = elapsed.as_secs_f64() * 1000.0; millis.push(ms); let row_count: usize = results.iter().map(|b| b.num_rows()).sum(); diff --git a/datafusion/core/tests/parquet/filter_pushdown.rs b/datafusion/core/tests/parquet/filter_pushdown.rs index eaca57d7e963c..0f13cbd655375 100644 --- a/datafusion/core/tests/parquet/filter_pushdown.rs +++ b/datafusion/core/tests/parquet/filter_pushdown.rs @@ -814,9 +814,11 @@ async fn pushed_down_predicate_reports_the_original_error() { async fn upfront_io_preserves_page_pruning_from_session_configuration() { use arrow::array::{Int64Array, StringArray, StructArray}; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion_common::tree_node::{Transformed, TreeNode}; + use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource_parquet::source::ParquetSource; + use datafusion_execution::memory_pool::GreedyMemoryPool; use parquet::arrow::ArrowWriter; use parquet::file::properties::EnabledStatistics; use std::sync::Arc; @@ -904,82 +906,104 @@ async fn upfront_io_preserves_page_pruning_from_session_configuration() { for (pushdown, progressive) in [(false, true), (true, true), (true, false)] { - let config = SessionConfig::new() - .with_target_partitions(1) - .with_batch_size(127); - let ctx = SessionContext::new_with_config(config); - for (key, value) in [ - ("pushdown_filters", pushdown), - ("progressive_io", progressive), - ] { - ctx.sql(&format!( - "SET datafusion.execution.parquet.{key} = {value}" - )) - .await - .unwrap() - .collect() + for budget in [0, 1 << 20] { + let config = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(127); + let ctx = SessionContext::new_with_config(config); + for (key, value) in [ + ("pushdown_filters", pushdown), + ("progressive_io", progressive), + ] { + ctx.sql(&format!( + "SET datafusion.execution.parquet.{key} = {value}" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + } + ctx.register_parquet( + "t", + path.to_str().unwrap(), + ParquetReadOptions::default(), + ) .await .unwrap(); - } - ctx.register_parquet( - "t", - path.to_str().unwrap(), - ParquetReadOptions::default(), - ) - .await - .unwrap(); - let key = if predicate_in_output { "key," } else { "" }; - let plan = ctx.sql(&format!("SELECT {key} payload, nested.value AS value FROM t WHERE key < 100 ORDER BY value")) - .await.unwrap().create_physical_plan().await.unwrap(); - let mut scans = 0; - plan.apply(|plan| { - if let Some(exec) = plan.downcast_ref::() - && let Some((_, source)) = - exec.downcast_to_file_source::() - { - scans += 1; - assert_eq!( - source.table_parquet_options().global.progressive_io, - progressive - ); + let key = if predicate_in_output { "key," } else { "" }; + let plan = ctx.sql(&format!("SELECT {key} payload, nested.value AS value FROM t WHERE key < 100 ORDER BY value")) + .await.unwrap().create_physical_plan().await.unwrap(); + let mut scans = 0; + let plan = plan + .transform_up(|plan| { + if let Some(exec) = plan.downcast_ref::() + && let Some((config, source)) = + exec.downcast_to_file_source::() + { + scans += 1; + assert_eq!( + source + .table_parquet_options() + .global + .progressive_io, + progressive + ); + let source = source.clone().with_row_group_prefetch( + budget, + Arc::new(GreedyMemoryPool::new(1 << 20)), + ); + let config = + FileScanConfigBuilder::from(config.clone()) + .with_source(Arc::new(source)) + .build(); + return Ok(Transformed::yes(Arc::new( + exec.clone().with_data_source(Arc::new(config)), + ))); + } + Ok(Transformed::no(plan)) + }) + .unwrap() + .data; + assert_eq!(scans, 1); + let batches = + collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let batch = concat_batches(&plan.schema(), &batches).unwrap(); + let values = batch + .column(batch.num_columns() - 1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values().as_ref(), expected.as_slice()); + if let Some(expected) = &expected_batch { + assert_eq!(&batch, expected); + } else { + expected_batch = Some(batch); } - Ok(TreeNodeRecursion::Continue) - }) - .unwrap(); - assert_eq!(scans, 1); - let batches = - collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); - let batch = concat_batches(&plan.schema(), &batches).unwrap(); - let values = batch - .column(batch.num_columns() - 1) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(values.values().as_ref(), expected.as_slice()); - if let Some(expected) = &expected_batch { - assert_eq!(&batch, expected); - } else { - expected_batch = Some(batch); - } - let metrics = TestParquetFile::parquet_metrics(&plan).unwrap(); - let bytes = get_value(&metrics, "bytes_scanned"); - if indexed && clustered { - assert!( - get_pruning_metrics(&metrics, "page_index_rows_pruned").0 > 0 - ); - } - if !predicate_in_output { - if let Some(expected) = baseline_bytes { - assert_eq!( - bytes, expected, - "clustered={clustered}, indexed={indexed}, pushdown={pushdown}, progressive={progressive}" + let metrics = TestParquetFile::parquet_metrics(&plan).unwrap(); + let bytes = get_value(&metrics, "bytes_scanned"); + if indexed && clustered { + assert!( + get_pruning_metrics(&metrics, "page_index_rows_pruned").0 + > 0 ); - } else { - baseline_bytes = Some(bytes); } - } - if pushdown { - assert!(get_value(&metrics, "pushdown_rows_pruned") > 0); + if !predicate_in_output { + if let Some(expected) = baseline_bytes { + assert_eq!( + bytes, expected, + "clustered={clustered}, indexed={indexed}, pushdown={pushdown}, progressive={progressive}, budget={budget}" + ); + } else { + baseline_bytes = Some(bytes); + } + } + if budget > 0 { + assert!(get_value(&metrics, "prefetch_bytes") > 0); + } + if pushdown { + assert!(get_value(&metrics, "pushdown_rows_pruned") > 0); + } } } } diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index a2589af19a6ee..df30b968b1027 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -63,6 +63,7 @@ criterion = { workspace = true } datafusion-functions = { workspace = true } datafusion-functions-nested = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true, features = ["rt-multi-thread"] } # Note: add additional linter rules in lib.rs. # Rust does not support workspace + new linter rules in subcrates yet diff --git a/datafusion/datasource-parquet/src/metrics.rs b/datafusion/datasource-parquet/src/metrics.rs index b82bd54839f4f..d1c05e6a1b78b 100644 --- a/datafusion/datasource-parquet/src/metrics.rs +++ b/datafusion/datasource-parquet/src/metrics.rs @@ -467,3 +467,26 @@ impl RowFilterSkippedFullyMatchedMetric { count.add(1); } } + +/// Per-stream counters registered in the scan's existing metric set. +#[derive(Clone)] +pub(crate) struct PrefetchMetrics { + pub bytes: Count, + pub row_groups: Count, + pub budget_skips: Count, + pub wait_time: Time, +} + +impl PrefetchMetrics { + pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + bytes: MetricBuilder::new(metrics).counter("prefetch_bytes", partition), + row_groups: MetricBuilder::new(metrics) + .counter("prefetch_row_groups", partition), + budget_skips: MetricBuilder::new(metrics) + .counter("prefetch_budget_skips", partition), + wait_time: MetricBuilder::new(metrics) + .subset_time("prefetch_wait_time", partition), + } + } +} diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 23f9a64ed2ad7..f6736152fe796 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/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, RgPlanEntry, - RowGroupPruner, + RowGroupPrefetchOptions, RowGroupPruner, }; use crate::row_group_filter::{RowGroupAccessPlanFilter, row_group_in_range}; use crate::{ @@ -236,6 +236,7 @@ fn validate_predicate_does_not_reference_virtual_columns( /// as an explicit state machine -- see [`ParquetOpenState`] for details. #[derive(Clone)] pub(super) struct ParquetMorselizer { + pub(crate) row_group_prefetch: Option, pub(crate) progressive_io: bool, /// Execution partition index pub(crate) partition_index: usize, @@ -422,6 +423,7 @@ impl fmt::Debug for ParquetOpenState { } struct PreparedParquetOpen { + row_group_prefetch: Option, progressive_io: bool, partition_index: usize, partitioned_file: PartitionedFile, @@ -856,6 +858,7 @@ impl ParquetMorselizer { metrics: self.metrics.clone(), parquet_file_reader_factory: Arc::clone(&self.parquet_file_reader_factory), async_file_reader, + row_group_prefetch: self.row_group_prefetch.clone(), progressive_io: self.progressive_io, batch_size: self.batch_size, logical_file_schema: Arc::clone(&logical_file_schema), @@ -1526,20 +1529,21 @@ impl RowGroupsPrunedParquetOpen { decoder_limit: prepared.limit, }; - let mut fetch_selections = if !prepared.progressive_io { - access_plan - .inner() - .iter() - .map(|access| match access { - crate::RowGroupAccess::Selection(selection) => { - Some(selection.clone()) - } - _ => None, - }) - .collect::>() - } else { - vec![] - }; + let mut fetch_selections = + if !prepared.progressive_io || prepared.row_group_prefetch.is_some() { + access_plan + .inner() + .iter() + .map(|access| match access { + crate::RowGroupAccess::Selection(selection) => { + Some(selection.clone()) + } + _ => None, + }) + .collect::>() + } else { + vec![] + }; let prepared_access_plan = prepare_access_plan(access_plan)?; // #24355: a row selection (from page-index pruning, or an externally // supplied `ParquetRowSelection`) is carried by the decoder as one @@ -1700,11 +1704,18 @@ impl RowGroupsPrunedParquetOpen { decoder: Some(decoder), active_reader: None, rg_plan, - reader: prepared.async_file_reader, + reader: Arc::new(tokio::sync::Mutex::new(prepared.async_file_reader)), + row_group_prefetch: prepared.row_group_prefetch, progressive_io: prepared.progressive_io, fetch_projection, upfront_row_group: None, parquet_metadata: Arc::clone(reader_metadata.metadata()), + pending_prefetch: None, + prefetch_metrics: crate::metrics::PrefetchMetrics::new( + &prepared.metrics, + prepared.partition_index, + ), + prefetch_reservation: None, decoder_projection, arrow_reader_metrics, predicate_cache_inner_records, @@ -2337,6 +2348,7 @@ mod test { )?; Ok(ParquetMorselizer { + row_group_prefetch: None, progressive_io: true, partition_index: self.partition_index, projection, @@ -3431,19 +3443,30 @@ mod test { .with_extension(access_plan); for progressive_io in [false, true] { - for (reverse, expected) in [ - (false, vec![3, 4, 5, 6, 7, 8, 9, 10]), - (true, vec![9, 10, 5, 6, 7, 8, 3, 4]), - ] { - let mut opener = ParquetMorselizerBuilder::new() - .with_store(Arc::clone(&store)) - .with_schema(Arc::clone(&schema)) - .with_projection_indices(&[0]) - .with_reverse_row_groups(reverse) - .build(); - opener.progressive_io = progressive_io; - let stream = open_file(&opener, file.clone()).await.unwrap(); - assert_eq!(collect_int32_values(stream).await, expected); + for prefetch in [false, true] { + for (reverse, expected) in [ + (false, vec![3, 4, 5, 6, 7, 8, 9, 10]), + (true, vec![9, 10, 5, 6, 7, 8, 3, 4]), + ] { + let mut opener = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_projection_indices(&[0]) + .with_reverse_row_groups(reverse) + .build(); + opener.progressive_io = progressive_io; + opener.row_group_prefetch = + prefetch.then(|| RowGroupPrefetchOptions { + max_bytes: 1 << 20, + memory_pool: Arc::new( + datafusion_execution::memory_pool::GreedyMemoryPool::new( + 1 << 20, + ), + ), + }); + let stream = open_file(&opener, file.clone()).await.unwrap(); + assert_eq!(collect_int32_values(stream).await, expected); + } } } } diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index cd1c6925cda39..037d020bb4519 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -35,9 +35,13 @@ //! The opener constructs both halves and hands the state off to //! [`PushDecoderStreamState::into_stream`] for consumption. +use bytes::Bytes; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use std::collections::VecDeque; use std::ops::Range; use std::sync::Arc; +use tokio::sync::Mutex; use arrow::array::RecordBatch; use arrow::datatypes::SchemaRef; @@ -279,6 +283,13 @@ impl RowGroupPruner { } } +/// Execution-local prefetch configuration, supplied by the embedding engine. +#[derive(Clone, Debug)] +pub(crate) struct RowGroupPrefetchOptions { + pub(crate) max_bytes: usize, + pub(crate) memory_pool: Arc, +} + /// Preserve the access plan's page pruning when fetching output and predicate /// columns together. Without a selection or offset index, fetch complete chunks. fn column_ranges( @@ -337,6 +348,53 @@ fn merge_ranges(mut ranges: Vec>) -> Vec> { ranges } +/// At most one future row group is in flight. The task owns the reservation so +/// cancellation keeps its bytes accounted until the I/O future is actually dropped. +pub(crate) struct PrefetchedRowGroup { + row_group: usize, + ranges: Vec>, + task: SpawnedTask, MemoryReservation)>>, +} + +impl PrefetchedRowGroup { + fn start( + entry: &RgPlanEntry, + metadata: &ParquetMetaData, + projection: &ProjectionMask, + options: &RowGroupPrefetchOptions, + reader: Arc>>, + metrics: &crate::metrics::PrefetchMetrics, + ) -> Option { + let ranges = merge_ranges(column_ranges(metadata, entry, projection)?); + let bytes = ranges.iter().try_fold(0usize, |total, range| { + total.checked_add(usize::try_from(range.end - range.start).ok()?) + })?; + if bytes == 0 || bytes > options.max_bytes { + metrics.budget_skips.add(1); + return None; + } + let reservation = MemoryConsumer::new("Parquet row-group prefetch") + .register(&options.memory_pool); + if reservation.try_grow(bytes).is_err() { + metrics.budget_skips.add(1); + return None; + } + let metrics = metrics.clone(); + let fetch_ranges = ranges.clone(); + let task = SpawnedTask::spawn(async move { + let data = reader.lock().await.get_byte_ranges(fetch_ranges).await?; + metrics.bytes.add(data.iter().map(Bytes::len).sum()); + metrics.row_groups.add(1); + Ok((data, reservation)) + }); + Some(Self { + row_group: entry.rg_index, + ranges, + task, + }) + } +} + /// State for a stream that decodes a single Parquet file using a push-based decoder. /// /// The [`transition`](Self::transition) method drives the decoder in a loop: it requests @@ -347,11 +405,15 @@ pub(crate) struct PushDecoderStreamState { pub(crate) decoder: Option, pub(crate) active_reader: Option, pub(crate) rg_plan: VecDeque, - pub(crate) reader: Box, + pub(crate) reader: Arc>>, + pub(crate) row_group_prefetch: Option, pub(crate) progressive_io: bool, pub(crate) fetch_projection: ProjectionMask, pub(crate) upfront_row_group: Option, pub(crate) parquet_metadata: Arc, + pub(crate) pending_prefetch: Option, + pub(crate) prefetch_metrics: crate::metrics::PrefetchMetrics, + pub(crate) prefetch_reservation: Option, /// Per-file projection: the mask installed on every decoder and the /// per-batch transform applied by [`Self::project_batch`]. pub(crate) decoder_projection: DecoderProjection, @@ -569,6 +631,41 @@ impl PushDecoderStreamState { } } + // Apply speculation only after runtime pruning has chosen the next + // row group. A pruned group's task is aborted and its bytes discarded. + if let Some(prefetch) = self.pending_prefetch.take() { + let decoder = self.decoder.as_mut().expect("decoder present"); + match decoder.peek_next_row_group() { + Ok(Some(next)) if next == prefetch.row_group => { + let result = { + let _timer = self.prefetch_metrics.wait_time.timer(); + prefetch.task.join_unwind().await + }; + match result { + Ok(Ok((data, reservation))) => { + if let Err(e) = decoder.push_ranges(prefetch.ranges, data) + { + return Some((Err(e.into()), self)); + } + self.upfront_row_group = Some(prefetch.row_group); + self.prefetch_reservation = Some(reservation); + } + // A speculative error must not fail a scan that would + // not need those bytes. Let demand reads retry normally. + Ok(Err(e)) => debug!("Parquet prefetch failed: {e}"), + Err(e) => { + return Some(( + Err(DataFusionError::External(Box::new(e))), + self, + )); + } + } + } + Ok(_) => {} + Err(e) => return Some((Err(e.into()), self)), + } + } + // Step 3: drive the decoder. let decoder = self.decoder.as_mut().expect("decoder present"); match decoder.try_next_reader() { @@ -604,6 +701,7 @@ impl PushDecoderStreamState { } if self.upfront_row_group != Some(row_group) { decoder.clear_all_ranges(); + self.prefetch_reservation = None; let entry = self.rg_plan.front().expect("current row group"); let Some(mut upfront_ranges) = column_ranges( @@ -633,6 +731,8 @@ impl PushDecoderStreamState { } let data = self .reader + .lock() + .await .get_byte_ranges(ranges.clone()) .await .map_err(DataFusionError::from); @@ -666,8 +766,39 @@ impl PushDecoderStreamState { self.active_reader = Some(reader); // The extracted reader now owns required bytes. Release any // unused speculation (e.g. pages removed by a row filter). - if !self.progressive_io { + if !self.progressive_io || self.prefetch_reservation.is_some() { decoder.clear_all_ranges(); + self.prefetch_reservation = None; + } + if let Some(options) = &self.row_group_prefetch { + match decoder.peek_next_row_group() { + Ok(Some(next)) => { + let entry = self + .rg_plan + .iter() + .find(|entry| entry.rg_index == next) + .expect("next row group is in the access plan"); + self.pending_prefetch = PrefetchedRowGroup::start( + entry, + &self.parquet_metadata, + if entry.fully_matched { + self.decoder_projection.projection_mask() + } else { + &self.fetch_projection + }, + options, + Arc::clone(&self.reader), + &self.prefetch_metrics, + ); + if self.pending_prefetch.is_some() { + // Let the I/O task start before this worker + // continues decoding or consuming batches. + tokio::task::yield_now().await; + } + } + Ok(None) => {} + Err(e) => return Some((Err(e.into()), self)), + } } } Ok(DecodeResult::Finished) => return None, @@ -905,6 +1036,11 @@ mod tests { struct ReadControl { calls: std::sync::atomic::AtomicUsize, bytes: std::sync::atomic::AtomicUsize, + started: tokio::sync::Notify, + release: tokio::sync::Notify, + block_second: bool, + fail_second: bool, + latency: std::time::Duration, } #[derive(Debug, Clone)] @@ -935,7 +1071,19 @@ mod tests { ranges.iter().map(|r| (r.end - r.start) as usize).sum(), Ordering::SeqCst, ); - self.control.calls.fetch_add(1, Ordering::SeqCst); + let call = self.control.calls.fetch_add(1, Ordering::SeqCst); + if call == 1 { + self.control.started.notify_one(); + if self.control.block_second { + self.control.release.notified().await; + } + if self.control.fail_second { + return Err(parquet::errors::ParquetError::General( + "injected prefetch failure".into(), + )); + } + } + tokio::time::sleep(self.control.latency).await; Ok(ranges .into_iter() .map(|range| { @@ -968,7 +1116,19 @@ mod tests { } } + fn prefetch_test_stream( + budget: usize, + pool: Arc, + control: Arc, + limit: Option, + predicate: Option>, + ) -> datafusion_execution::SendableRecordBatchStream { + io_test_stream(budget, pool, control, limit, predicate, true) + } + fn io_test_stream( + budget: usize, + pool: Arc, control: Arc, limit: Option, predicate: Option>, @@ -985,6 +1145,7 @@ mod tests { let file = PartitionedFile::new("prefetch.parquet", data.len() as u64); let mut source = crate::source::ParquetSource::new(schema) .with_progressive_io(progressive) + .with_row_group_prefetch(budget, pool) // Page selections disable runtime pruning (#24355). Page-pruned // I/O is covered by the session-configuration integration test. .with_enable_page_index(false) @@ -1013,36 +1174,259 @@ mod tests { .unwrap() } + async fn assert_pool_released(pool: &Arc) { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while pool.reserved() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + #[tokio::test] - async fn upfront_reads_include_predicate_only_columns_and_skip_empty_groups() { + async fn prefetch_overlaps_decode_and_preserves_order() { + use datafusion_execution::memory_pool::GreedyMemoryPool; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl { + block_second: true, + ..Default::default() + }); + let mut stream = prefetch_test_stream( + 1 << 20, + Arc::clone(&pool), + Arc::clone(&control), + None, + None, + ); + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), 100); + // No further polling of the scan: next-RG I/O must start independently. + tokio::time::timeout( + std::time::Duration::from_secs(5), + control.started.notified(), + ) + .await + .unwrap(); + assert!(pool.reserved() > 0); + let mut batches = vec![first]; + // The current reader must keep producing while the next fetch is blocked. + for _ in 1..10 { + batches.push( + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(), + ); + } + control.release.notify_one(); + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, (0..3000).collect::>()); + assert_pool_released(&pool).await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn prefetch_starts_while_consumer_uses_worker() { + use datafusion_common::instant::Instant; + use datafusion_execution::memory_pool::GreedyMemoryPool; use std::sync::atomic::Ordering; - // A predicate-only column empties the first, middle, or last group - // without row-group statistics pruning. Compare every returned value. - for empty_group in 0..3i64 { - for progressive in [false, true] { - let predicate = Arc::new(BinaryExpr::new( - Arc::new(BinaryExpr::new( - Arc::new(Column::new("w", 1)), - Operator::Divide, - lit(1000i64), - )), - Operator::NotEq, - lit(empty_group), - )) as Arc; - - let control = Arc::new(ReadControl::default()); - let mut stream = io_test_stream( - Arc::clone(&control), + + SpawnedTask::spawn(async { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl { + block_second: true, + ..Default::default() + }); + let mut stream = prefetch_test_stream( + 1 << 20, + Arc::clone(&pool), + Arc::clone(&control), + None, + None, + ); + stream.next().await.unwrap().unwrap(); + // A consumer can keep doing CPU work after receiving a batch. + // Waiting on a Notify here would yield and hide a delayed I/O start. + let deadline = Instant::now() + std::time::Duration::from_millis(100); + while control.calls.load(Ordering::SeqCst) < 2 && Instant::now() < deadline { + std::hint::spin_loop(); + } + assert_eq!( + control.calls.load(Ordering::SeqCst), + 2, + "prefetch must start while the consumer is using a runtime worker" + ); + drop(stream); + assert_pool_released(&pool).await; + }) + .join_unwind() + .await + .unwrap(); + } + + #[tokio::test] + async fn prefetch_budget_pool_pressure_limit_and_cancellation() { + use datafusion_execution::memory_pool::GreedyMemoryPool; + use std::sync::atomic::Ordering; + for (budget, capacity, limit) in [ + (0, 1 << 20, None), + (1, 1 << 20, None), + (1 << 20, 0, None), + (1 << 20, 1 << 20, Some(100)), + ] { + let pool: Arc = Arc::new(GreedyMemoryPool::new(capacity)); + let control = Arc::new(ReadControl::default()); + let mut stream = prefetch_test_stream( + budget, + Arc::clone(&pool), + Arc::clone(&control), + limit, + None, + ); + assert_eq!(stream.next().await.unwrap().unwrap().num_rows(), 100); + tokio::task::yield_now().await; + assert_eq!(pool.reserved(), 0); + assert_eq!(control.calls.load(Ordering::SeqCst), 1); + drop(stream); + } + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl { + block_second: true, + ..Default::default() + }); + let mut stream = prefetch_test_stream( + 1 << 20, + Arc::clone(&pool), + Arc::clone(&control), + None, + None, + ); + stream.next().await.unwrap().unwrap(); + tokio::time::timeout( + std::time::Duration::from_secs(5), + control.started.notified(), + ) + .await + .unwrap(); + assert!(pool.reserved() > 0); + drop(stream); // The blocked read must be aborted without releasing its gate. + assert_pool_released(&pool).await; + } + + #[tokio::test] + async fn prefetch_failure_retries_on_demand() { + use datafusion_execution::memory_pool::GreedyMemoryPool; + use std::sync::atomic::Ordering; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl { + fail_second: true, + ..Default::default() + }); + let mut stream = prefetch_test_stream( + 1 << 20, + Arc::clone(&pool), + Arc::clone(&control), + None, + None, + ); + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch.unwrap().num_rows(); + } + assert_eq!(rows, 3000); + assert_eq!(control.calls.load(Ordering::SeqCst), 4); + assert_pool_released(&pool).await; + } + + #[tokio::test] + async fn prefetch_cancels_a_row_group_pruned_while_decoding() { + use datafusion_execution::memory_pool::GreedyMemoryPool; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl { + block_second: true, + ..Default::default() + }); + let dynamic = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("v", 0))], + gt_predicate(-1), + )); + let mut stream = prefetch_test_stream( + 1 << 20, + Arc::clone(&pool), + Arc::clone(&control), + None, + Some(Arc::clone(&dynamic) as _), + ); + stream.next().await.unwrap().unwrap(); + tokio::time::timeout( + std::time::Duration::from_secs(5), + control.started.notified(), + ) + .await + .unwrap(); + dynamic.update(gt_predicate(2500)).unwrap(); + // RG1 is now prunable. Its blocked prefetch must be cancelled, not awaited. + let rows = tokio::time::timeout(std::time::Duration::from_secs(5), async { + let mut rows = 100; + while let Some(batch) = stream.next().await { + rows += batch.unwrap().num_rows(); + } + rows + }) + .await + .unwrap(); + assert_eq!(rows, 1499); // RG0 already active; 499 rows in RG2 pass the filter. + assert_pool_released(&pool).await; + } + + #[tokio::test] + async fn prefetch_with_row_filter_matches_demand_reads() { + use datafusion_execution::memory_pool::GreedyMemoryPool; + // The modulo filter empties RG1 without statistics pruning, exercising + // the decoder advancing past a prefetched group without yielding a reader. + let modulo = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("v", 0)), + Operator::Modulo, + lit(2000i64), + )), + Operator::Lt, + lit(1000i64), + )) as Arc; + for (predicate, expected) in [ + (gt_predicate(1500), (1501..3000).collect::>()), + (modulo, (0..1000).chain(2000..3000).collect()), + ] { + for budget in [0, 1 << 20] { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let mut stream = prefetch_test_stream( + budget, + Arc::clone(&pool), + Arc::new(ReadControl::default()), None, - Some(predicate), - progressive, + Some(Arc::clone(&predicate)), ); let mut values = Vec::new(); while let Some(batch) = stream.next().await { - let batch = batch.unwrap(); - assert_eq!(batch.num_columns(), 1); values.extend_from_slice( batch + .unwrap() .column(0) .as_any() .downcast_ref::() @@ -1050,20 +1434,77 @@ mod tests { .values(), ); } - assert_eq!( - values, - (0..3000) - .filter(|v| v / 1000 != empty_group) - .collect::>() - ); - let expected_calls = if !progressive { 3 } else { 5 }; - assert_eq!(control.calls.load(Ordering::SeqCst), expected_calls); + assert_eq!(values, expected); + assert_pool_released(&pool).await; + } + } + } + + #[tokio::test] + async fn upfront_reads_include_predicate_only_columns_and_skip_empty_groups() { + use datafusion_execution::memory_pool::GreedyMemoryPool; + use std::sync::atomic::Ordering; + // A predicate-only column empties the first, middle, or last group + // without row-group statistics pruning. Compare every returned value. + for empty_group in 0..3i64 { + for progressive in [false, true] { + for budget in [0, 1 << 20] { + let predicate = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("w", 1)), + Operator::Divide, + lit(1000i64), + )), + Operator::NotEq, + lit(empty_group), + )) as Arc; + let pool: Arc = + Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl::default()); + let mut stream = io_test_stream( + budget, + Arc::clone(&pool), + Arc::clone(&control), + None, + Some(predicate), + progressive, + ); + let mut values = Vec::new(); + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + assert_eq!(batch.num_columns(), 1); + values.extend_from_slice( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + } + assert_eq!( + values, + (0..3000) + .filter(|v| v / 1000 != empty_group) + .collect::>() + ); + let expected_calls = if !progressive { + 3 + } else if budget > 0 && empty_group != 1 { + 4 + } else { + 5 + }; + assert_eq!(control.calls.load(Ordering::SeqCst), expected_calls); + assert_pool_released(&pool).await; + } } } } #[tokio::test] async fn io_policies_preserve_fully_matched_filter_suppression() { + use datafusion_execution::memory_pool::GreedyMemoryPool; use std::sync::atomic::Ordering; let (_, metadata, _) = build_three_rg_file_data(); @@ -1079,40 +1520,80 @@ mod tests { .sum::() + metadata.row_group(1).column(1).byte_range().1; for progressive in [true, false] { - let control = Arc::new(ReadControl::default()); - let predicate = Arc::new(BinaryExpr::new( - Arc::new(Column::new("w", 1)), - operator, - lit(1500i64), - )); - let mut stream = io_test_stream( - Arc::clone(&control), - None, - Some(predicate), - progressive, - ); - let mut values = Vec::new(); - while let Some(batch) = stream.next().await { - values.extend_from_slice( - batch - .unwrap() - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values(), + for budget in [0, 1 << 20] { + let pool: Arc = + Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl::default()); + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("w", 1)), + operator, + lit(1500i64), + )); + let mut stream = io_test_stream( + budget, + Arc::clone(&pool), + Arc::clone(&control), + None, + Some(predicate), + progressive, ); + let mut values = Vec::new(); + while let Some(batch) = stream.next().await { + values.extend_from_slice( + batch + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + } + assert_eq!(values, expected); + assert_eq!( + control.bytes.load(Ordering::SeqCst) as u64, + expected_bytes, + "predicate-only chunks in fully matched groups must stay unread" + ); + assert_pool_released(&pool).await; } - assert_eq!(values, expected); - assert_eq!( - control.bytes.load(Ordering::SeqCst) as u64, - expected_bytes, - "predicate-only chunks in fully matched groups must stay unread" - ); } } } + /// Controlled scheduling experiment, not a production throughput benchmark. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[ignore = "manual benchmark with simulated I/O and batch processing latency"] + async fn prefetch_latency_benchmark() { + use datafusion_common::instant::Instant; + use datafusion_execution::memory_pool::GreedyMemoryPool; + use std::time::Duration; + for budget in [0, 1 << 20] { + let mut times = Vec::new(); + for _ in 0..5 { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let control = Arc::new(ReadControl { + latency: Duration::from_millis(40), + ..Default::default() + }); + let mut stream = + prefetch_test_stream(budget, Arc::clone(&pool), control, None, None); + let start = Instant::now(); + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch.unwrap().num_rows(); + // Simulate synchronous downstream processing on this worker. + std::thread::sleep(Duration::from_millis(4)); + } + assert_eq!(rows, 3000); + times.push(start.elapsed()); + assert_pool_released(&pool).await; + } + times.sort(); + println!("prefetch budget={budget}, median={:?}", times[2]); + } + } + /// Create a fresh `(creation_errors, evaluation_errors)` counter pair /// for tests. The names mirror the two metrics /// [`RowGroupPruner::new`] consumes — predicate construction is @@ -1273,6 +1754,7 @@ mod tests { let err = PushDecoderStreamState::advance_rg_plan_to(&mut plan, 5, &mut byte_progress) .expect_err("a target absent from the plan must be an internal error"); + assert!( err.to_string().contains("diverged"), "expected a divergence internal error, got: {err}", diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index f045bcda53c4b..1add32d85c7ca 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -25,6 +25,7 @@ use crate::ParquetFileReaderFactory; use crate::opener::ParquetMorselizer; use crate::opener::build_pruning_predicates; use crate::opener::build_virtual_columns_state; +use crate::push_decoder::RowGroupPrefetchOptions; use crate::row_filter::can_expr_be_pushed_down_with_schemas; use arrow_schema::Fields; use arrow_schema::extension::ExtensionType; @@ -35,6 +36,7 @@ use datafusion_common::config::EncryptionFactoryOptions; use datafusion_datasource::as_file_source; use datafusion_datasource::file_stream::FileOpener; use datafusion_datasource::morsel::Morselizer; +use datafusion_execution::memory_pool::MemoryPool; use arrow::array::timezone::Tz; use arrow::datatypes::TimeUnit; @@ -318,6 +320,7 @@ pub struct ParquetSource { /// Sort order driving `PreparedAccessPlan::reorder_by_statistics` /// in the opener. sort_order_for_reorder: Option, + row_group_prefetch: Option, } impl ParquetSource { @@ -344,9 +347,35 @@ impl ParquetSource { encryption_factory: None, reverse_row_groups: false, sort_order_for_reorder: None, + row_group_prefetch: None, } } + /// Prefetch one upcoming row group's output and predicate pages selected at file + /// open while decoding the current group. Disabled by default; a zero budget disables it. + /// + /// `max_bytes` bounds additional compressed bytes per file stream, not the + /// current reader's memory. Prefetch is skipped if a complete fetched row + /// group does not fit or `memory_pool` cannot reserve its bytes. Required + /// reads continue normally. Use the execution's memory pool to account for + /// concurrent scans together. + /// + /// This can read extra bytes when row filtering or a later dynamic + /// predicate eliminates prefetched data. Output order is unchanged. Dropping + /// the stream cancels its background I/O. This execution-local option is not + /// serialized in physical plans; set it on the executing ParquetSource. + pub fn with_row_group_prefetch( + mut self, + max_bytes: usize, + memory_pool: Arc, + ) -> Self { + self.row_group_prefetch = (max_bytes > 0).then_some(RowGroupPrefetchOptions { + max_bytes, + memory_pool, + }); + self + } + /// Fetch pages progressively as decoding and row filtering require /// them (the default). When false, the first demand read for each row group /// fetches the output and predicate pages selected at file open together. This reduces @@ -645,6 +674,7 @@ impl FileSource for ParquetSource { Ok(Box::new(ParquetMorselizer { partition_index: partition, + row_group_prefetch: self.row_group_prefetch.clone(), progressive_io: self.table_parquet_options.global.progressive_io, projection: self.projection.clone(), batch_size: self