Exact layout reader row count (#5088)

This was an unused idea from a while back. Putting them back to exact.

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-layout/src/layout.rs

@@ -77,11 +77,6 @@ pub enum LayoutChildType {
     /// A layout child that represents a single field of data.
     /// Contains the field name of the child.
     Field(FieldName),
-    // A layout child that contains a subset of the fields of the parent layout.
-    // Contains a mask over the fields of the parent layout.
-    // TODO(ngates): FieldMask API needs fixing before we enable this. We also don't yet have a
-    //  use-case for this.
-    // Mask(Vec<FieldMask>),
 }
 
 impl LayoutChildType {

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

@@ -9,7 +9,6 @@ use futures::future::BoxFuture;
 use futures::stream::FuturesOrdered;
 use futures::{FutureExt, TryStreamExt};
 use vortex_array::arrays::ChunkedArray;
-use vortex_array::stats::Precision;
 use vortex_array::{ArrayRef, MaskFuture};
 use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexExpect, VortexResult, vortex_panic};
@@ -142,8 +141,8 @@ impl LayoutReader for ChunkedReader {
         self.layout.dtype()
     }
 
-    fn row_count(&self) -> Precision<u64> {
-        Precision::Exact(self.layout.row_count())
+    fn row_count(&self) -> u64 {
+        self.layout.row_count()
     }
 
     fn register_splits(

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

@@ -8,7 +8,6 @@ use std::sync::{Arc, OnceLock};
 use futures::future::BoxFuture;
 use futures::{FutureExt, TryFutureExt, try_join};
 use vortex_array::compute::{MinMaxResult, min_max, take};
-use vortex_array::stats::Precision;
 use vortex_array::{ArrayRef, MaskFuture};
 use vortex_dict::DictArray;
 use vortex_dtype::{DType, FieldMask};
@@ -105,8 +104,8 @@ impl LayoutReader for DictReader {
         self.layout.dtype()
     }
 
-    fn row_count(&self) -> Precision<u64> {
-        Precision::Exact(self.layout.row_count())
+    fn row_count(&self) -> u64 {
+        self.layout.row_count()
     }
 
     fn register_splits(

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

@@ -9,7 +9,6 @@ use futures::FutureExt;
 use futures::future::BoxFuture;
 use vortex_array::compute::filter;
 use vortex_array::serde::ArrayParts;
-use vortex_array::stats::Precision;
 use vortex_array::{Array, ArrayRef, MaskFuture};
 use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexExpect, VortexResult, VortexUnwrap as _};
@@ -78,8 +77,8 @@ impl LayoutReader for FlatReader {
         self.layout.dtype()
     }
 
-    fn row_count(&self) -> Precision<u64> {
-        Precision::Exact(self.layout.row_count())
+    fn row_count(&self) -> u64 {
+        self.layout.row_count()
     }
 
     fn register_splits(

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

@@ -13,7 +13,6 @@ pub use expr::*;
 use futures::FutureExt;
 use futures::future::BoxFuture;
 use vortex_array::compute::filter;
-use vortex_array::stats::Precision;
 use vortex_array::{ArrayRef, IntoArray, MaskFuture};
 use vortex_dtype::{DType, FieldMask, FieldName, Nullability, PType};
 use vortex_error::{VortexExpect, VortexResult};
@@ -130,7 +129,7 @@ impl LayoutReader for RowIdxLayoutReader {
         self.child.dtype()
     }
 
-    fn row_count(&self) -> Precision<u64> {
+    fn row_count(&self) -> u64 {
         self.child.row_count()
     }
 

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

@@ -7,7 +7,6 @@ use std::sync::Arc;
 
 use itertools::Itertools;
 use vortex_array::MaskFuture;
-use vortex_array::stats::Precision;
 use vortex_dtype::{DType, FieldMask, FieldName, StructFields};
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 use vortex_expr::transform::immediate_access::annotate_scope_access;
@@ -169,8 +168,8 @@ impl LayoutReader for StructReader {
         self.layout.dtype()
     }
 
-    fn row_count(&self) -> Precision<u64> {
-        Precision::Exact(self.layout.row_count())
+    fn row_count(&self) -> u64 {
+        self.layout.row_count()
     }
 
     fn register_splits(

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

@@ -9,7 +9,6 @@ use futures::future::{BoxFuture, Shared};
 use futures::{FutureExt, TryFutureExt};
 use itertools::Itertools;
 use parking_lot::RwLock;
-use vortex_array::stats::Precision;
 use vortex_array::{ArrayRef, MaskFuture, ToCanonical};
 use vortex_buffer::BitBufferMut;
 use vortex_dtype::{DType, FieldMask, FieldPath, FieldPathSet};
@@ -186,8 +185,8 @@ impl LayoutReader for ZonedReader {
         self.layout.dtype()
     }
 
-    fn row_count(&self) -> Precision<u64> {
-        Precision::exact(self.layout.row_count())
+    fn row_count(&self) -> u64 {
+        self.layout.row_count()
     }
 
     fn register_splits(

vortex-layout/src/reader.rs

@@ -7,7 +7,6 @@ use std::sync::Arc;
 
 use futures::future::BoxFuture;
 use once_cell::sync::OnceCell;
-use vortex_array::stats::Precision;
 use vortex_array::{ArrayRef, MaskFuture};
 use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexResult, vortex_bail};
@@ -28,9 +27,8 @@ pub trait LayoutReader: 'static + Send + Sync {
     /// Returns the un-projected dtype of the layout reader.
     fn dtype(&self) -> &DType;
 
-    /// Returns the number of rows in the layout reader.
-    /// An inexact count may be larger or smaller than the actual row count.
-    fn row_count(&self) -> Precision<u64>;
+    /// Returns the number of rows in the layout.
+    fn row_count(&self) -> u64;
 
     /// Register the splits of this layout reader.
     // TODO(ngates): this is a temporary API until we make layout readers stream based.

vortex-scan/src/split_by.rs

@@ -4,9 +4,8 @@
 use std::collections::BTreeSet;
 use std::iter::once;
 
-use vortex_array::stats::StatBound;
 use vortex_dtype::FieldMask;
-use vortex_error::{VortexResult, vortex_err};
+use vortex_error::VortexResult;
 use vortex_layout::LayoutReader;
 
 /// Defines how the Vortex file is split into batches for reading.
@@ -40,9 +39,7 @@ impl SplitBy {
                 row_splits
             }
             SplitBy::RowCount(n) => {
-                let row_count = *layout_reader.row_count().to_exact().ok_or_else(|| {
-                    vortex_err!("Cannot split layout by row count, row count is not exact")
-                })?;
+                let row_count = layout_reader.row_count();
                 (0..row_count).step_by(n).chain(once(row_count)).collect()
             }
         })