register_splits to get both offset and relative row_range

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

vortex-cuda/src/layout.rs

@@ -49,6 +49,7 @@ use vortex::layout::LayoutReader;
 use vortex::layout::LayoutReaderRef;
 use vortex::layout::LayoutRef;
 use vortex::layout::LayoutStrategy;
+use vortex::layout::SplitRange;
 use vortex::layout::VTable;
 use vortex::layout::layouts::SharedArrayFuture;
 use vortex::layout::segments::SegmentId;
@@ -284,10 +285,11 @@ impl LayoutReader for CudaFlatReader {
     fn register_splits(
         &self,
         _field_mask: &[FieldMask],
-        row_range: &Range<u64>,
+        split_range: &SplitRange,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        splits.insert(row_range.start + self.layout.row_count);
+        split_range.check_bounds(self.layout.row_count)?;
+        splits.insert(split_range.root_row_range().end);
         Ok(())
     }
 

vortex-file/public-api.lock

@@ -90,7 +90,7 @@ pub fn vortex_file::v2::FileStatsLayoutReader::projection_evaluation(&self, &cor
 
 pub fn vortex_file::v2::FileStatsLayoutReader::pruning_evaluation(&self, &core::ops::range::Range<u64>, &vortex_array::expr::expression::Expression, vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::mask_future::MaskFuture>
 
-pub fn vortex_file::v2::FileStatsLayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &core::ops::range::Range<u64>, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
+pub fn vortex_file::v2::FileStatsLayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &vortex_layout::reader::SplitRange, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
 
 pub fn vortex_file::v2::FileStatsLayoutReader::row_count(&self) -> u64
 

vortex-file/src/v2/file_stats_reader.rs

@@ -32,6 +32,7 @@ use vortex_error::VortexResult;
 use vortex_layout::ArrayFuture;
 use vortex_layout::LayoutReader;
 use vortex_layout::LayoutReaderRef;
+use vortex_layout::SplitRange;
 use vortex_mask::Mask;
 use vortex_session::VortexSession;
 use vortex_utils::aliases::dash_map::DashMap;
@@ -156,10 +157,10 @@ impl LayoutReader for FileStatsLayoutReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_range: &Range<u64>,
+        split_range: &SplitRange,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.child.register_splits(field_mask, row_range, splits)
+        self.child.register_splits(field_mask, split_range, splits)
     }
 
     fn pruning_evaluation(

vortex-layout/public-api.lock

@@ -636,7 +636,7 @@ pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::projection_evaluatio
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::pruning_evaluation(&self, &core::ops::range::Range<u64>, &vortex_array::expr::expression::Expression, vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::mask_future::MaskFuture>
 
-pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &core::ops::range::Range<u64>, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
+pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &vortex_layout::SplitRange, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::row_count(&self) -> u64
 
@@ -1712,6 +1712,42 @@ pub fn vortex_layout::LazyReaderChildren::get(&self, usize) -> vortex_error::Vor
 
 pub fn vortex_layout::LazyReaderChildren::new(alloc::sync::Arc<dyn vortex_layout::LayoutChildren>, alloc::vec::Vec<vortex_array::dtype::DType>, alloc::vec::Vec<alloc::sync::Arc<str>>, alloc::sync::Arc<dyn vortex_layout::segments::SegmentSource>, vortex_session::VortexSession) -> Self
 
+pub struct vortex_layout::SplitRange
+
+impl vortex_layout::SplitRange
+
+pub fn vortex_layout::SplitRange::check_bounds(&self, u64) -> vortex_error::VortexResult<()>
+
+pub fn vortex_layout::SplitRange::is_empty(&self) -> bool
+
+pub fn vortex_layout::SplitRange::len(&self) -> u64
+
+pub fn vortex_layout::SplitRange::root(core::ops::range::Range<u64>) -> vortex_error::VortexResult<Self>
+
+pub fn vortex_layout::SplitRange::root_row_range(&self) -> core::ops::range::Range<u64>
+
+pub fn vortex_layout::SplitRange::row_offset(&self) -> u64
+
+pub fn vortex_layout::SplitRange::row_range(&self) -> &core::ops::range::Range<u64>
+
+pub fn vortex_layout::SplitRange::try_new(u64, core::ops::range::Range<u64>) -> vortex_error::VortexResult<Self>
+
+impl core::clone::Clone for vortex_layout::SplitRange
+
+pub fn vortex_layout::SplitRange::clone(&self) -> vortex_layout::SplitRange
+
+impl core::cmp::Eq for vortex_layout::SplitRange
+
+impl core::cmp::PartialEq for vortex_layout::SplitRange
+
+pub fn vortex_layout::SplitRange::eq(&self, &vortex_layout::SplitRange) -> bool
+
+impl core::fmt::Debug for vortex_layout::SplitRange
+
+pub fn vortex_layout::SplitRange::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::marker::StructuralPartialEq for vortex_layout::SplitRange
+
 pub trait vortex_layout::ArrayFutureExt
 
 pub fn vortex_layout::ArrayFutureExt::masked(self, vortex_array::mask_future::MaskFuture) -> Self
@@ -1846,7 +1882,7 @@ pub fn vortex_layout::LayoutReader::projection_evaluation(&self, &core::ops::ran
 
 pub fn vortex_layout::LayoutReader::pruning_evaluation(&self, &core::ops::range::Range<u64>, &vortex_array::expr::expression::Expression, vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::mask_future::MaskFuture>
 
-pub fn vortex_layout::LayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &core::ops::range::Range<u64>, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
+pub fn vortex_layout::LayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &vortex_layout::SplitRange, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
 
 pub fn vortex_layout::LayoutReader::row_count(&self) -> u64
 
@@ -1864,7 +1900,7 @@ pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::projection_evaluatio
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::pruning_evaluation(&self, &core::ops::range::Range<u64>, &vortex_array::expr::expression::Expression, vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::mask_future::MaskFuture>
 
-pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &core::ops::range::Range<u64>, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
+pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::register_splits(&self, &[vortex_array::dtype::field_mask::FieldMask], &vortex_layout::SplitRange, &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::row_count(&self) -> u64
 

vortex-layout/src/layouts/chunked/reader.rs

@@ -10,7 +10,6 @@ use futures::FutureExt;
 use futures::TryStreamExt;
 use futures::future::BoxFuture;
 use futures::stream::FuturesOrdered;
-use itertools::Itertools;
 use tracing::trace;
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
@@ -30,6 +29,7 @@ use crate::LayoutReaderRef;
 use crate::LazyReaderChildren;
 use crate::layouts::chunked::ChunkedLayout;
 use crate::reader::LayoutReader;
+use crate::reader::SplitRange;
 use crate::segments::SegmentSource;
 
 /// A [`LayoutReader`] for chunked layouts.
@@ -107,10 +107,11 @@ impl ChunkedReader {
     fn ranges<'a>(
         &'a self,
         row_range: &'a Range<u64>,
-    ) -> impl Iterator<Item = (usize, Range<u64>, Range<usize>)> + 'a {
+    ) -> impl Iterator<Item = (usize, u64, Range<u64>, Range<usize>)> + 'a {
         self.chunk_range(row_range).map(move |chunk_idx| {
             // Figure out the chunk row range relative to the mask's row range.
             let chunk_row_range = self.chunk_offset(chunk_idx)..self.chunk_offset(chunk_idx + 1);
+            let chunk_start = chunk_row_range.start;
 
             // Find the intersection of the mask and the chunk row ranges.
             let intersecting_row_range =
@@ -146,7 +147,7 @@ impl ChunkedReader {
                 .vortex_expect("Chunk range calculation overflow");
             let chunk_range = chunk_relative_start..chunk_relative_end;
 
-            (chunk_idx, chunk_range, mask_range)
+            (chunk_idx, chunk_start, chunk_range, mask_range)
         })
     }
 }
@@ -167,37 +168,32 @@ impl LayoutReader for ChunkedReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_range: &Range<u64>,
+        split_range: &SplitRange,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        if row_range.is_empty() {
+        split_range.check_bounds(self.layout.row_count())?;
+
+        if split_range.is_empty() {
             return Ok(());
         }
 
-        for (index, (&start, &end)) in self
-            .chunk_offsets
-            .iter()
-            .tuple_windows::<(_, _)>()
-            .enumerate()
-        {
-            if end < row_range.start {
-                continue;
-            }
+        for (chunk_idx, chunk_start, child_range, _) in self.ranges(split_range.row_range()) {
+            let child = self.chunk_reader(chunk_idx)?;
+            let child_row_offset = split_range
+                .row_offset()
+                .checked_add(chunk_start)
+                .vortex_expect("Chunked layout split offset overflow");
+            let child_split_range = SplitRange::try_new(child_row_offset, child_range)?;
 
-            if start >= row_range.end {
-                break;
-            }
-
-            // Child overlaps in whole or in part with split
-            let child = self.chunk_reader(index)?;
-            let child_range =
-                std::cmp::max(row_range.start, start)..std::cmp::min(row_range.end, end);
-
-            // Register any splits from the child
-            child.register_splits(field_mask, &child_range, splits)?;
+            child.register_splits(field_mask, &child_split_range, splits)?;
 
             // Register the split indicating the end of this chunk
-            splits.insert(child_range.end);
+            splits.insert(
+                split_range
+                    .row_offset()
+                    .checked_add(chunk_start + child_split_range.row_range().end)
+                    .vortex_expect("Chunked layout split offset overflow"),
+            );
         }
 
         Ok(())
@@ -215,7 +211,7 @@ impl LayoutReader for ChunkedReader {
 
         let mut chunk_evals = vec![];
 
-        for (chunk_idx, chunk_range, mask_range) in self.ranges(row_range) {
+        for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) {
             let chunk_reader = self.chunk_reader(chunk_idx)?;
             let chunk_eval = chunk_reader
                 .pruning_evaluation(&chunk_range, expr, mask.slice(mask_range))
@@ -261,7 +257,7 @@ impl LayoutReader for ChunkedReader {
 
         let mut chunk_evals = vec![];
 
-        for (chunk_idx, chunk_range, mask_range) in self.ranges(row_range) {
+        for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) {
             let chunk_reader = self.chunk_reader(chunk_idx)?;
             let chunk_eval = chunk_reader
                 .filter_evaluation(&chunk_range, expr, mask.slice(mask_range))
@@ -301,7 +297,7 @@ impl LayoutReader for ChunkedReader {
 
         let mut chunk_evals = vec![];
 
-        for (chunk_idx, chunk_range, mask_range) in self.ranges(row_range) {
+        for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) {
             let chunk_reader = self.chunk_reader(chunk_idx)?;
             let chunk_eval = chunk_reader
                 .projection_evaluation(&chunk_range, expr, mask.slice(mask_range))
@@ -343,17 +339,25 @@ mod test {
     use vortex_array::MaskFuture;
     use vortex_array::assert_arrays_eq;
     use vortex_array::dtype::DType;
+    use vortex_array::dtype::FieldMask;
     use vortex_array::dtype::Nullability::NonNullable;
     use vortex_array::dtype::PType;
     use vortex_array::expr::root;
     use vortex_buffer::buffer;
     use vortex_io::runtime::single::block_on;
     use vortex_io::session::RuntimeSessionExt;
+    use vortex_session::registry::ReadContext;
 
+    use crate::IntoLayout;
     use crate::LayoutRef;
     use crate::LayoutStrategy;
+    use crate::OwnedLayoutChildren;
+    use crate::layouts::chunked::ChunkedLayout;
     use crate::layouts::chunked::writer::ChunkedLayoutStrategy;
+    use crate::layouts::flat::FlatLayout;
     use crate::layouts::flat::writer::FlatLayoutStrategy;
+    use crate::scan::split_by::SplitBy;
+    use crate::segments::SegmentId;
     use crate::segments::SegmentSource;
     use crate::segments::TestSegments;
     use crate::sequence::SequenceId;
@@ -395,6 +399,52 @@ mod test {
         (segments, layout)
     }
 
+    fn nested_chunked_layout() -> LayoutRef {
+        let dtype = DType::Primitive(PType::U8, NonNullable);
+        let ctx = ReadContext::new([]);
+        let flat = |segment_id| {
+            FlatLayout::new(5, dtype.clone(), SegmentId::from(segment_id), ctx.clone())
+                .into_layout()
+        };
+        let nested = |first_segment_id| {
+            ChunkedLayout::new(
+                10,
+                dtype.clone(),
+                OwnedLayoutChildren::layout_children(vec![
+                    flat(first_segment_id),
+                    flat(first_segment_id + 1),
+                ]),
+            )
+            .into_layout()
+        };
+
+        ChunkedLayout::new(
+            30,
+            dtype.clone(),
+            OwnedLayoutChildren::layout_children(vec![nested(0), nested(2), nested(4)]),
+        )
+        .into_layout()
+    }
+
+    #[rstest]
+    #[case(0..30, [0, 5, 10, 15, 20, 25, 30])]
+    #[case(7..23, [7, 10, 15, 20, 23])]
+    fn test_nested_chunked_layout_splits(
+        #[case] row_range: std::ops::Range<u64>,
+        #[case] expected: impl IntoIterator<Item = u64>,
+    ) {
+        let layout = nested_chunked_layout();
+        let reader = layout
+            .new_reader("".into(), Arc::new(TestSegments::default()), &SESSION)
+            .unwrap();
+
+        let splits = SplitBy::Layout
+            .splits(reader.as_ref(), &row_range, &[FieldMask::All])
+            .unwrap();
+
+        assert_eq!(splits, expected.into_iter().collect());
+    }
+
     #[rstest]
     fn test_chunked_evaluator(
         #[from(chunked_layout)] (segments, layout): (Arc<dyn SegmentSource>, LayoutRef),

vortex-layout/src/layouts/dict/reader.rs

@@ -32,6 +32,7 @@ use vortex_utils::aliases::dash_map::DashMap;
 use super::DictLayout;
 use crate::LayoutReader;
 use crate::LayoutReaderRef;
+use crate::SplitRange;
 use crate::layouts::SharedArrayFuture;
 use crate::segments::SegmentSource;
 
@@ -166,10 +167,10 @@ impl LayoutReader for DictReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_range: &Range<u64>,
+        split_range: &SplitRange,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.codes.register_splits(field_mask, row_range, splits)
+        self.codes.register_splits(field_mask, split_range, splits)
     }
 
     fn pruning_evaluation(

vortex-layout/src/layouts/flat/reader.rs

@@ -21,9 +21,10 @@ use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_session::VortexSession;
 
-use crate::LayoutReader;
 use crate::layouts::SharedArrayFuture;
 use crate::layouts::flat::FlatLayout;
+use crate::reader::LayoutReader;
+use crate::reader::SplitRange;
 use crate::segments::SegmentSource;
 
 /// The threshold of mask density below which we will evaluate the expression only over the
@@ -103,10 +104,11 @@ impl LayoutReader for FlatReader {
     fn register_splits(
         &self,
         _field_mask: &[FieldMask],
-        row_range: &Range<u64>,
+        split_range: &SplitRange,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        splits.insert(row_range.start + self.layout.row_count);
+        split_range.check_bounds(self.layout.row_count)?;
+        splits.insert(split_range.root_row_range().end);
         Ok(())
     }
 

vortex-layout/src/layouts/row_idx/mod.rs

@@ -43,6 +43,7 @@ use vortex_utils::aliases::dash_map::DashMap;
 
 use crate::ArrayFuture;
 use crate::LayoutReader;
+use crate::SplitRange;
 use crate::layouts::partitioned::PartitionedExprEval;
 
 pub struct RowIdxLayoutReader {
@@ -170,10 +171,10 @@ impl LayoutReader for RowIdxLayoutReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_range: &Range<u64>,
+        split_range: &SplitRange,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.child.register_splits(field_mask, row_range, splits)
+        self.child.register_splits(field_mask, split_range, splits)
     }
 
     fn pruning_evaluation(