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 ;
3940use std:: sync:: Arc ;
4041
4142use 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) ]
427512mod 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
0 commit comments