only register_splits for active row_range in partitions

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

bench-vortex/src/realnest/gharchive.rs

@@ -230,6 +230,10 @@ impl Benchmark for GithubArchive {
     fn data_url(&self) -> &Url {
         &self.data_url
     }
+
+    fn expected_row_counts(&self) -> Option<&[usize]> {
+        Some(&[1, 2, 100, 10, 82468])
+    }
 }
 
 pub async fn register_table(

vortex-file/src/file.rs

@@ -137,8 +137,9 @@ impl VortexFile {
     }
 
     pub fn splits(&self) -> VortexResult<Vec<Range<u64>>> {
+        let reader = self.layout_reader()?;
         Ok(SplitBy::Layout
-            .splits(self.layout_reader()?.as_ref(), &[FieldMask::All])?
+            .splits(reader.as_ref(), &(0..reader.row_count()), &[FieldMask::All])?
             .into_iter()
             .tuple_windows()
             .map(|(start, end)| start..end)

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

@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use futures::future::BoxFuture;
 use futures::stream::FuturesOrdered;
-use futures::{FutureExt, StreamExt, TryStreamExt};
+use futures::{FutureExt, TryStreamExt};
 use vortex_array::arrays::ChunkedArray;
 use vortex_array::{ArrayRef, MaskFuture};
 use vortex_dtype::{DType, FieldMask};
@@ -149,16 +149,23 @@ impl LayoutReader for ChunkedReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        let mut offset = row_offset;
+        let mut offset = row_range.start;
         for i in 0..self.layout.nchildren() {
+            // Bail early if the row range only overlaps with a subset of the chunks
+            if offset >= row_range.end {
+                break;
+            }
+
             let child = self.chunk_reader(i)?;
-            child.register_splits(field_mask, offset, splits)?;
-            offset += self.layout.child(i)?.row_count();
+            let child_range = offset..offset + child.row_count();
+            child.register_splits(field_mask, &child_range, splits)?;
+            offset = child_range.end;
             splits.insert(offset);
         }
+
         Ok(())
     }
 

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

@@ -111,10 +111,10 @@ impl LayoutReader for DictReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.codes.register_splits(field_mask, row_offset, splits)
+        self.codes.register_splits(field_mask, row_range, splits)
     }
 
     fn pruning_evaluation(

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

@@ -84,10 +84,10 @@ impl LayoutReader for FlatReader {
     fn register_splits(
         &self,
         _field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        splits.insert(row_offset + self.layout.row_count());
+        splits.insert(row_range.end);
         Ok(())
     }
 

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

@@ -136,10 +136,10 @@ impl LayoutReader for RowIdxLayoutReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.child.register_splits(field_mask, row_offset, splits)
+        self.child.register_splits(field_mask, row_range, splits)
     }
 
     fn pruning_evaluation(

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

@@ -8,7 +8,6 @@ use std::sync::Arc;
 use futures::try_join;
 use itertools::Itertools;
 use vortex_array::arrays::StructArray;
-use vortex_array::stats::Precision;
 use vortex_array::vtable::ValidityHelper;
 use vortex_array::{ArrayRef, IntoArray, MaskFuture, ToCanonical};
 use vortex_dtype::{DType, FieldMask, FieldName, Nullability, StructFields};
@@ -116,11 +115,6 @@ impl StructReader {
             idx
         };
 
-        let field_dtype = self
-            .struct_fields()
-            .field_by_index(idx)
-            .ok_or_else(|| vortex_err!("Missing field {idx}"))?;
-        let name = &self.struct_fields().names()[idx];
         self.lazy_children.get(child_index)
     }
 
@@ -210,15 +204,15 @@ impl LayoutReader for StructReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
         // In the case of an empty struct, we need to register the end split.
-        splits.insert(row_offset + self.layout.row_count);
+        splits.insert(row_range.end);
 
         self.layout.matching_fields(field_mask, |mask, idx| {
             self.field_reader_by_index(idx)?
-                .register_splits(&[mask], row_offset, splits)
+                .register_splits(&[mask], row_range, splits)
         })
     }
 

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

@@ -194,11 +194,11 @@ impl LayoutReader for ZonedReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
         self.data_child()?
-            .register_splits(field_mask, row_offset, splits)
+            .register_splits(field_mask, row_range, splits)
     }
 
     fn pruning_evaluation(

vortex-layout/src/reader.rs

@@ -36,7 +36,7 @@ pub trait LayoutReader: 'static + Send + Sync {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()>;
 
@@ -104,7 +104,6 @@ pub struct LazyReaderChildren {
     dtypes: Vec<DType>,
     names: Vec<Arc<str>>,
     segment_source: Arc<dyn SegmentSource>,
-
     // TODO(ngates): we may want a hash map of some sort here?
     cache: Vec<OnceCell<LayoutReaderRef>>,
 }

vortex-scan/src/scan_builder.rs

@@ -237,7 +237,15 @@ impl<A: 'static + Send> ScanBuilder<A> {
             if let Some(ranges) = attempt_split_ranges(&self.selection, self.row_range.as_ref()) {
                 Splits::Ranges(ranges)
             } else {
-                Splits::Natural(self.split_by.splits(layout_reader.as_ref(), &field_mask)?)
+                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,
+                )?)
             };
 
         Ok(RepeatedScan::new(