@@ -44,7 +44,6 @@ use datafusion_pruning::FilePruner;
4444use futures:: StreamExt ;
4545use futures:: TryStreamExt ;
4646use futures:: stream:: BoxStream ;
47- use itertools:: Itertools ;
4847use object_store:: path:: Path ;
4948use tokio:: sync:: OnceCell ;
5049use tracing:: Instrument ;
@@ -79,8 +78,52 @@ use crate::persistent::cache::CachedVortexMetadata;
7978use crate :: persistent:: stream:: PrunableStream ;
8079use crate :: reader:: VortexReaderFactory ;
8180
82- pub ( crate ) type NaturalSplitRanges = Arc < [ Range < u64 > ] > ;
83- pub ( crate ) type NaturalSplitCache = DashMap < Path , Arc < OnceCell < NaturalSplitRanges > > > ;
81+ pub ( crate ) type NaturalSplitCache = DashMap < Path , Arc < OnceCell < Arc < NaturalSplits > > > > ;
82+
83+ /// A file's natural split boundaries, with each split's byte assignment precomputed.
84+ ///
85+ /// `row_boundaries` is the sorted, deduplicated split point list covering the whole file, shaped
86+ /// so it can be handed straight to [`ScanBuilder::with_natural_splits`]. `assignment_bytes` holds
87+ /// the byte that owns each split (see [`split_assignment_byte`]) and is sorted, so translating a
88+ /// DataFusion byte range into rows is two binary searches rather than a scan over every split.
89+ #[ derive( Debug ) ]
90+ pub ( crate ) struct NaturalSplits {
91+ row_boundaries : Arc < [ u64 ] > ,
92+ assignment_bytes : Box < [ u64 ] > ,
93+ }
94+
95+ impl NaturalSplits {
96+ fn new ( row_boundaries : Arc < [ u64 ] > , total_size : u64 ) -> Self {
97+ let row_count = row_boundaries. last ( ) . copied ( ) . unwrap_or_default ( ) ;
98+ let assignment_bytes = if row_count == 0 {
99+ Box :: default ( )
100+ } else {
101+ row_boundaries
102+ . windows ( 2 )
103+ . enumerate ( )
104+ . map ( |( idx, boundaries) | {
105+ split_assignment_byte (
106+ idx,
107+ & ( boundaries[ 0 ] ..boundaries[ 1 ] ) ,
108+ row_count,
109+ total_size,
110+ )
111+ } )
112+ . collect ( )
113+ } ;
114+
115+ debug_assert ! ( assignment_bytes. is_sorted( ) ) ;
116+ debug_assert_eq ! (
117+ assignment_bytes. len( ) + usize :: from( !row_boundaries. is_empty( ) ) ,
118+ row_boundaries. len( )
119+ ) ;
120+
121+ Self {
122+ row_boundaries,
123+ assignment_bytes,
124+ }
125+ }
126+ }
84127
85128/// Creates morsel planners for Vortex files.
86129pub struct VortexMorselizer {
@@ -114,7 +157,7 @@ pub struct VortexMorselizer {
114157 /// a file reader the first time we read a file.
115158 pub layout_readers : Arc < DashMap < Path , Weak < dyn LayoutReader > > > ,
116159 /// Shared full-file natural split ranges keyed by file path.
117- pub natural_split_ranges : Arc < NaturalSplitCache > ,
160+ pub ( crate ) natural_splits : Arc < NaturalSplitCache > ,
118161 /// Whether the query has output ordering specified
119162 pub has_output_ordering : bool ,
120163
@@ -139,7 +182,7 @@ impl std::fmt::Debug for VortexMorselizer {
139182 . field ( "limit" , & self . limit )
140183 . field ( "df_metrics" , & self . df_metrics )
141184 . field ( "layout_readers" , & self . layout_readers )
142- . field ( "natural_split_ranges " , & self . natural_split_ranges )
185+ . field ( "natural_splits " , & self . natural_splits )
143186 . field ( "has_output_ordering" , & self . has_output_ordering )
144187 . field ( "file_metadata_cache" , & self . file_metadata_cache )
145188 . field ( "projection_pushdown" , & self . projection_pushdown )
@@ -194,7 +237,7 @@ impl VortexMorselPlanner {
194237 let unified_file_schema = Arc :: clone ( morselizer. table_schema . file_schema ( ) ) ;
195238 let limit = morselizer. limit ;
196239 let layout_readers = Arc :: clone ( & morselizer. layout_readers ) ;
197- let natural_split_ranges = Arc :: clone ( & morselizer. natural_split_ranges ) ;
240+ let natural_splits = Arc :: clone ( & morselizer. natural_splits ) ;
198241 let has_output_ordering = morselizer. has_output_ordering ;
199242 let scan_concurrency = morselizer. scan_concurrency ;
200243
@@ -273,7 +316,7 @@ impl VortexMorselPlanner {
273316 unified_file_schema,
274317 limit,
275318 layout_readers,
276- natural_split_ranges ,
319+ natural_splits ,
277320 has_output_ordering,
278321 expr_convertor,
279322 projection_pushdown,
@@ -299,7 +342,7 @@ struct FileOpenState {
299342 unified_file_schema : SchemaRef ,
300343 limit : Option < u64 > ,
301344 layout_readers : Arc < DashMap < Path , Weak < dyn LayoutReader > > > ,
302- natural_split_ranges : Arc < NaturalSplitCache > ,
345+ natural_splits : Arc < NaturalSplitCache > ,
303346 has_output_ordering : bool ,
304347 expr_convertor : Arc < dyn ExpressionConvertor > ,
305348 projection_pushdown : bool ,
@@ -326,13 +369,15 @@ enum State {
326369 vxf : VortexFile ,
327370 layout_reader : Arc < dyn LayoutReader > ,
328371 byte_range : Range < u64 > ,
329- split_ranges : Arc < OnceCell < NaturalSplitRanges > > ,
372+ split_ranges : Arc < OnceCell < Arc < NaturalSplits > > > ,
330373 } ,
331374 PrepareScan {
332375 state : FileOpenState ,
333376 vxf : VortexFile ,
334377 layout_reader : Arc < dyn LayoutReader > ,
335378 row_range : Option < Range < u64 > > ,
379+ /// Set alongside `row_range`, so the scan can skip re-walking the layout for splits.
380+ natural_splits : Option < Arc < NaturalSplits > > ,
336381 } ,
337382 PreparedScan {
338383 scan : RepeatedScan ,
@@ -457,7 +502,7 @@ impl std::fmt::Debug for FileOpenState {
457502 . field ( "unified_file_schema" , & self . unified_file_schema )
458503 . field ( "limit" , & self . limit )
459504 . field ( "layout_readers" , & self . layout_readers )
460- . field ( "natural_split_ranges " , & self . natural_split_ranges )
505+ . field ( "natural_splits " , & self . natural_splits )
461506 . field ( "has_output_ordering" , & self . has_output_ordering )
462507 . field ( "expr_convertor" , & "<expr_convertor>" )
463508 . field ( "projection_pushdown" , & self . projection_pushdown )
@@ -557,13 +602,14 @@ impl MorselPlanner for VortexMorselPlanner {
557602 vxf,
558603 layout_reader,
559604 row_range: None ,
605+ natural_splits: None ,
560606 } ,
561607 } ,
562608 ) ] ) ) ) ;
563609 } ;
564610
565611 let split_ranges = natural_split_cell_for_file (
566- state. natural_split_ranges . as_ref ( ) ,
612+ state. natural_splits . as_ref ( ) ,
567613 & state. file . object_meta . location ,
568614 ) ;
569615
@@ -591,22 +637,24 @@ impl MorselPlanner for VortexMorselPlanner {
591637 Ok ( Some (
592638 MorselPlan :: new ( ) . with_pending_planner (
593639 async move {
594- let split_ranges = split_ranges
595- . get_or_try_init ( || async {
596- compute_natural_split_ranges ( layout_reader. as_ref ( ) )
597- } )
598- . await ?;
640+ let natural_splits = Arc :: clone (
641+ split_ranges
642+ . get_or_try_init ( || async {
643+ compute_natural_splits ( layout_reader. as_ref ( ) , total_size)
644+ } )
645+ . await ?,
646+ ) ;
599647
600648 let new_state = match split_aligned_row_range (
601649 byte_range,
602- total_size,
603- split_ranges. as_ref ( ) ,
650+ natural_splits. as_ref ( ) ,
604651 ) {
605652 Some ( row_range) => State :: PrepareScan {
606653 state,
607654 vxf,
608655 layout_reader,
609656 row_range : Some ( row_range) ,
657+ natural_splits : Some ( natural_splits) ,
610658 } ,
611659 None => State :: Done ,
612660 } ;
@@ -623,6 +671,7 @@ impl MorselPlanner for VortexMorselPlanner {
623671 vxf,
624672 layout_reader,
625673 row_range,
674+ natural_splits,
626675 } => {
627676 let FileOpenState {
628677 file,
@@ -723,6 +772,10 @@ impl MorselPlanner for VortexMorselPlanner {
723772 if let Some ( row_range) = row_range {
724773 scan_builder = scan_builder. with_row_range ( row_range) ;
725774 }
775+ if let Some ( natural_splits) = natural_splits {
776+ scan_builder = scan_builder
777+ . with_natural_splits ( Arc :: clone ( & natural_splits. row_boundaries ) ) ;
778+ }
726779
727780 let filter = filter
728781 . and_then ( |f| {
@@ -899,10 +952,10 @@ fn partial_file_byte_range(file: &PartitionedFile) -> DFResult<Option<Range<u64>
899952}
900953
901954fn natural_split_cell_for_file (
902- natural_split_ranges : & NaturalSplitCache ,
955+ natural_splits : & NaturalSplitCache ,
903956 path : & Path ,
904- ) -> Arc < OnceCell < NaturalSplitRanges > > {
905- match natural_split_ranges . entry ( path. clone ( ) ) {
957+ ) -> Arc < OnceCell < Arc < NaturalSplits > > > {
958+ match natural_splits . entry ( path. clone ( ) ) {
906959 Entry :: Occupied ( entry) => Arc :: clone ( entry. get ( ) ) ,
907960 Entry :: Vacant ( entry) => {
908961 let split_ranges = Arc :: new ( OnceCell :: new ( ) ) ;
@@ -912,52 +965,46 @@ fn natural_split_cell_for_file(
912965 }
913966}
914967
915- fn compute_natural_split_ranges ( layout_reader : & dyn LayoutReader ) -> DFResult < NaturalSplitRanges > {
968+ fn compute_natural_splits (
969+ layout_reader : & dyn LayoutReader ,
970+ total_size : u64 ,
971+ ) -> DFResult < Arc < NaturalSplits > > {
916972 let row_count = layout_reader. row_count ( ) ;
917973 let row_range = 0 ..row_count;
918- let split_points : Vec < _ > = SplitBy :: Layout
974+ let row_boundaries = SplitBy :: Layout
919975 . splits ( layout_reader, & row_range, & [ FieldMask :: All ] )
920- . map_err ( |e| exec_datafusion_err ! ( "Failed to compute Vortex natural splits: {e}" ) ) ?
921- . into_iter ( )
922- . tuple_windows ( )
923- . map ( |( s, e) | s..e)
924- . collect :: < Vec < _ > > ( ) ;
976+ . map_err ( |e| exec_datafusion_err ! ( "Failed to compute Vortex natural splits: {e}" ) ) ?;
925977
926- Ok ( split_points. into ( ) )
978+ Ok ( Arc :: new ( NaturalSplits :: new (
979+ row_boundaries. into ( ) ,
980+ total_size,
981+ ) ) )
927982}
928983
929984/// Translate a DataFusion byte range to the contiguous natural split ranges it owns.
930985/// Most splits are assigned by midpoint, but the leading split stays with the range that owns
931986/// byte 0 so a tiny first byte range still claims the first rows.
932987fn split_aligned_row_range (
933988 byte_range : Range < u64 > ,
934- total_size : u64 ,
935- split_ranges : & [ Range < u64 > ] ,
989+ natural_splits : & NaturalSplits ,
936990) -> Option < Range < u64 > > {
937991 if byte_range. start >= byte_range. end {
938992 return None ;
939993 }
940994
941- let row_count = split_ranges. last ( ) . map ( |split| split. end ) ?;
942- if row_count == 0 {
995+ let first_split = natural_splits
996+ . assignment_bytes
997+ . partition_point ( |& assignment_byte| assignment_byte < byte_range. start ) ;
998+ let after_last_split = natural_splits
999+ . assignment_bytes
1000+ . partition_point ( |& assignment_byte| assignment_byte < byte_range. end ) ;
1001+ if first_split == after_last_split {
9431002 return None ;
9441003 }
9451004
946- let mut owned_splits = split_ranges
947- . iter ( )
948- . enumerate ( )
949- . filter_map ( |( idx, split_range) | {
950- let assignment_byte = split_assignment_byte ( idx, split_range, row_count, total_size) ;
951- byte_range. contains ( & assignment_byte) . then_some ( split_range)
952- } ) ;
953-
954- let first_split = owned_splits. next ( ) ?;
955- let mut row_range = first_split. start ..first_split. end ;
956- for split_range in owned_splits {
957- row_range. end = split_range. end ;
958- }
959-
960- Some ( row_range)
1005+ Some (
1006+ natural_splits. row_boundaries [ first_split] ..natural_splits. row_boundaries [ after_last_split] ,
1007+ )
9611008}
9621009
9631010fn split_assignment_byte (
0 commit comments