feat: VortexDataset.get_fragments yields one fragment per split (#4386)

Signed-off-by: Daniel King <[email protected]>

Author: danking

vortex-file/src/file.rs

@@ -6,18 +6,19 @@
 //! The `VortexFile` provides methods for accessing file metadata, creating segment sources for reading
 //! data from the file, and initiating scans to read the file's contents into memory as Vortex arrays.
 
+use std::ops::Range;
 use std::sync::Arc;
 
 use vortex_array::ArrayRef;
 use vortex_array::stats::StatsSet;
-use vortex_dtype::{DType, Field, FieldPath, FieldPathSet};
+use vortex_dtype::{DType, Field, FieldMask, FieldPath, FieldPathSet};
 use vortex_error::VortexResult;
 use vortex_expr::pruning::checked_pruning_expr;
 use vortex_expr::{ExprRef, Scope};
 use vortex_layout::LayoutReader;
 use vortex_layout::segments::SegmentSource;
 use vortex_metrics::VortexMetrics;
-use vortex_scan::ScanBuilder;
+use vortex_scan::{ScanBuilder, SplitBy};
 use vortex_utils::aliases::hash_map::HashMap;
 
 use crate::footer::Footer;
@@ -133,4 +134,8 @@ impl VortexFile {
             .as_constant()
             .is_some_and(|result| result.as_bool().value() == Some(true)))
     }
+
+    pub fn splits(&self) -> VortexResult<Vec<Range<u64>>> {
+        SplitBy::Layout.splits(self.layout_reader()?.as_ref(), &[FieldMask::All])
+    }
 }

vortex-python/python/vortex/_lib/dataset.pyi

@@ -8,16 +8,27 @@ from .expr import Expr
 
 class VortexDataset:
     def to_array(
-        self, columns: list[str] | list[int] | None = None, row_filter: Expr | None = None, indices: Array | None = None
+        self,
+        columns: list[str] | list[int] | None = None,
+        row_filter: Expr | None = None,
+        indices: Array | None = None,
+        row_range: tuple[int, int] | None = None,
     ) -> Array: ...
     def to_record_batch_reader(
         self,
         columns: list[str] | list[int] | None = None,
         row_filter: Expr | None = None,
         indices: Array | None = None,
         split_by: int | None = None,
+        row_range: tuple[int, int] | None = None,
     ) -> pyarrow.RecordBatchReader: ...
-    def count_rows(self, row_filter: Expr | None = None, split_by: int | None = None) -> int: ...
+    def count_rows(
+        self,
+        row_filter: Expr | None = None,
+        split_by: int | None = None,
+        row_range: tuple[int, int] | None = None,
+    ) -> int: ...
     def schema(self) -> pyarrow.Schema: ...
+    def splits(self) -> list[tuple[int, int]]: ...
 
 def dataset_from_url(url: str) -> VortexDataset: ...

vortex-python/python/vortex/_lib/file.pyi

@@ -45,5 +45,6 @@ class VortexFile:
     ) -> pa.RecordBatchReader: ...
     def to_dataset(self) -> VortexDataset: ...
     def to_polars(self) -> pl.LazyFrame: ...
+    def splits(self) -> list[tuple[int, int]]: ...
 
 def open(path: str, *, without_segment_cache: bool = False) -> VortexFile: ...

vortex-python/python/vortex/arrays.py

@@ -41,7 +41,10 @@ def __ArrowDtype_type_patched(self: pandas.ArrowDtype):
 
 
 def empty_arrow_table(schema: pyarrow.Schema) -> pyarrow.Table:
-    return pyarrow.Table.from_arrays([pyarrow.array([], type=t) for t in schema], schema=schema)  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
+    def empty_array(f: pyarrow.Field[pyarrow.DataType]) -> pyarrow.Array[pyarrow.Scalar[pyarrow.DataType]]:
+        return pyarrow.array([], type=f.type)
+
+    return pyarrow.Table.from_arrays([empty_array(field) for field in schema], schema=schema)  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
 
 
 def arrow_table_from_struct_array(