Layout reader updates (#3400)

See #3382

---------

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

Author: gatesn

vortex-layout/src/layout.rs

@@ -1,12 +1,11 @@
 use std::any::Any;
-use std::collections::BTreeSet;
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
 use arcref::ArcRef;
 use itertools::Itertools;
 use vortex_array::{ArrayContext, SerializeMetadata};
-use vortex_dtype::{DType, FieldMask, FieldName};
+use vortex_dtype::{DType, FieldName};
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 
 use crate::segments::{SegmentId, SegmentSource};
@@ -48,13 +47,6 @@ pub trait Layout: 'static + Send + Sync + Debug + private::Sealed {
     /// Get the segment IDs for this layout.
     fn segment_ids(&self) -> Vec<SegmentId>;
 
-    fn register_splits(
-        &self,
-        field_mask: &[FieldMask],
-        row_offset: u64,
-        splits: &mut BTreeSet<u64>,
-    ) -> VortexResult<()>;
-
     fn new_reader(
         &self,
         name: &Arc<str>,
@@ -249,15 +241,6 @@ impl<V: VTable> Layout for LayoutAdapter<V> {
         V::segment_ids(&self.0)
     }
 
-    fn register_splits(
-        &self,
-        field_mask: &[FieldMask],
-        row_offset: u64,
-        splits: &mut BTreeSet<u64>,
-    ) -> VortexResult<()> {
-        V::register_splits(&self.0, field_mask, row_offset, splits)
-    }
-
     fn new_reader(
         &self,
         name: &Arc<str>,

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

@@ -1,11 +1,10 @@
 mod reader;
 pub mod writer;
 
-use std::collections::BTreeSet;
 use std::sync::Arc;
 
 use vortex_array::{ArrayContext, DeserializeMetadata, EmptyMetadata};
-use vortex_dtype::{DType, FieldMask};
+use vortex_dtype::DType;
 use vortex_error::VortexResult;
 
 use crate::children::LayoutChildren;
@@ -58,22 +57,6 @@ impl VTable for ChunkedVTable {
         LayoutChildType::Chunk((idx, layout.chunk_offsets[idx]))
     }
 
-    fn register_splits(
-        layout: &Self::Layout,
-        field_mask: &[FieldMask],
-        row_offset: u64,
-        splits: &mut BTreeSet<u64>,
-    ) -> VortexResult<()> {
-        let mut offset = row_offset;
-        for i in 0..layout.nchildren() {
-            let child = layout.child(i)?;
-            child.register_splits(field_mask, offset, splits)?;
-            offset += child.row_count();
-            splits.insert(offset);
-        }
-        Ok(())
-    }
-
     fn new_reader(
         layout: &Self::Layout,
         name: &Arc<str>,

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

@@ -1,4 +1,5 @@
-use std::ops::{Deref, Range};
+use std::collections::BTreeSet;
+use std::ops::Range;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -7,8 +8,9 @@ use futures::stream::FuturesOrdered;
 use futures::{FutureExt, TryStreamExt};
 use itertools::Itertools;
 use vortex_array::arrays::ChunkedArray;
+use vortex_array::stats::Precision;
 use vortex_array::{ArrayContext, ArrayRef};
-use vortex_dtype::DType;
+use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_expr::ExprRef;
 use vortex_mask::Mask;
@@ -17,7 +19,7 @@ use crate::layouts::chunked::ChunkedLayout;
 use crate::reader::LayoutReader;
 use crate::segments::SegmentSource;
 use crate::{
-    ArrayEvaluation, Layout, LayoutReaderRef, LazyReaderChildren, MaskEvaluation, PruningEvaluation,
+    ArrayEvaluation, LayoutReaderRef, LazyReaderChildren, MaskEvaluation, PruningEvaluation,
 };
 
 /// A [`LayoutReader`] for chunked layouts.
@@ -111,19 +113,35 @@ impl ChunkedReader {
     }
 }
 
-impl Deref for ChunkedReader {
-    type Target = dyn Layout;
-
-    fn deref(&self) -> &Self::Target {
-        self.layout.as_ref()
-    }
-}
-
 impl LayoutReader for ChunkedReader {
     fn name(&self) -> &Arc<str> {
         &self.name
     }
 
+    fn dtype(&self) -> &DType {
+        self.layout.dtype()
+    }
+
+    fn row_count(&self) -> Precision<u64> {
+        Precision::Exact(self.layout.row_count())
+    }
+
+    fn register_splits(
+        &self,
+        field_mask: &[FieldMask],
+        row_offset: u64,
+        splits: &mut BTreeSet<u64>,
+    ) -> VortexResult<()> {
+        let mut offset = row_offset;
+        for i in 0..self.layout.nchildren() {
+            let child = self.chunk_reader(i)?;
+            child.register_splits(field_mask, offset, splits)?;
+            offset += self.layout.child(i)?.row_count();
+            splits.insert(offset);
+        }
+        Ok(())
+    }
+
     fn pruning_evaluation(
         &self,
         row_range: &Range<u64>,

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

@@ -1,12 +1,11 @@
 mod reader;
 pub mod writer;
 
-use std::collections::BTreeSet;
 use std::sync::Arc;
 
 use reader::DictReader;
 use vortex_array::{ArrayContext, DeserializeMetadata, ProstMetadata};
-use vortex_dtype::{DType, FieldMask, PType};
+use vortex_dtype::{DType, PType};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_panic};
 
 use crate::children::LayoutChildren;
@@ -68,15 +67,6 @@ impl VTable for DictVTable {
         }
     }
 
-    fn register_splits(
-        layout: &Self::Layout,
-        field_mask: &[FieldMask],
-        row_offset: u64,
-        splits: &mut BTreeSet<u64>,
-    ) -> VortexResult<()> {
-        layout.codes.register_splits(field_mask, row_offset, splits)
-    }
-
     fn new_reader(
         layout: &Self::Layout,
         name: &Arc<str>,

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

@@ -1,13 +1,16 @@
-use std::ops::{BitAnd, Deref, Range};
+use std::collections::BTreeSet;
+use std::ops::{BitAnd, Range};
 use std::sync::{Arc, OnceLock};
 
 use async_trait::async_trait;
 use dashmap::DashMap;
 use futures::{FutureExt, join};
 use vortex_array::arrays::StructArray;
 use vortex_array::compute::{MinMaxResult, filter, min_max};
+use vortex_array::stats::Precision;
 use vortex_array::{Array, ArrayContext, ArrayRef, ToCanonical};
 use vortex_dict::DictArray;
+use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_expr::{ExprRef, Identity};
 use vortex_mask::Mask;
@@ -16,7 +19,7 @@ use super::DictLayout;
 use crate::layouts::SharedArrayFuture;
 use crate::segments::SegmentSource;
 use crate::{
-    ArrayEvaluation, Layout, LayoutReader, LayoutReaderRef, MaskEvaluation, NoOpPruningEvaluation,
+    ArrayEvaluation, LayoutReader, LayoutReaderRef, MaskEvaluation, NoOpPruningEvaluation,
     PruningEvaluation,
 };
 
@@ -25,6 +28,8 @@ pub struct DictReader {
     #[allow(dead_code)] // Typically used for logging
     name: Arc<str>,
 
+    /// Length of the values array
+    values_len: usize,
     /// Cached dict values array
     values_array: OnceLock<SharedArrayFuture>,
     /// Cache of expression evaluation results on the values array by expression
@@ -34,21 +39,14 @@ pub struct DictReader {
     codes: LayoutReaderRef,
 }
 
-impl Deref for DictReader {
-    type Target = dyn Layout;
-
-    fn deref(&self) -> &Self::Target {
-        self.layout.deref()
-    }
-}
-
 impl DictReader {
     pub(super) fn try_new(
         layout: DictLayout,
         name: Arc<str>,
         segment_source: &Arc<dyn SegmentSource>,
         ctx: &ArrayContext,
     ) -> VortexResult<Self> {
+        let values_len = usize::try_from(layout.values.row_count())?;
         let values =
             layout
                 .values
@@ -61,6 +59,7 @@ impl DictReader {
         Ok(Self {
             layout,
             name,
+            values_len,
             values_array: Default::default(),
             values_evals: Default::default(),
             values,
@@ -71,21 +70,18 @@ impl DictReader {
     fn values_array(&self) -> SharedArrayFuture {
         // We capture the name, so it may be wrong if we re-use the same reader within multiple
         // different parent readers. But that's rare...
+        let values_len = self.values_len;
         self.values_array
             .get_or_init(move || {
-                let values_len = self.values.row_count();
                 let eval = self
                     .values
-                    .projection_evaluation(&(0..values_len), &Identity::new_expr())
+                    .projection_evaluation(&(0..values_len as u64), &Identity::new_expr())
                     .vortex_expect("must construct dict values array evaluation");
 
                 async move {
-                    eval.invoke(Mask::new_true(
-                        usize::try_from(values_len)
-                            .vortex_expect("dict values length must fit in u32"),
-                    ))
-                    .await
-                    .map_err(Arc::new)
+                    eval.invoke(Mask::new_true(values_len))
+                        .await
+                        .map_err(Arc::new)
                 }
                 .boxed()
                 .shared()
@@ -111,6 +107,23 @@ impl LayoutReader for DictReader {
         &self.name
     }
 
+    fn dtype(&self) -> &DType {
+        self.layout.dtype()
+    }
+
+    fn row_count(&self) -> Precision<u64> {
+        Precision::Exact(self.layout.row_count())
+    }
+
+    fn register_splits(
+        &self,
+        field_mask: &[FieldMask],
+        row_offset: u64,
+        splits: &mut BTreeSet<u64>,
+    ) -> VortexResult<()> {
+        self.codes.register_splits(field_mask, row_offset, splits)
+    }
+
     fn pruning_evaluation(
         &self,
         _row_range: &Range<u64>,

vortex-layout/src/layouts/filter.rs

@@ -1,5 +1,6 @@
+use std::collections::BTreeSet;
 use std::iter;
-use std::ops::{BitAnd, Deref, Range};
+use std::ops::{BitAnd, Range};
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -8,14 +9,14 @@ use dashmap::DashMap;
 use itertools::Itertools;
 use parking_lot::RwLock;
 use sketches_ddsketch::DDSketch;
+use vortex_array::stats::Precision;
+use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
 use vortex_expr::ExprRef;
 use vortex_expr::forms::cnf::cnf;
 use vortex_mask::Mask;
 
-use crate::{
-    ArrayEvaluation, Layout, LayoutReader, LayoutReaderRef, MaskEvaluation, PruningEvaluation,
-};
+use crate::{ArrayEvaluation, LayoutReader, LayoutReaderRef, MaskEvaluation, PruningEvaluation};
 
 /// The selectivity histogram quantile to use for reordering conjuncts. Where 0 == no rows match.
 const DEFAULT_SELECTIVITY_QUANTILE: f64 = 0.1;
@@ -40,19 +41,29 @@ impl FilterLayoutReader {
     }
 }
 
-impl Deref for FilterLayoutReader {
-    type Target = dyn Layout;
-
-    fn deref(&self) -> &Self::Target {
-        self.child.deref()
-    }
-}
-
 impl LayoutReader for FilterLayoutReader {
     fn name(&self) -> &Arc<str> {
         self.child.name()
     }
 
+    fn dtype(&self) -> &DType {
+        self.child.dtype()
+    }
+
+    fn row_count(&self) -> Precision<u64> {
+        self.child.row_count()
+    }
+
+    fn register_splits(
+        &self,
+        field_mask: &[FieldMask],
+        row_offset: u64,
+        splits: &mut BTreeSet<u64>,
+    ) -> VortexResult<()> {
+        // Pass-through the splits to the child layout reader.
+        self.child.register_splits(field_mask, row_offset, splits)
+    }
+
     fn pruning_evaluation(
         &self,
         row_range: &Range<u64>,

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

@@ -1,11 +1,10 @@
 mod reader;
 pub mod writer;
 
-use std::collections::BTreeSet;
 use std::sync::Arc;
 
 use vortex_array::{ArrayContext, DeserializeMetadata, EmptyMetadata};
-use vortex_dtype::{DType, FieldMask};
+use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail, vortex_panic};
 
 use crate::children::LayoutChildren;
@@ -58,21 +57,6 @@ impl VTable for FlatVTable {
         vortex_panic!("Flat layout has no children");
     }
 
-    fn register_splits(
-        layout: &Self::Layout,
-        field_mask: &[FieldMask],
-        row_offset: u64,
-        splits: &mut BTreeSet<u64>,
-    ) -> VortexResult<()> {
-        for path in field_mask {
-            if path.matches_root() {
-                splits.insert(row_offset + layout.row_count());
-                break;
-            }
-        }
-        Ok(())
-    }
-
     fn new_reader(
         layout: &Self::Layout,
         name: &Arc<str>,