Skip to content

Commit 4b2a246

Browse files
AdamGSclaude
andcommitted
Reuse precomputed natural splits in the Morselizer
The Morselizer carried a copy of the natural-split helpers from the opener it replaced, which predated the precomputed-assignment form on the base branch. Port that forward: cache a NaturalSplits (row boundaries plus each split's assignment byte) instead of raw ranges, translate a byte range with two binary searches instead of scanning every split, and hand the boundaries to the scan so prepare() does not re-walk the layout per partition. The async OnceCell cache the Morselizer introduced is kept. Signed-off-by: Adam Gutglick <adam@spiraldb.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 651f117 commit 4b2a246

3 files changed

Lines changed: 130 additions & 62 deletions

File tree

‎vortex-datafusion/src/persistent/morsel/mod.rs‎

Lines changed: 95 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@ use datafusion_pruning::FilePruner;
4444
use futures::StreamExt;
4545
use futures::TryStreamExt;
4646
use futures::stream::BoxStream;
47-
use itertools::Itertools;
4847
use object_store::path::Path;
4948
use tokio::sync::OnceCell;
5049
use tracing::Instrument;
@@ -79,8 +78,52 @@ use crate::persistent::cache::CachedVortexMetadata;
7978
use crate::persistent::stream::PrunableStream;
8079
use 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.
86129
pub 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

901954
fn 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.
932987
fn 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

9631010
fn split_assignment_byte(

0 commit comments

Comments
 (0)