From 2d0806b7886730b095586e59974dbd675811e6eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sun, 12 Jul 2026 14:17:18 +0200 Subject: [PATCH 1/2] Prefetch Parquet row groups with a bounded budget --- datafusion/common/src/config.rs | 7 + .../common/src/file_options/parquet_writer.rs | 3 + .../datasource-parquet/src/opener/mod.rs | 8 + .../datasource-parquet/src/push_decoder.rs | 229 +++++++++++++++++- datafusion/datasource-parquet/src/source.rs | 1 + .../proto/datafusion_common.proto | 6 +- datafusion/proto-common/src/from_proto/mod.rs | 15 ++ .../proto-common/src/generated/pbjson.rs | 24 ++ .../proto-common/src/generated/prost.rs | 7 + datafusion/proto-common/src/to_proto/mod.rs | 1 + .../src/generated/datafusion_proto_common.rs | 7 + .../proto/src/logical_plan/file_formats.rs | 8 + .../test_files/information_schema.slt | 2 + docs/source/user-guide/configs.md | 1 + 14 files changed, 315 insertions(+), 4 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 454af28c14b4a..29a15593e10d9 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1133,6 +1133,13 @@ config_namespace! { /// the hint, two reads will still be performed. pub metadata_size_hint: Option, default = Some(512 * 1024) + /// (reading) If specified, the parquet reader will prefetch data for + /// subsequent row groups when the projected column chunks fit within + /// this many bytes. The required ranges for the current row group are + /// always read, even when they exceed this value. None disables data + /// prefetching. + pub prefetch_size: Option, default = None + /// (reading) If true, filter expressions are be applied during the parquet decoding operation to /// reduce the number of rows decoded. This optimization is sometimes called "late materialization". pub pushdown_filters: bool, default = false diff --git a/datafusion/common/src/file_options/parquet_writer.rs b/datafusion/common/src/file_options/parquet_writer.rs index 320bfcf33e488..80857cc457d51 100644 --- a/datafusion/common/src/file_options/parquet_writer.rs +++ b/datafusion/common/src/file_options/parquet_writer.rs @@ -235,6 +235,7 @@ impl ParquetOptions { pruning: _, skip_metadata: _, metadata_size_hint: _, + prefetch_size: _, pushdown_filters: _, reorder_filters: _, force_filter_selections: _, // not used for writer props @@ -491,6 +492,7 @@ mod tests { pruning: defaults.pruning, skip_metadata: defaults.skip_metadata, metadata_size_hint: defaults.metadata_size_hint, + prefetch_size: defaults.prefetch_size, pushdown_filters: defaults.pushdown_filters, reorder_filters: defaults.reorder_filters, force_filter_selections: defaults.force_filter_selections, @@ -610,6 +612,7 @@ mod tests { pruning: global_options_defaults.pruning, skip_metadata: global_options_defaults.skip_metadata, metadata_size_hint: global_options_defaults.metadata_size_hint, + prefetch_size: global_options_defaults.prefetch_size, pushdown_filters: global_options_defaults.pushdown_filters, reorder_filters: global_options_defaults.reorder_filters, force_filter_selections: global_options_defaults.force_filter_selections, diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 87ec341f590da..aad3c8c3e0481 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -252,6 +252,8 @@ pub(super) struct ParquetMorselizer { /// Optional hint for how large the initial request to read parquet metadata /// should be pub metadata_size_hint: Option, + /// Maximum number of bytes buffered by speculative row-group prefetching. + pub prefetch_size: Option, /// Metrics for reporting pub metrics: ExecutionPlanMetricsSet, /// Factory for instantiating parquet reader @@ -425,6 +427,7 @@ struct PreparedParquetOpen { baseline_metrics: BaselineMetrics, file_pruner: Option, metadata_size_hint: Option, + prefetch_size: Option, metrics: ExecutionPlanMetricsSet, parquet_file_reader_factory: Arc, async_file_reader: Box, @@ -828,6 +831,7 @@ impl ParquetMorselizer { baseline_metrics, file_pruner, metadata_size_hint, + prefetch_size: self.prefetch_size, metrics: self.metrics.clone(), parquet_file_reader_factory: Arc::clone(&self.parquet_file_reader_factory), async_file_reader, @@ -1482,6 +1486,9 @@ impl RowGroupsPrunedParquetOpen { active_reader: None, rg_plan, reader: prepared.async_file_reader, + parquet_metadata: Arc::clone(&file_metadata), + prefetch_size: prepared.prefetch_size, + prefetched_row_groups: std::collections::HashSet::new(), decoder_projection, arrow_reader_metrics, predicate_cache_inner_records, @@ -2014,6 +2021,7 @@ mod test { predicate: self.predicate, table_schema, metadata_size_hint: self.metadata_size_hint, + prefetch_size: None, metrics: self.metrics, parquet_file_reader_factory: self .parquet_file_reader_factory diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 31bd365a4631d..c67d087ff6b19 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -35,7 +35,8 @@ //! The opener constructs both halves and hands the state off to //! [`PushDecoderStreamState::into_stream`] for consumption. -use std::collections::VecDeque; +use std::collections::{HashSet, VecDeque}; +use std::ops::Range; use std::sync::Arc; use arrow::array::RecordBatch; @@ -244,6 +245,14 @@ pub(crate) struct PushDecoderStreamState { pub(crate) active_reader: Option, pub(crate) rg_plan: VecDeque, pub(crate) reader: Box, + /// Parquet metadata used to identify projected column chunks belonging to + /// subsequent row groups. + pub(crate) parquet_metadata: Arc, + /// Maximum bytes that may be staged in the push decoder. Required ranges + /// are always fetched, even when they exceed this budget. + pub(crate) prefetch_size: Option, + /// Row groups whose projected ranges were already fetched speculatively. + pub(crate) prefetched_row_groups: HashSet, /// Per-file projection: the mask installed on every decoder and the /// per-batch transform applied by [`Self::project_batch`]. pub(crate) decoder_projection: DecoderProjection, @@ -373,6 +382,20 @@ impl PushDecoderStreamState { let decoder = self.decoder.as_mut().expect("decoder present"); match decoder.try_next_reader() { Ok(DecodeResult::NeedsData(ranges)) => { + let buffered_bytes = self + .decoder + .as_ref() + .expect("decoder present") + .buffered_bytes(); + let ranges = prefetch_row_group_ranges( + ranges, + buffered_bytes, + self.prefetch_size, + &self.rg_plan, + self.decoder_projection.projection_mask(), + &self.parquet_metadata, + &mut self.prefetched_row_groups, + ); let data = self .reader .get_byte_ranges(ranges.clone()) @@ -423,6 +446,68 @@ impl PushDecoderStreamState { } } +/// Append projected column chunks from subsequent row groups to a decoder +/// request while staying within `prefetch_size`. +/// +/// The first entry in `rg_plan` is the row group responsible for `ranges`. +/// Complete projected ranges for later row groups are added in scan order so +/// the push decoder can stage them for future calls to `try_next_reader`. +fn prefetch_row_group_ranges( + mut ranges: Vec>, + buffered_bytes: u64, + prefetch_size: Option, + rg_plan: &VecDeque, + projection: &ProjectionMask, + metadata: &ParquetMetaData, + prefetched_row_groups: &mut HashSet, +) -> Vec> { + let Some(prefetch_size) = prefetch_size.filter(|size| *size > 0) else { + return ranges; + }; + + let requested_bytes = ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + let mut staged_bytes = buffered_bytes.saturating_add(requested_bytes); + let budget = prefetch_size as u64; + if staged_bytes >= budget { + return ranges; + } + + for entry in rg_plan.iter().skip(1) { + if prefetched_row_groups.contains(&entry.rg_index) { + continue; + } + + let row_group = metadata.row_group(entry.rg_index); + let row_group_ranges = row_group + .columns() + .iter() + .enumerate() + .filter(|(column_idx, _)| projection.leaf_included(*column_idx)) + .map(|(_, column)| { + let (start, len) = column.byte_range(); + start..start + len + }) + .collect::>(); + let row_group_bytes = row_group_ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + + if staged_bytes.saturating_add(row_group_bytes) > budget { + break; + } + + ranges.extend(row_group_ranges); + staged_bytes += row_group_bytes; + prefetched_row_groups.insert(entry.rg_index); + } + + ranges +} + #[cfg(test)] mod tests { use super::*; @@ -444,6 +529,11 @@ mod tests { /// column statistics are disjoint: RG0 → 0..1000, RG1 → 1000..2000, /// RG2 → 2000..3000. Returns (metadata, schema). fn build_three_rg_file() -> (Arc, SchemaRef) { + let (_, metadata, schema) = build_three_rg_file_data(); + (metadata, schema) + } + + fn build_three_rg_file_data() -> (Bytes, Arc, SchemaRef) { let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); let mut buf = Vec::new(); let props = WriterProperties::builder() @@ -474,12 +564,145 @@ mod tests { reason = "we want a single range covering the whole file" )] let ranges = vec![0..len]; - md.push_ranges(ranges, vec![file]).unwrap(); + md.push_ranges(ranges, vec![file.clone()]).unwrap(); let DecodeResult::Data(meta) = md.try_decode().unwrap() else { panic!("decoding metadata"); }; assert_eq!(meta.num_row_groups(), 3, "test fixture must have 3 RGs"); - (Arc::new(meta), schema) + (file, Arc::new(meta), schema) + } + + fn column_range(metadata: &ParquetMetaData, row_group: usize) -> Range { + let (start, len) = metadata.row_group(row_group).column(0).byte_range(); + start..start + len + } + + #[test] + fn prefetch_packs_complete_row_groups_within_budget() { + let (metadata, _) = build_three_rg_file(); + let rg0 = column_range(&metadata, 0); + let rg1 = column_range(&metadata, 1); + let rg2 = column_range(&metadata, 2); + let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect(); + let budget = ((rg0.end - rg0.start) + (rg1.end - rg1.start)) as usize; + let mut prefetched = HashSet::new(); + + let ranges = prefetch_row_group_ranges( + vec![rg0.clone()], + 0, + Some(budget), + &rg_plan, + &ProjectionMask::all(), + &metadata, + &mut prefetched, + ); + + assert_eq!(ranges, vec![rg0, rg1]); + assert_eq!(prefetched, HashSet::from([1])); + assert!(!ranges.contains(&rg2)); + } + + #[test] + fn prefetched_bytes_are_staged_for_the_next_push_decoder_reader() { + let (file, metadata, _) = build_three_rg_file_data(); + let mut decoder = + ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata)) + .unwrap() + .build() + .unwrap(); + let requested = match decoder.try_next_reader().unwrap() { + DecodeResult::NeedsData(ranges) => ranges, + other => panic!("expected initial byte request, got {other:?}"), + }; + let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect(); + let rg1 = column_range(&metadata, 1); + let budget = requested + .iter() + .map(|range| range.end - range.start) + .sum::() + + rg1.end + - rg1.start; + let mut prefetched = HashSet::new(); + let ranges = prefetch_row_group_ranges( + requested, + decoder.buffered_bytes(), + Some(budget as usize), + &rg_plan, + &ProjectionMask::all(), + &metadata, + &mut prefetched, + ); + let data = ranges + .iter() + .map(|range| file.slice(range.start as usize..range.end as usize)) + .collect(); + decoder.push_ranges(ranges, data).unwrap(); + + let DecodeResult::Data(first_reader) = decoder.try_next_reader().unwrap() else { + panic!("first row group should be ready"); + }; + assert_eq!( + first_reader + .map(|batch| batch.unwrap().num_rows()) + .sum::(), + 1000 + ); + + let DecodeResult::Data(second_reader) = decoder.try_next_reader().unwrap() else { + panic!("prefetched second row group should not require more I/O"); + }; + assert_eq!( + second_reader + .map(|batch| batch.unwrap().num_rows()) + .sum::(), + 1000 + ); + } + + #[test] + fn prefetch_accounts_for_already_buffered_bytes() { + let (metadata, _) = build_three_rg_file(); + let rg0 = column_range(&metadata, 0); + let rg1 = column_range(&metadata, 1); + let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect(); + let requested = rg0.end - rg0.start; + let next = rg1.end - rg1.start; + let budget = (requested + next) as usize; + let mut prefetched = HashSet::new(); + + let ranges = prefetch_row_group_ranges( + vec![rg0.clone()], + 1, + Some(budget), + &rg_plan, + &ProjectionMask::all(), + &metadata, + &mut prefetched, + ); + + assert_eq!(ranges, vec![rg0]); + assert!(prefetched.is_empty()); + } + + #[test] + fn prefetch_disabled_leaves_request_unchanged() { + let (metadata, _) = build_three_rg_file(); + let rg0 = column_range(&metadata, 0); + let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect(); + let mut prefetched = HashSet::new(); + + let ranges = prefetch_row_group_ranges( + vec![rg0.clone()], + 0, + None, + &rg_plan, + &ProjectionMask::all(), + &metadata, + &mut prefetched, + ); + + assert_eq!(ranges, vec![rg0]); + assert!(prefetched.is_empty()); } /// Create a fresh `(creation_errors, evaluation_errors)` counter pair diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 3443b08475e0d..6fe96cadedaf6 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -631,6 +631,7 @@ impl FileSource for ParquetSource { predicate: self.predicate.clone(), table_schema: self.table_schema.clone(), metadata_size_hint: self.metadata_size_hint, + prefetch_size: self.table_parquet_options.global.prefetch_size, metrics: self.metrics().clone(), parquet_file_reader_factory, pushdown_filters: self.pushdown_filters(), diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 7fff5b6b715ff..bd156bb76b7a3 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -631,6 +631,10 @@ message ParquetOptions { uint64 max_row_group_bytes = 37; } + oneof prefetch_size_opt { + uint64 prefetch_size = 38; + } + ParquetCdcOptions content_defined_chunking = 35; // Optional timezone applied to INT96-coerced timestamps when `coerce_int96` @@ -715,4 +719,4 @@ enum MetricCategory { message ExplainAnalyzeCategoriesNode { bool all = 1; repeated MetricCategory only = 2; -} \ No newline at end of file +} diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 97cc9af230105..560eec6b084a9 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1060,6 +1060,11 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { protobuf::parquet_options::MetadataSizeHintOpt::MetadataSizeHint(v) => Some(v as usize), }) .unwrap_or(None), + prefetch_size: value.prefetch_size_opt.map(|opt| match opt { + protobuf::parquet_options::PrefetchSizeOpt::PrefetchSize(v) => { + Some(v as usize) + } + }).unwrap_or(None), pushdown_filters: value.pushdown_filters, reorder_filters: value.reorder_filters, force_filter_selections: value.force_filter_selections, @@ -1394,6 +1399,16 @@ mod tests { ); } + #[test] + fn test_parquet_options_prefetch_size_round_trip() { + let opts = ParquetOptions { + prefetch_size: Some(20 * 1024 * 1024), + ..ParquetOptions::default() + }; + let recovered = parquet_options_proto_round_trip(opts); + assert_eq!(recovered.prefetch_size, Some(20 * 1024 * 1024)); + } + #[test] fn test_table_parquet_options_coerce_int96_tz_round_trip() { let mut opts = TableParquetOptions::default(); diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 963faa5a3e9cb..3ce954a25afde 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -6451,6 +6451,9 @@ impl serde::Serialize for ParquetOptions { if self.max_row_group_bytes_opt.is_some() { len += 1; } + if self.prefetch_size_opt.is_some() { + len += 1; + } if self.coerce_int96_tz_opt.is_some() { len += 1; } @@ -6631,6 +6634,15 @@ impl serde::Serialize for ParquetOptions { } } } + if let Some(v) = self.prefetch_size_opt.as_ref() { + match v { + parquet_options::PrefetchSizeOpt::PrefetchSize(v) => { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("prefetchSize", ToString::to_string(&v).as_str())?; + } + } + } if let Some(v) = self.coerce_int96_tz_opt.as_ref() { match v { parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(v) => { @@ -6713,6 +6725,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "maxPredicateCacheSize", "max_row_group_bytes", "maxRowGroupBytes", + "prefetch_size", + "prefetchSize", "coerce_int96_tz", "coerceInt96Tz", ]; @@ -6753,6 +6767,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { CoerceInt96, MaxPredicateCacheSize, MaxRowGroupBytes, + PrefetchSize, CoerceInt96Tz, } impl<'de> serde::Deserialize<'de> for GeneratedField { @@ -6809,6 +6824,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "coerceInt96" | "coerce_int96" => Ok(GeneratedField::CoerceInt96), "maxPredicateCacheSize" | "max_predicate_cache_size" => Ok(GeneratedField::MaxPredicateCacheSize), "maxRowGroupBytes" | "max_row_group_bytes" => Ok(GeneratedField::MaxRowGroupBytes), + "prefetchSize" | "prefetch_size" => Ok(GeneratedField::PrefetchSize), "coerceInt96Tz" | "coerce_int96_tz" => Ok(GeneratedField::CoerceInt96Tz), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } @@ -6863,6 +6879,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut coerce_int96_opt__ = None; let mut max_predicate_cache_size_opt__ = None; let mut max_row_group_bytes_opt__ = None; + let mut prefetch_size_opt__ = None; let mut coerce_int96_tz_opt__ = None; while let Some(k) = map_.next_key()? { match k { @@ -7084,6 +7101,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { } max_row_group_bytes_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(x.0)); } + GeneratedField::PrefetchSize => { + if prefetch_size_opt__.is_some() { + return Err(serde::de::Error::duplicate_field("prefetchSize")); + } + prefetch_size_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::PrefetchSizeOpt::PrefetchSize(x.0)); + } GeneratedField::CoerceInt96Tz => { if coerce_int96_tz_opt__.is_some() { return Err(serde::de::Error::duplicate_field("coerceInt96Tz")); @@ -7127,6 +7150,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { coerce_int96_opt: coerce_int96_opt__, max_predicate_cache_size_opt: max_predicate_cache_size_opt__, max_row_group_bytes_opt: max_row_group_bytes_opt__, + prefetch_size_opt: prefetch_size_opt__, coerce_int96_tz_opt: coerce_int96_tz_opt__, }) } diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index 93b97c4f1376c..55b89eeecbad0 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -904,6 +904,8 @@ pub struct ParquetOptions { pub max_row_group_bytes_opt: ::core::option::Option< parquet_options::MaxRowGroupBytesOpt, >, + #[prost(oneof = "parquet_options::PrefetchSizeOpt", tags = "38")] + pub prefetch_size_opt: ::core::option::Option, /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default @@ -973,6 +975,11 @@ pub mod parquet_options { #[prost(uint64, tag = "37")] MaxRowGroupBytes(u64), } + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum PrefetchSizeOpt { + #[prost(uint64, tag = "38")] + PrefetchSize(u64), + } /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index d2e1ca50c812d..11e8c8e928bde 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -909,6 +909,7 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { pruning: value.pruning, skip_metadata: value.skip_metadata, metadata_size_hint_opt: value.metadata_size_hint.map(|v| protobuf::parquet_options::MetadataSizeHintOpt::MetadataSizeHint(v as u64)), + prefetch_size_opt: value.prefetch_size.map(|v| protobuf::parquet_options::PrefetchSizeOpt::PrefetchSize(v as u64)), pushdown_filters: value.pushdown_filters, reorder_filters: value.reorder_filters, force_filter_selections: value.force_filter_selections, diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index 93b97c4f1376c..55b89eeecbad0 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -904,6 +904,8 @@ pub struct ParquetOptions { pub max_row_group_bytes_opt: ::core::option::Option< parquet_options::MaxRowGroupBytesOpt, >, + #[prost(oneof = "parquet_options::PrefetchSizeOpt", tags = "38")] + pub prefetch_size_opt: ::core::option::Option, /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default @@ -973,6 +975,11 @@ pub mod parquet_options { #[prost(uint64, tag = "37")] MaxRowGroupBytes(u64), } + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum PrefetchSizeOpt { + #[prost(uint64, tag = "38")] + PrefetchSize(u64), + } /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default diff --git a/datafusion/proto/src/logical_plan/file_formats.rs b/datafusion/proto/src/logical_plan/file_formats.rs index 8940b16bf83f5..3ad63e9de94d3 100644 --- a/datafusion/proto/src/logical_plan/file_formats.rs +++ b/datafusion/proto/src/logical_plan/file_formats.rs @@ -407,6 +407,9 @@ mod parquet { metadata_size_hint_opt: global_options.global.metadata_size_hint.map(|size| { parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size as u64) }), + prefetch_size_opt: global_options.global.prefetch_size.map(|size| { + parquet_options::PrefetchSizeOpt::PrefetchSize(size as u64) + }), pushdown_filters: global_options.global.pushdown_filters, reorder_filters: global_options.global.reorder_filters, force_filter_selections: global_options.global.force_filter_selections, @@ -543,6 +546,11 @@ mod parquet { *size as usize } }), + prefetch_size: proto.prefetch_size_opt.as_ref().map(|opt| match opt { + parquet_options::PrefetchSizeOpt::PrefetchSize(size) => { + *size as usize + } + }), pushdown_filters: proto.pushdown_filters, reorder_filters: proto.reorder_filters, force_filter_selections: proto.force_filter_selections, diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index bf45564e26333..fe2414fd9f07b 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -259,6 +259,7 @@ datafusion.execution.parquet.max_row_group_size 1048576 datafusion.execution.parquet.maximum_buffered_record_batches_per_stream 2 datafusion.execution.parquet.maximum_parallel_row_group_writers 1 datafusion.execution.parquet.metadata_size_hint 524288 +datafusion.execution.parquet.prefetch_size NULL datafusion.execution.parquet.pruning true datafusion.execution.parquet.pushdown_filters false datafusion.execution.parquet.reorder_filters false @@ -418,6 +419,7 @@ datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum datafusion.execution.parquet.maximum_buffered_record_batches_per_stream 2 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. datafusion.execution.parquet.maximum_parallel_row_group_writers 1 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. datafusion.execution.parquet.metadata_size_hint 524288 (reading) If specified, the parquet reader will try and fetch the last `size_hint` bytes of the parquet file optimistically. If not specified, two reads are required: One read to fetch the 8-byte parquet footer and another to fetch the metadata length encoded in the footer Default setting to 512 KiB, which should be sufficient for most parquet files, it can reduce one I/O operation per parquet file. If the metadata is larger than the hint, two reads will still be performed. +datafusion.execution.parquet.prefetch_size NULL (reading) If specified, the parquet reader will prefetch data for subsequent row groups when the projected column chunks fit within this many bytes. The required ranges for the current row group are always read, even when they exceed this value. None disables data prefetching. datafusion.execution.parquet.pruning true (reading) If true, the parquet reader attempts to skip entire row groups based on the predicate in the query and the metadata (min/max values) stored in the parquet file datafusion.execution.parquet.pushdown_filters false (reading) If true, filter expressions are be applied during the parquet decoding operation to reduce the number of rows decoded. This optimization is sometimes called "late materialization". datafusion.execution.parquet.reorder_filters false (reading) If true, filter expressions evaluated during the parquet decoding operation will be reordered heuristically to minimize the cost of evaluation. If false, the filters are applied in the same order as written in the query diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 03340c366d70f..7a908fbbc2d60 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -84,6 +84,7 @@ The following configuration settings are available: | datafusion.execution.parquet.pruning | true | (reading) If true, the parquet reader attempts to skip entire row groups based on the predicate in the query and the metadata (min/max values) stored in the parquet file | | datafusion.execution.parquet.skip_metadata | true | (reading) If true, the parquet reader skip the optional embedded metadata that may be in the file Schema. This setting can help avoid schema conflicts when querying multiple parquet files with schemas containing compatible types but different metadata | | datafusion.execution.parquet.metadata_size_hint | 524288 | (reading) If specified, the parquet reader will try and fetch the last `size_hint` bytes of the parquet file optimistically. If not specified, two reads are required: One read to fetch the 8-byte parquet footer and another to fetch the metadata length encoded in the footer Default setting to 512 KiB, which should be sufficient for most parquet files, it can reduce one I/O operation per parquet file. If the metadata is larger than the hint, two reads will still be performed. | +| datafusion.execution.parquet.prefetch_size | NULL | (reading) If specified, the parquet reader will prefetch data for subsequent row groups when the projected column chunks fit within this many bytes. The required ranges for the current row group are always read, even when they exceed this value. None disables data prefetching. | | datafusion.execution.parquet.pushdown_filters | false | (reading) If true, filter expressions are be applied during the parquet decoding operation to reduce the number of rows decoded. This optimization is sometimes called "late materialization". | | datafusion.execution.parquet.reorder_filters | false | (reading) If true, filter expressions evaluated during the parquet decoding operation will be reordered heuristically to minimize the cost of evaluation. If false, the filters are applied in the same order as written in the query | | datafusion.execution.parquet.force_filter_selections | false | (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. | From 838732c4b1afb877dc9f5d86f2233e63c12411e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 4 Aug 2026 09:00:05 +0200 Subject: [PATCH 2/2] Use datafusion_common::HashSet for prefetched row groups Switch `prefetched_row_groups` to the `datafusion_common` HashSet re-export so the opener no longer needs an inline `std::collections` path, and drop a redundant assertion in the prefetch budget test. Co-Authored-By: Claude Opus 5 --- datafusion/datasource-parquet/src/opener/mod.rs | 2 +- datafusion/datasource-parquet/src/push_decoder.rs | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index bc649efcc5e42..bb15177bf0240 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1488,7 +1488,7 @@ impl RowGroupsPrunedParquetOpen { reader: prepared.async_file_reader, parquet_metadata: Arc::clone(&file_metadata), prefetch_size: prepared.prefetch_size, - prefetched_row_groups: std::collections::HashSet::new(), + prefetched_row_groups: HashSet::new(), decoder_projection, arrow_reader_metrics, predicate_cache_inner_records, diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index c67d087ff6b19..6207d802ba514 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -35,7 +35,7 @@ //! The opener constructs both halves and hands the state off to //! [`PushDecoderStreamState::into_stream`] for consumption. -use std::collections::{HashSet, VecDeque}; +use std::collections::VecDeque; use std::ops::Range; use std::sync::Arc; @@ -54,7 +54,7 @@ use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder}; use parquet::file::metadata::ParquetMetaData; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::{DataFusionError, HashSet, Result}; use datafusion_physical_expr::expressions::DynamicFilterTracking; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_plan::metrics::{BaselineMetrics, Count, Gauge}; @@ -582,7 +582,6 @@ mod tests { let (metadata, _) = build_three_rg_file(); let rg0 = column_range(&metadata, 0); let rg1 = column_range(&metadata, 1); - let rg2 = column_range(&metadata, 2); let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect(); let budget = ((rg0.end - rg0.start) + (rg1.end - rg1.start)) as usize; let mut prefetched = HashSet::new(); @@ -597,9 +596,9 @@ mod tests { &mut prefetched, ); + // RG2 does not fit in the budget, so it is left out entirely. assert_eq!(ranges, vec![rg0, rg1]); assert_eq!(prefetched, HashSet::from([1])); - assert!(!ranges.contains(&rg2)); } #[test]