diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index d0ffb472ebe..9c28a5622e8 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -37,18 +37,15 @@ use futures::FutureExt; use futures::StreamExt; use futures::TryStreamExt; use futures::stream; -use itertools::Itertools; use object_store::path::Path; use tracing::Instrument; use vortex::array::VortexSessionExecute; -use vortex::dtype::FieldMask; use vortex::error::VortexError; use vortex::error::VortexExpect; use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; -use vortex::layout::scan::split_by::SplitBy; use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; @@ -98,8 +95,8 @@ pub(crate) struct VortexOpener { /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache /// a file reader the first time we read a file. pub layout_readers: Arc>>, - /// Shared full-file natural split ranges keyed by file path. - pub natural_split_ranges: Arc]>>>, + /// Shared full-file natural splits keyed by file path. + pub natural_splits: Arc>>, /// Whether the query has output ordering specified pub has_output_ordering: bool, @@ -140,7 +137,7 @@ impl FileOpener for VortexOpener { let unified_file_schema = Arc::clone(self.table_schema.file_schema()); let limit = self.limit; let layout_readers = Arc::clone(&self.layout_readers); - let natural_split_ranges = Arc::clone(&self.natural_split_ranges); + let natural_splits = Arc::clone(&self.natural_splits); let has_output_ordering = self.has_output_ordering; let scan_concurrency = self.scan_concurrency; @@ -360,34 +357,6 @@ impl FileOpener for VortexOpener { scan_builder = vortex_plan.apply_to_builder(scan_builder); } - if let Some(file_range) = file.range { - let byte_range = Range { - start: u64::try_from(file_range.start) - .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, - end: u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, - }; - if byte_range.start != 0 || byte_range.end != file.object_meta.size { - // Full-file scans already cover every natural split. Only translate the - // byte range back into row boundaries when DataFusion has trimmed the file. - let natural_split_ranges = natural_split_ranges_for_file( - natural_split_ranges.as_ref(), - &file.object_meta.location, - &layout_reader, - )?; - - let Some(row_range) = split_aligned_row_range( - byte_range, - file.object_meta.size, - natural_split_ranges.as_ref(), - ) else { - return Ok(stream::empty().boxed()); - }; - - scan_builder = scan_builder.with_row_range(row_range); - } - } - let filter = filter .and_then(|f| { // Verify that all filters we've accepted from DataFusion get pushed down. @@ -430,11 +399,46 @@ impl FileOpener for VortexOpener { scan_builder = scan_builder.with_concurrency(concurrency); } + // Set before the byte-range translation below, which computes natural splits for + // the fields the scan's projection and filter reference. + scan_builder = scan_builder + .with_projection(scan_projection) + .with_some_filter(filter); + + if let Some(file_range) = file.range { + let byte_range = Range { + start: u64::try_from(file_range.start) + .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, + end: u64::try_from(file_range.end) + .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + // Full-file scans already cover every natural split. Only translate the + // byte range back into row boundaries when DataFusion has trimmed the file. + let natural_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + &scan_builder, + file.object_meta.size, + )?; + + let Some(row_range) = + split_aligned_row_range(byte_range, natural_splits.as_ref()) + else { + return Ok(stream::empty().boxed()); + }; + + scan_builder = scan_builder + .with_row_range(row_range) + // Hand the shared full-file boundaries back to the scan so prepare() + // skips its own layout walk. + .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + } + } + let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); let stream = scan_builder .with_metrics_registry(metrics_registry) - .with_projection(scan_projection) - .with_some_filter(filter) .with_ordered(has_output_ordering) .map(move |chunk| { let mut ctx = session.create_execution_ctx(); @@ -482,38 +486,98 @@ impl FileOpener for VortexOpener { } } -fn natural_split_ranges_for_file( - natural_split_ranges: &DashMap]>>, - path: &Path, - layout_reader: &Arc, -) -> DFResult]>> { - if let Some(split_ranges) = natural_split_ranges.get(path) { - return Ok(Arc::clone(split_ranges.value())); +/// A file's natural split boundaries plus the precomputed byte each split is assigned to, +/// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row +/// boundaries with a binary search instead of re-projecting every split per partition. +/// +/// The boundaries are computed for the fields referenced by the scan's projection and filter. +/// All partitions translate through the first opener's cached entry (the cache lives on the +/// source, so projection and filter are fixed for its lifetime), which keeps the byte ranges +/// tiling the file's rows exactly once. +#[derive(Debug)] +pub(crate) struct NaturalSplits { + /// Sorted row boundaries of the natural splits; split `i` covers + /// `row_boundaries[i]..row_boundaries[i + 1]`. Shared so partitions can hand the + /// boundaries back to the scan via [`ScanBuilder::with_natural_splits`], skipping the + /// per-partition layout walk in `prepare`. + row_boundaries: Arc<[u64]>, + /// For each split, the byte a DataFusion byte range must contain to own it (see + /// [`split_assignment_byte`]); one entry per split, sorted because split midpoints + /// increase monotonically under the row-to-byte projection. + assignment_bytes: Box<[u64]>, +} + +impl NaturalSplits { + fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self { + let row_count = row_boundaries.last().copied().unwrap_or_default(); + let assignment_bytes = if row_count == 0 { + Box::default() + } else { + row_boundaries + .windows(2) + .enumerate() + .map(|(idx, boundaries)| { + split_assignment_byte( + idx, + &(boundaries[0]..boundaries[1]), + row_count, + total_size, + ) + }) + .collect() + }; + + debug_assert!(assignment_bytes.is_sorted()); + debug_assert_eq!( + assignment_bytes.len() + usize::from(!row_boundaries.is_empty()), + row_boundaries.len() + ); + + Self { + row_boundaries, + assignment_bytes, + } } +} - let split_ranges = compute_natural_split_ranges(layout_reader.as_ref())?; +/// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. +fn natural_splits_for_file( + natural_splits: &DashMap>, + path: &Path, + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + if let Some(splits) = natural_splits.get(path) { + return Ok(Arc::clone(splits.value())); + } - match natural_split_ranges.entry(path.clone()) { + // Compute while holding the entry so concurrent partitions opening the same file wait + // for the winner instead of all walking the layout tree; the redundant walks contend on + // the lazily-initialized layout children and dominate the cost of the computation itself. + match natural_splits.entry(path.clone()) { Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), Entry::Vacant(entry) => { - entry.insert(Arc::clone(&split_ranges)); - Ok(split_ranges) + let splits = compute_natural_splits(scan_builder, total_size)?; + entry.insert(Arc::clone(&splits)); + Ok(splits) } } } -fn compute_natural_split_ranges(layout_reader: &dyn LayoutReader) -> DFResult]>> { - let row_count = layout_reader.row_count(); - let row_range = 0..row_count; - let split_points: Vec<_> = SplitBy::Layout - .splits(layout_reader, &row_range, &[FieldMask::All]) - .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))? - .into_iter() - .tuple_windows() - .map(|(s, e)| s..e) - .collect::>(); - - Ok(split_points.into()) +/// Walk the layout tree to compute the file's full natural split boundaries for the fields +/// referenced by the scan's projection and filter. +fn compute_natural_splits( + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + let row_boundaries = scan_builder + .full_file_splits() + .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; + + Ok(Arc::new(NaturalSplits::new( + row_boundaries.into(), + total_size, + ))) } /// Translate a DataFusion byte range to the contiguous natural split ranges it owns. @@ -521,33 +585,25 @@ fn compute_natural_split_ranges(layout_reader: &dyn LayoutReader) -> DFResult, - total_size: u64, - split_ranges: &[Range], + natural_splits: &NaturalSplits, ) -> Option> { if byte_range.start >= byte_range.end { return None; } - let row_count = split_ranges.last().map(|split| split.end)?; - if row_count == 0 { + let first_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.start); + let after_last_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.end); + if first_split == after_last_split { return None; } - let mut owned_splits = split_ranges - .iter() - .enumerate() - .filter_map(|(idx, split_range)| { - let assignment_byte = split_assignment_byte(idx, split_range, row_count, total_size); - byte_range.contains(&assignment_byte).then_some(split_range) - }); - - let first_split = owned_splits.next()?; - let mut row_range = first_split.start..first_split.end; - for split_range in owned_splits { - row_range.end = split_range.end; - } - - Some(row_range) + Some( + natural_splits.row_boundaries[first_split]..natural_splits.row_boundaries[after_last_split], + ) } fn split_assignment_byte( @@ -675,6 +731,15 @@ mod tests { } } + fn natural_splits(total_size: u64, split_ranges: &[Range]) -> NaturalSplits { + let mut row_boundaries = Vec::with_capacity(split_ranges.len() + 1); + if let Some(first) = split_ranges.first() { + row_boundaries.push(first.start); + row_boundaries.extend(split_ranges.iter().map(|range| range.end)); + } + NaturalSplits::new(row_boundaries.into(), total_size) + } + #[rstest] #[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] #[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] @@ -688,7 +753,7 @@ mod tests { #[case] expected: Option>, ) { assert_eq!( - split_aligned_row_range(byte_range, total_size, &split_ranges), + split_aligned_row_range(byte_range, &natural_splits(total_size, &split_ranges)), expected ); } @@ -697,10 +762,11 @@ mod tests { fn test_split_aligned_ranges_cover_splits_exactly_once() { let split_ranges = vec![0..1, 1..4, 4..10, 10..13]; let byte_ranges = [0..4, 4..8, 8..12, 12..16]; + let natural_splits = natural_splits(16, &split_ranges); let assigned = byte_ranges .into_iter() - .filter_map(|byte_range| split_aligned_row_range(byte_range, 16, &split_ranges)) + .filter_map(|byte_range| split_aligned_row_range(byte_range, &natural_splits)) .collect::>(); assert_eq!(assigned, vec![0..4, 4..10, 10..13]); @@ -731,6 +797,15 @@ mod tests { } } + #[test] + fn test_split_aligned_row_range_keeps_colliding_assignments_together() { + let natural_splits = natural_splits(2, &[0..1, 1..2, 2..3, 3..4]); + + assert_eq!(natural_splits.assignment_bytes.as_ref(), [0, 0, 1, 1]); + assert_eq!(split_aligned_row_range(0..1, &natural_splits), Some(0..2)); + assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4)); + } + async fn write_arrow_to_vortex( object_store: Arc, path: &str, @@ -767,7 +842,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1098,7 +1173,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1185,7 +1260,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1342,7 +1417,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1402,7 +1477,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, @@ -1611,7 +1686,7 @@ mod tests { metrics_registry: Arc::new(DefaultMetricsRegistry::default()), df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), - natural_split_ranges: Default::default(), + natural_splits: Default::default(), has_output_ordering: false, expression_convertor: Arc::new(DefaultExpressionConvertor::default()), file_metadata_cache: None, diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 34746c0ea50..81fcf937955 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt::Formatter; -use std::ops::Range; use std::sync::Arc; use std::sync::Weak; @@ -36,6 +35,7 @@ use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; +use super::opener::NaturalSplits; use super::opener::VortexOpener; use crate::VortexTableOptions; use crate::convert::exprs::DefaultExpressionConvertor; @@ -195,8 +195,8 @@ pub struct VortexSource { /// /// Sharing the readers allows us to only read every layout once from the file, even across partitions. layout_readers: Arc>>, - /// Shared full-file natural split ranges keyed by path. - natural_split_ranges: Arc]>>>, + /// Shared full-file natural splits keyed by path. + natural_splits: Arc>>, expression_convertor: Arc, pub(crate) vortex_reader_factory: Option>, pub(crate) ordered: bool, @@ -229,7 +229,7 @@ impl VortexSource { vortex_predicate: None, df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), - natural_split_ranges: Arc::new(DashMap::default()), + natural_splits: Arc::new(DashMap::default()), expression_convertor, vortex_reader_factory: None, vx_metrics_registry: Arc::new(DefaultMetricsRegistry::default()), @@ -356,7 +356,7 @@ impl VortexSource { metrics_registry: Arc::clone(&self.vx_metrics_registry), df_metrics: self.df_metrics.clone(), layout_readers: Arc::clone(&self.layout_readers), - natural_split_ranges: Arc::clone(&self.natural_split_ranges), + natural_splits: Arc::clone(&self.natural_splits), has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered, expression_convertor: Arc::clone(&self.expression_convertor), file_metadata_cache: self.file_metadata_cache.clone(), diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 10e6c0a1b18..413761b8103 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -144,7 +144,7 @@ impl RepeatedScan { if range.is_empty() { return Ok(Vec::new()); } - let lo = vec.partition_point(|&x| x < range.start); + let lo = vec.partition_point(|&x| x <= range.start); let hi = vec.partition_point(|&x| x < range.end); Either::Right( iter::once(range.start) diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 16a8758c16d..bc077586073 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -71,6 +71,9 @@ pub struct ScanBuilder { selection: Selection, /// How to split the file for concurrent processing. split_by: SplitBy, + /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare) + /// uses them instead of walking the layout. + natural_splits: Option>, /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. @@ -97,6 +100,7 @@ impl ScanBuilder { row_range: None, selection: Default::default(), split_by: SplitBy::Layout, + natural_splits: None, // We default to four tasks per worker thread, which allows for some I/O lookahead // without too much impact on work-stealing. concurrency: 4, @@ -191,6 +195,40 @@ impl ScanBuilder { self } + /// Supply precomputed full-file natural split boundaries (see + /// [`full_file_splits`](Self::full_file_splits)) so [`prepare`](Self::prepare) reuses them + /// instead of walking the layout. Callers translating external partitions into row ranges + /// can compute the boundaries once per file and share them across partitions. + /// + /// Takes precedence over [`with_split_by`](Self::with_split_by); boundaries outside the + /// scan's row range are clamped during execution. + pub fn with_natural_splits(mut self, boundaries: Arc<[u64]>) -> Self { + self.natural_splits = Some(boundaries); + self + } + + /// Compute the full-file natural split boundaries for the fields referenced by this scan's + /// projection and filter, ignoring any configured row range. + /// + /// These are the boundaries [`prepare`](Self::prepare) derives for a whole-file scan; hand + /// them back via [`with_natural_splits`](Self::with_natural_splits) to skip the layout walk + /// in `prepare`. + pub fn full_file_splits(&self) -> VortexResult> { + let dtype = self.layout_reader.dtype(); + let bound_projection = self.projection.optimize_recursive(dtype)?.bind(dtype)?; + let bound_filter = self + .filter + .as_ref() + .map(|f| f.optimize_recursive(dtype)?.bind(dtype)) + .transpose()?; + let field_mask = referenced_field_masks(&bound_projection, bound_filter.as_ref())?; + self.split_by.splits( + self.layout_reader.as_ref(), + &(0..self.layout_reader.row_count()), + &field_mask, + ) + } + /// Returns the per-worker row-split concurrency. pub fn concurrency(&self) -> usize { self.concurrency @@ -253,6 +291,7 @@ impl ScanBuilder { row_range: self.row_range, selection: self.selection, split_by: self.split_by, + natural_splits: self.natural_splits, concurrency: self.concurrency, metrics_registry: self.metrics_registry, file_stats: self.file_stats, @@ -297,22 +336,24 @@ impl ScanBuilder { .map(|expr| expr.bind(layout_reader.dtype())) .transpose()?; - // Construct field masks and compute the row splits of the scan. - let field_mask = referenced_field_masks(&bound_projection, bound_filter.as_ref())?; - + // Compute the row splits of the scan. let splits = if let Some(ranges) = attempt_split_ranges(&self.selection, self.row_range.as_ref()) { Splits::Ranges(ranges) + } else if let Some(boundaries) = self.natural_splits { + // Caller-supplied full-file boundaries; execution clamps them to the row range. + Splits::Natural(boundaries) } else { + let field_mask = referenced_field_masks(&bound_projection, bound_filter.as_ref())?; let split_range = self .row_range .clone() .unwrap_or_else(|| 0..layout_reader.row_count()); - Splits::Natural(self.split_by.splits( - layout_reader.as_ref(), - &split_range, - &field_mask, - )?) + Splits::Natural( + self.split_by + .splits(layout_reader.as_ref(), &split_range, &field_mask)? + .into(), + ) }; Ok(RepeatedScan::new( @@ -773,6 +814,47 @@ mod test { Ok(()) } + #[test] + fn supplied_natural_splits_skip_layout_walk() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let calls = Arc::new(AtomicUsize::new(0)); + let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls))); + + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + + let stream = ScanBuilder::new(session, reader) + .with_natural_splits(vec![0u64, 2, 4].into()) + .with_row_range(1..4) + .into_stream()?; + let mut iter = runtime.block_on_stream(stream); + + let mut chunks = Vec::new(); + for chunk in &mut iter { + let prim = chunk?.execute::(&mut ctx)?; + chunks.push(prim.into_buffer::().to_vec()); + } + + assert_eq!(calls.load(Ordering::Relaxed), 0); + // Supplied full-file boundaries [0, 2, 4] clamped to rows 1..4. + assert_eq!(chunks, [vec![1], vec![2, 3]]); + + Ok(()) + } + + #[test] + fn full_file_splits_ignore_row_range() -> VortexResult<()> { + let calls = Arc::new(AtomicUsize::new(0)); + let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls))); + + let splits = ScanBuilder::new(SCAN_SESSION.clone(), reader) + .with_row_range(1..3) + .full_file_splits()?; + + assert_eq!(splits, [0, 1, 2, 3, 4]); + Ok(()) + } + #[derive(Debug)] struct BlockingSplitsLayoutReader { name: Arc, diff --git a/vortex-layout/src/scan/splits.rs b/vortex-layout/src/scan/splits.rs index cf5fa7a0759..a5dcdea8868 100644 --- a/vortex-layout/src/scan/splits.rs +++ b/vortex-layout/src/scan/splits.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::Range; +use std::sync::Arc; use vortex_scan::selection::Selection; @@ -18,8 +19,8 @@ pub enum Splits { /// Natural splits computed by the layout reader (e.g., computing splits across different-sized /// column chunks). /// - /// The vec is sorted in ascending order and deduplicated. - Natural(Vec), + /// The boundaries are sorted in ascending order and deduplicated. + Natural(Arc<[u64]>), /// Exact split ranges. ///