Scaffolding for layout stats (#1885)

Author: gatesn

vortex-file/src/v2/file.rs

@@ -10,7 +10,7 @@ use vortex_array::ContextRef;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_io::VortexReadAt;
-use vortex_layout::{LayoutData, LayoutReader};
+use vortex_layout::{ExprEvaluator, LayoutData, LayoutReader};
 use vortex_scan::Scan;
 
 use crate::v2::segments::cache::SegmentCache;
@@ -56,7 +56,7 @@ impl<R: VortexReadAt + Unpin> VortexFile<R> {
 
             let eval = scan
                 .range_scan(row_range)?
-                .evaluate_async(reader.evaluator());
+                .evaluate_async(|row_mask, expr| reader.evaluate_expr(row_mask, expr));
             pin_mut!(eval);
 
             poll_fn(|cx| {

…-layout/src/layouts/chunked/evaluator.rs …-layout/src/layouts/chunked/eval_expr.rsvortex-layout/src/layouts/chunked/evaluator.rs renamed to vortex-layout/src/layouts/chunked/eval_expr.rs vortex-layout/src/layouts/chunked/evaluator.rs renamed to vortex-layout/src/layouts/chunked/eval_expr.rs

@@ -7,14 +7,19 @@ use vortex_error::VortexResult;
 use vortex_expr::pruning::PruningPredicate;
 use vortex_expr::ExprRef;
 use vortex_scalar::Scalar;
-use vortex_scan::{AsyncEvaluator, RowMask};
+use vortex_scan::RowMask;
 
 use crate::layouts::chunked::reader::ChunkedReader;
 use crate::reader::LayoutScanExt;
+use crate::ExprEvaluator;
 
 #[async_trait(?Send)]
-impl AsyncEvaluator for ChunkedReader {
-    async fn evaluate(self: &Self, row_mask: RowMask, expr: ExprRef) -> VortexResult<ArrayData> {
+impl ExprEvaluator for ChunkedReader {
+    async fn evaluate_expr(
+        self: &Self,
+        row_mask: RowMask,
+        expr: ExprRef,
+    ) -> VortexResult<ArrayData> {
         // Compute the result dtype of the expression.
         let dtype = expr
             .evaluate(&Canonical::empty(self.dtype())?.into_array())?
@@ -74,7 +79,7 @@ impl AsyncEvaluator for ChunkedReader {
                 .shift(chunk_range.start)?;
 
             let expr = expr.clone();
-            chunks.push(chunk_reader.evaluate(chunk_mask, expr).boxed_local());
+            chunks.push(chunk_reader.evaluate_expr(chunk_mask, expr).boxed_local());
         }
 
         // Wait for all chunks to be evaluated
@@ -130,7 +135,7 @@ mod test {
             let result = layout
                 .reader(segments, Default::default())
                 .unwrap()
-                .evaluate(
+                .evaluate_expr(
                     RowMask::new_valid_between(0, layout.row_count()),
                     Identity::new_expr(),
                 )
@@ -155,7 +160,7 @@ mod test {
             let expr = gt(Identity::new_expr(), lit(7));
 
             let result = reader
-                .evaluate(RowMask::new_valid_between(0, row_count), expr.clone())
+                .evaluate_expr(RowMask::new_valid_between(0, row_count), expr.clone())
                 .await
                 .unwrap();
             let result = ChunkedArray::try_from(result).unwrap();

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

@@ -0,0 +1,18 @@
+use async_trait::async_trait;
+use vortex_array::stats::{Stat, StatsSet};
+use vortex_dtype::FieldPath;
+use vortex_error::VortexResult;
+
+use crate::layouts::chunked::reader::ChunkedReader;
+use crate::StatsEvaluator;
+
+#[async_trait(?Send)]
+impl StatsEvaluator for ChunkedReader {
+    async fn evaluate_stats(
+        &self,
+        field_paths: &[FieldPath],
+        _stats: &[Stat],
+    ) -> VortexResult<Vec<StatsSet>> {
+        Ok(vec![StatsSet::default(); field_paths.len()])
+    }
+}

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

@@ -1,6 +1,6 @@
-mod evaluator;
+mod eval_expr;
+mod eval_stats;
 mod reader;
-// mod stats;
 pub mod stats_table;
 pub mod writer;
 

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

@@ -5,13 +5,13 @@ use vortex_array::stats::{stats_from_bitset_bytes, Stat};
 use vortex_array::ContextRef;
 use vortex_error::{vortex_panic, VortexResult};
 use vortex_expr::Identity;
-use vortex_scan::{AsyncEvaluator, RowMask};
+use vortex_scan::RowMask;
 
 use crate::layouts::chunked::stats_table::StatsTable;
 use crate::layouts::chunked::ChunkedLayout;
 use crate::reader::LayoutReader;
 use crate::segments::AsyncSegmentReader;
-use crate::{LayoutData, LayoutEncoding};
+use crate::{ExprEvaluator, LayoutData, LayoutEncoding};
 
 #[derive(Clone)]
 pub struct ChunkedReader {
@@ -83,7 +83,7 @@ impl ChunkedReader {
 
                         let stats_array = stats_layout
                             .reader(self.segments.clone(), self.ctx.clone())?
-                            .evaluate(
+                            .evaluate_expr(
                                 RowMask::new_valid_between(0, nchunks as u64),
                                 Identity::new_expr(),
                             )
@@ -119,8 +119,4 @@ impl LayoutReader for ChunkedReader {
     fn layout(&self) -> &LayoutData {
         &self.layout
     }
-
-    fn evaluator(&self) -> &dyn AsyncEvaluator {
-        self
-    }
 }

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

@@ -7,15 +7,19 @@ use vortex_array::ArrayData;
 use vortex_error::{vortex_err, VortexExpect, VortexResult};
 use vortex_expr::ExprRef;
 use vortex_flatbuffers::{array as fba, FlatBuffer};
-use vortex_scan::{AsyncEvaluator, RowMask};
+use vortex_scan::RowMask;
 
 use crate::layouts::flat::reader::FlatReader;
 use crate::reader::LayoutScanExt;
-use crate::LayoutReader;
+use crate::{ExprEvaluator, LayoutReader};
 
 #[async_trait(?Send)]
-impl AsyncEvaluator for FlatReader {
-    async fn evaluate(self: &Self, row_mask: RowMask, expr: ExprRef) -> VortexResult<ArrayData> {
+impl ExprEvaluator for FlatReader {
+    async fn evaluate_expr(
+        self: &Self,
+        row_mask: RowMask,
+        expr: ExprRef,
+    ) -> VortexResult<ArrayData> {
         // Fetch all the array buffers.
         let mut buffers = try_join_all(
             self.layout()
@@ -80,7 +84,7 @@ mod test {
             let result = layout
                 .reader(Arc::new(segments), Default::default())
                 .unwrap()
-                .evaluate(
+                .evaluate_expr(
                     RowMask::new_valid_between(0, layout.row_count()),
                     Identity::new_expr(),
                 )
@@ -106,7 +110,7 @@ mod test {
             let result = layout
                 .reader(Arc::new(segments), Default::default())
                 .unwrap()
-                .evaluate(RowMask::new_valid_between(0, layout.row_count()), expr)
+                .evaluate_expr(RowMask::new_valid_between(0, layout.row_count()), expr)
                 .await
                 .unwrap()
                 .into_bool()

…tex-layout/src/layouts/flat/evaluator.rs …tex-layout/src/layouts/flat/eval_expr.rsvortex-layout/src/layouts/flat/evaluator.rs renamed to vortex-layout/src/layouts/flat/eval_expr.rs vortex-layout/src/layouts/flat/evaluator.rs renamed to vortex-layout/src/layouts/flat/eval_expr.rs

@@ -0,0 +1,18 @@
+use async_trait::async_trait;
+use vortex_array::stats::{Stat, StatsSet};
+use vortex_dtype::FieldPath;
+use vortex_error::VortexResult;
+
+use crate::layouts::flat::reader::FlatReader;
+use crate::StatsEvaluator;
+
+#[async_trait(?Send)]
+impl StatsEvaluator for FlatReader {
+    async fn evaluate_stats(
+        &self,
+        field_paths: &[FieldPath],
+        _stats: &[Stat],
+    ) -> VortexResult<Vec<StatsSet>> {
+        Ok(vec![StatsSet::default(); field_paths.len()])
+    }
+}

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

@@ -1,6 +1,7 @@
-mod evaluator;
+mod eval_expr;
 mod reader;
 // mod stats;
+mod eval_stats;
 pub mod writer;
 
 use std::collections::BTreeSet;

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

@@ -2,7 +2,6 @@ use std::sync::Arc;
 
 use vortex_array::ContextRef;
 use vortex_error::{vortex_panic, VortexResult};
-use vortex_scan::AsyncEvaluator;
 
 use crate::layouts::flat::FlatLayout;
 use crate::reader::LayoutReader;
@@ -45,8 +44,4 @@ impl LayoutReader for FlatReader {
     fn layout(&self) -> &LayoutData {
         &self.layout
     }
-
-    fn evaluator(&self) -> &dyn AsyncEvaluator {
-        self
-    }
 }