Skip to content

Commit 4d02e1e

Browse files
committed
Prefetch Parquet row groups with a bounded budget
1 parent 2880e10 commit 4d02e1e

15 files changed

Lines changed: 686 additions & 4 deletions

File tree

‎benchmarks/parquet_prefetch_benchmark.md‎

Lines changed: 371 additions & 0 deletions
Large diffs are not rendered by default.

‎datafusion/common/src/config.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1133,6 +1133,13 @@ config_namespace! {
11331133
/// the hint, two reads will still be performed.
11341134
pub metadata_size_hint: Option<usize>, default = Some(512 * 1024)
11351135

1136+
/// (reading) If specified, the parquet reader will prefetch data for
1137+
/// subsequent row groups when the projected column chunks fit within
1138+
/// this many bytes. The required ranges for the current row group are
1139+
/// always read, even when they exceed this value. None disables data
1140+
/// prefetching.
1141+
pub prefetch_size: Option<usize>, default = None
1142+
11361143
/// (reading) If true, filter expressions are be applied during the parquet decoding operation to
11371144
/// reduce the number of rows decoded. This optimization is sometimes called "late materialization".
11381145
pub pushdown_filters: bool, default = false

‎datafusion/common/src/file_options/parquet_writer.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,7 @@ impl ParquetOptions {
235235
pruning: _,
236236
skip_metadata: _,
237237
metadata_size_hint: _,
238+
prefetch_size: _,
238239
pushdown_filters: _,
239240
reorder_filters: _,
240241
force_filter_selections: _, // not used for writer props
@@ -491,6 +492,7 @@ mod tests {
491492
pruning: defaults.pruning,
492493
skip_metadata: defaults.skip_metadata,
493494
metadata_size_hint: defaults.metadata_size_hint,
495+
prefetch_size: defaults.prefetch_size,
494496
pushdown_filters: defaults.pushdown_filters,
495497
reorder_filters: defaults.reorder_filters,
496498
force_filter_selections: defaults.force_filter_selections,
@@ -610,6 +612,7 @@ mod tests {
610612
pruning: global_options_defaults.pruning,
611613
skip_metadata: global_options_defaults.skip_metadata,
612614
metadata_size_hint: global_options_defaults.metadata_size_hint,
615+
prefetch_size: global_options_defaults.prefetch_size,
613616
pushdown_filters: global_options_defaults.pushdown_filters,
614617
reorder_filters: global_options_defaults.reorder_filters,
615618
force_filter_selections: global_options_defaults.force_filter_selections,

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,8 @@ pub(super) struct ParquetMorselizer {
252252
/// Optional hint for how large the initial request to read parquet metadata
253253
/// should be
254254
pub metadata_size_hint: Option<usize>,
255+
/// Maximum number of bytes buffered by speculative row-group prefetching.
256+
pub prefetch_size: Option<usize>,
255257
/// Metrics for reporting
256258
pub metrics: ExecutionPlanMetricsSet,
257259
/// Factory for instantiating parquet reader
@@ -425,6 +427,7 @@ struct PreparedParquetOpen {
425427
baseline_metrics: BaselineMetrics,
426428
file_pruner: Option<FilePruner>,
427429
metadata_size_hint: Option<usize>,
430+
prefetch_size: Option<usize>,
428431
metrics: ExecutionPlanMetricsSet,
429432
parquet_file_reader_factory: Arc<dyn ParquetFileReaderFactory>,
430433
async_file_reader: Box<dyn AsyncFileReader>,
@@ -828,6 +831,7 @@ impl ParquetMorselizer {
828831
baseline_metrics,
829832
file_pruner,
830833
metadata_size_hint,
834+
prefetch_size: self.prefetch_size,
831835
metrics: self.metrics.clone(),
832836
parquet_file_reader_factory: Arc::clone(&self.parquet_file_reader_factory),
833837
async_file_reader,
@@ -1482,6 +1486,9 @@ impl RowGroupsPrunedParquetOpen {
14821486
active_reader: None,
14831487
rg_plan,
14841488
reader: prepared.async_file_reader,
1489+
parquet_metadata: Arc::clone(&file_metadata),
1490+
prefetch_size: prepared.prefetch_size,
1491+
prefetched_row_groups: std::collections::HashSet::new(),
14851492
decoder_projection,
14861493
arrow_reader_metrics,
14871494
predicate_cache_inner_records,
@@ -2014,6 +2021,7 @@ mod test {
20142021
predicate: self.predicate,
20152022
table_schema,
20162023
metadata_size_hint: self.metadata_size_hint,
2024+
prefetch_size: None,
20172025
metrics: self.metrics,
20182026
parquet_file_reader_factory: self
20192027
.parquet_file_reader_factory

‎datafusion/datasource-parquet/src/push_decoder.rs‎

Lines changed: 226 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@
3535
//! The opener constructs both halves and hands the state off to
3636
//! [`PushDecoderStreamState::into_stream`] for consumption.
3737
38-
use std::collections::VecDeque;
38+
use std::collections::{HashSet, VecDeque};
39+
use std::ops::Range;
3940
use std::sync::Arc;
4041

4142
use arrow::array::RecordBatch;
@@ -244,6 +245,14 @@ pub(crate) struct PushDecoderStreamState {
244245
pub(crate) active_reader: Option<ParquetRecordBatchReader>,
245246
pub(crate) rg_plan: VecDeque<RgPlanEntry>,
246247
pub(crate) reader: Box<dyn AsyncFileReader>,
248+
/// Parquet metadata used to identify projected column chunks belonging to
249+
/// subsequent row groups.
250+
pub(crate) parquet_metadata: Arc<ParquetMetaData>,
251+
/// Maximum bytes that may be staged in the push decoder. Required ranges
252+
/// are always fetched, even when they exceed this budget.
253+
pub(crate) prefetch_size: Option<usize>,
254+
/// Row groups whose projected ranges were already fetched speculatively.
255+
pub(crate) prefetched_row_groups: HashSet<usize>,
247256
/// Per-file projection: the mask installed on every decoder and the
248257
/// per-batch transform applied by [`Self::project_batch`].
249258
pub(crate) decoder_projection: DecoderProjection,
@@ -373,6 +382,20 @@ impl PushDecoderStreamState {
373382
let decoder = self.decoder.as_mut().expect("decoder present");
374383
match decoder.try_next_reader() {
375384
Ok(DecodeResult::NeedsData(ranges)) => {
385+
let buffered_bytes = self
386+
.decoder
387+
.as_ref()
388+
.expect("decoder present")
389+
.buffered_bytes();
390+
let ranges = prefetch_row_group_ranges(
391+
ranges,
392+
buffered_bytes,
393+
self.prefetch_size,
394+
&self.rg_plan,
395+
self.decoder_projection.projection_mask(),
396+
&self.parquet_metadata,
397+
&mut self.prefetched_row_groups,
398+
);
376399
let data = self
377400
.reader
378401
.get_byte_ranges(ranges.clone())
@@ -423,6 +446,68 @@ impl PushDecoderStreamState {
423446
}
424447
}
425448

449+
/// Append projected column chunks from subsequent row groups to a decoder
450+
/// request while staying within `prefetch_size`.
451+
///
452+
/// The first entry in `rg_plan` is the row group responsible for `ranges`.
453+
/// Complete projected ranges for later row groups are added in scan order so
454+
/// the push decoder can stage them for future calls to `try_next_reader`.
455+
fn prefetch_row_group_ranges(
456+
mut ranges: Vec<Range<u64>>,
457+
buffered_bytes: u64,
458+
prefetch_size: Option<usize>,
459+
rg_plan: &VecDeque<RgPlanEntry>,
460+
projection: &ProjectionMask,
461+
metadata: &ParquetMetaData,
462+
prefetched_row_groups: &mut HashSet<usize>,
463+
) -> Vec<Range<u64>> {
464+
let Some(prefetch_size) = prefetch_size.filter(|size| *size > 0) else {
465+
return ranges;
466+
};
467+
468+
let requested_bytes = ranges
469+
.iter()
470+
.map(|range| range.end - range.start)
471+
.sum::<u64>();
472+
let mut staged_bytes = buffered_bytes.saturating_add(requested_bytes);
473+
let budget = prefetch_size as u64;
474+
if staged_bytes >= budget {
475+
return ranges;
476+
}
477+
478+
for entry in rg_plan.iter().skip(1) {
479+
if prefetched_row_groups.contains(&entry.rg_index) {
480+
continue;
481+
}
482+
483+
let row_group = metadata.row_group(entry.rg_index);
484+
let row_group_ranges = row_group
485+
.columns()
486+
.iter()
487+
.enumerate()
488+
.filter(|(column_idx, _)| projection.leaf_included(*column_idx))
489+
.map(|(_, column)| {
490+
let (start, len) = column.byte_range();
491+
start..start + len
492+
})
493+
.collect::<Vec<_>>();
494+
let row_group_bytes = row_group_ranges
495+
.iter()
496+
.map(|range| range.end - range.start)
497+
.sum::<u64>();
498+
499+
if staged_bytes.saturating_add(row_group_bytes) > budget {
500+
break;
501+
}
502+
503+
ranges.extend(row_group_ranges);
504+
staged_bytes += row_group_bytes;
505+
prefetched_row_groups.insert(entry.rg_index);
506+
}
507+
508+
ranges
509+
}
510+
426511
#[cfg(test)]
427512
mod tests {
428513
use super::*;
@@ -444,6 +529,11 @@ mod tests {
444529
/// column statistics are disjoint: RG0 → 0..1000, RG1 → 1000..2000,
445530
/// RG2 → 2000..3000. Returns (metadata, schema).
446531
fn build_three_rg_file() -> (Arc<ParquetMetaData>, SchemaRef) {
532+
let (_, metadata, schema) = build_three_rg_file_data();
533+
(metadata, schema)
534+
}
535+
536+
fn build_three_rg_file_data() -> (Bytes, Arc<ParquetMetaData>, SchemaRef) {
447537
let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
448538
let mut buf = Vec::new();
449539
let props = WriterProperties::builder()
@@ -474,12 +564,145 @@ mod tests {
474564
reason = "we want a single range covering the whole file"
475565
)]
476566
let ranges = vec![0..len];
477-
md.push_ranges(ranges, vec![file]).unwrap();
567+
md.push_ranges(ranges, vec![file.clone()]).unwrap();
478568
let DecodeResult::Data(meta) = md.try_decode().unwrap() else {
479569
panic!("decoding metadata");
480570
};
481571
assert_eq!(meta.num_row_groups(), 3, "test fixture must have 3 RGs");
482-
(Arc::new(meta), schema)
572+
(file, Arc::new(meta), schema)
573+
}
574+
575+
fn column_range(metadata: &ParquetMetaData, row_group: usize) -> Range<u64> {
576+
let (start, len) = metadata.row_group(row_group).column(0).byte_range();
577+
start..start + len
578+
}
579+
580+
#[test]
581+
fn prefetch_packs_complete_row_groups_within_budget() {
582+
let (metadata, _) = build_three_rg_file();
583+
let rg0 = column_range(&metadata, 0);
584+
let rg1 = column_range(&metadata, 1);
585+
let rg2 = column_range(&metadata, 2);
586+
let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect();
587+
let budget = ((rg0.end - rg0.start) + (rg1.end - rg1.start)) as usize;
588+
let mut prefetched = HashSet::new();
589+
590+
let ranges = prefetch_row_group_ranges(
591+
vec![rg0.clone()],
592+
0,
593+
Some(budget),
594+
&rg_plan,
595+
&ProjectionMask::all(),
596+
&metadata,
597+
&mut prefetched,
598+
);
599+
600+
assert_eq!(ranges, vec![rg0, rg1]);
601+
assert_eq!(prefetched, HashSet::from([1]));
602+
assert!(!ranges.contains(&rg2));
603+
}
604+
605+
#[test]
606+
fn prefetched_bytes_are_staged_for_the_next_push_decoder_reader() {
607+
let (file, metadata, _) = build_three_rg_file_data();
608+
let mut decoder =
609+
ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata))
610+
.unwrap()
611+
.build()
612+
.unwrap();
613+
let requested = match decoder.try_next_reader().unwrap() {
614+
DecodeResult::NeedsData(ranges) => ranges,
615+
other => panic!("expected initial byte request, got {other:?}"),
616+
};
617+
let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect();
618+
let rg1 = column_range(&metadata, 1);
619+
let budget = requested
620+
.iter()
621+
.map(|range| range.end - range.start)
622+
.sum::<u64>()
623+
+ rg1.end
624+
- rg1.start;
625+
let mut prefetched = HashSet::new();
626+
let ranges = prefetch_row_group_ranges(
627+
requested,
628+
decoder.buffered_bytes(),
629+
Some(budget as usize),
630+
&rg_plan,
631+
&ProjectionMask::all(),
632+
&metadata,
633+
&mut prefetched,
634+
);
635+
let data = ranges
636+
.iter()
637+
.map(|range| file.slice(range.start as usize..range.end as usize))
638+
.collect();
639+
decoder.push_ranges(ranges, data).unwrap();
640+
641+
let DecodeResult::Data(first_reader) = decoder.try_next_reader().unwrap() else {
642+
panic!("first row group should be ready");
643+
};
644+
assert_eq!(
645+
first_reader
646+
.map(|batch| batch.unwrap().num_rows())
647+
.sum::<usize>(),
648+
1000
649+
);
650+
651+
let DecodeResult::Data(second_reader) = decoder.try_next_reader().unwrap() else {
652+
panic!("prefetched second row group should not require more I/O");
653+
};
654+
assert_eq!(
655+
second_reader
656+
.map(|batch| batch.unwrap().num_rows())
657+
.sum::<usize>(),
658+
1000
659+
);
660+
}
661+
662+
#[test]
663+
fn prefetch_accounts_for_already_buffered_bytes() {
664+
let (metadata, _) = build_three_rg_file();
665+
let rg0 = column_range(&metadata, 0);
666+
let rg1 = column_range(&metadata, 1);
667+
let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect();
668+
let requested = rg0.end - rg0.start;
669+
let next = rg1.end - rg1.start;
670+
let budget = (requested + next) as usize;
671+
let mut prefetched = HashSet::new();
672+
673+
let ranges = prefetch_row_group_ranges(
674+
vec![rg0.clone()],
675+
1,
676+
Some(budget),
677+
&rg_plan,
678+
&ProjectionMask::all(),
679+
&metadata,
680+
&mut prefetched,
681+
);
682+
683+
assert_eq!(ranges, vec![rg0]);
684+
assert!(prefetched.is_empty());
685+
}
686+
687+
#[test]
688+
fn prefetch_disabled_leaves_request_unchanged() {
689+
let (metadata, _) = build_three_rg_file();
690+
let rg0 = column_range(&metadata, 0);
691+
let rg_plan = (0..3).map(|rg_index| RgPlanEntry { rg_index }).collect();
692+
let mut prefetched = HashSet::new();
693+
694+
let ranges = prefetch_row_group_ranges(
695+
vec![rg0.clone()],
696+
0,
697+
None,
698+
&rg_plan,
699+
&ProjectionMask::all(),
700+
&metadata,
701+
&mut prefetched,
702+
);
703+
704+
assert_eq!(ranges, vec![rg0]);
705+
assert!(prefetched.is_empty());
483706
}
484707

485708
/// Create a fresh `(creation_errors, evaluation_errors)` counter pair

‎datafusion/datasource-parquet/src/source.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -631,6 +631,7 @@ impl FileSource for ParquetSource {
631631
predicate: self.predicate.clone(),
632632
table_schema: self.table_schema.clone(),
633633
metadata_size_hint: self.metadata_size_hint,
634+
prefetch_size: self.table_parquet_options.global.prefetch_size,
634635
metrics: self.metrics().clone(),
635636
parquet_file_reader_factory,
636637
pushdown_filters: self.pushdown_filters(),

‎datafusion/proto-common/proto/datafusion_common.proto‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -631,6 +631,10 @@ message ParquetOptions {
631631
uint64 max_row_group_bytes = 37;
632632
}
633633

634+
oneof prefetch_size_opt {
635+
uint64 prefetch_size = 38;
636+
}
637+
634638
ParquetCdcOptions content_defined_chunking = 35;
635639

636640
// Optional timezone applied to INT96-coerced timestamps when `coerce_int96`
@@ -715,4 +719,4 @@ enum MetricCategory {
715719
message ExplainAnalyzeCategoriesNode {
716720
bool all = 1;
717721
repeated MetricCategory only = 2;
718-
}
722+
}

0 commit comments

Comments
 (0)