Precompute natural split assignment bytes

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-datafusion/src/persistent/opener.rs

@@ -37,7 +37,6 @@ 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;
@@ -98,8 +97,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<DashMap<Path, Weak<dyn LayoutReader>>>,
-    /// Shared full-file natural split ranges keyed by file path.
-    pub natural_split_ranges: Arc<DashMap<Path, Arc<[Range<u64>]>>>,
+    /// Shared full-file natural splits keyed by file path.
+    pub natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
     /// Whether the query has output ordering specified
     pub has_output_ordering: bool,
 
@@ -140,7 +139,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;
 
@@ -370,17 +369,16 @@ impl FileOpener for VortexOpener {
                 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(),
+                    let natural_splits = natural_splits_for_file(
+                        natural_splits.as_ref(),
                         &file.object_meta.location,
                         &layout_reader,
+                        file.object_meta.size,
                     )?;
 
-                    let Some(row_range) = split_aligned_row_range(
-                        byte_range,
-                        file.object_meta.size,
-                        natural_split_ranges.as_ref(),
-                    ) else {
+                    let Some(row_range) =
+                        split_aligned_row_range(byte_range, natural_splits.as_ref())
+                    else {
                         return Ok(stream::empty().boxed());
                     };
 
@@ -482,72 +480,108 @@ impl FileOpener for VortexOpener {
     }
 }
 
-fn natural_split_ranges_for_file(
-    natural_split_ranges: &DashMap<Path, Arc<[Range<u64>]>>,
+#[derive(Debug)]
+pub(crate) struct NaturalSplits {
+    row_boundaries: Arc<[u64]>,
+    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,
+        }
+    }
+}
+
+fn natural_splits_for_file(
+    natural_splits: &DashMap<Path, Arc<NaturalSplits>>,
     path: &Path,
     layout_reader: &Arc<dyn LayoutReader>,
-) -> DFResult<Arc<[Range<u64>]>> {
-    if let Some(split_ranges) = natural_split_ranges.get(path) {
-        return Ok(Arc::clone(split_ranges.value()));
+    total_size: u64,
+) -> DFResult<Arc<NaturalSplits>> {
+    if let Some(splits) = natural_splits.get(path) {
+        return Ok(Arc::clone(splits.value()));
     }
 
-    let split_ranges = compute_natural_split_ranges(layout_reader.as_ref())?;
-
-    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(layout_reader.as_ref(), total_size)?;
+            entry.insert(Arc::clone(&splits));
+            Ok(splits)
         }
     }
 }
 
-fn compute_natural_split_ranges(layout_reader: &dyn LayoutReader) -> DFResult<Arc<[Range<u64>]>> {
+fn compute_natural_splits(
+    layout_reader: &dyn LayoutReader,
+    total_size: u64,
+) -> DFResult<Arc<NaturalSplits>> {
     let row_count = layout_reader.row_count();
     let row_range = 0..row_count;
-    let split_points: Vec<_> = SplitBy::Layout
+    let row_boundaries = 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::<Vec<_>>();
+        .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?;
 
-    Ok(split_points.into())
+    Ok(Arc::new(NaturalSplits::new(
+        row_boundaries.into(),
+        total_size,
+    )))
 }
 
 /// Translate a DataFusion byte range to the contiguous natural split ranges it owns.
 /// Most splits are assigned by midpoint, but the leading split stays with the range that owns
 /// byte 0 so a tiny first byte range still claims the first rows.
 fn split_aligned_row_range(
     byte_range: Range<u64>,
-    total_size: u64,
-    split_ranges: &[Range<u64>],
+    natural_splits: &NaturalSplits,
 ) -> Option<Range<u64>> {
     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 +709,15 @@ mod tests {
         }
     }
 
+    fn natural_splits(total_size: u64, split_ranges: &[Range<u64>]) -> 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 +731,7 @@ mod tests {
         #[case] expected: Option<Range<u64>>,
     ) {
         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 +740,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::<Vec<_>>();
 
         assert_eq!(assigned, vec![0..4, 4..10, 10..13]);
@@ -731,6 +775,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<dyn ObjectStore>,
         path: &str,
@@ -767,7 +820,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 +1151,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 +1238,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 +1395,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 +1455,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 +1664,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,

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<DashMap<Path, Weak<dyn LayoutReader>>>,
-    /// Shared full-file natural split ranges keyed by path.
-    natural_split_ranges: Arc<DashMap<Path, Arc<[Range<u64>]>>>,
+    /// Shared full-file natural splits keyed by path.
+    natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
     expression_convertor: Arc<dyn ExpressionConvertor>,
     pub(crate) vortex_reader_factory: Option<Arc<dyn VortexReaderFactory>>,
     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(),

vortex-layout/src/scan/repeated_scan.rs

@@ -144,7 +144,7 @@ impl<A: 'static + Send> RepeatedScan<A> {
                         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)

vortex-layout/src/scan/scan_builder.rs

@@ -302,11 +302,11 @@ impl<A: 'static + Send> ScanBuilder<A> {
                     .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(

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;
 
@@ -19,7 +20,7 @@ pub enum Splits {
     /// column chunks).
     ///
     /// The vec is sorted in ascending order and deduplicated.
-    Natural(Vec<u64>),
+    Natural(Arc<[u64]>),
 
     /// Exact split ranges.
     ///