vortex_layout/
reader.rs

1use std::ops::Range;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use futures::FutureExt;
6use futures::future::{BoxFuture, Shared};
7use vortex_array::ArrayRef;
8use vortex_dtype::DType;
9use vortex_error::{SharedVortexResult, VortexError, VortexResult};
10use vortex_expr::ExprRef;
11use vortex_mask::Mask;
12
13use crate::Layout;
14
15/// A [`LayoutReader`] is an instance of a [`Layout`] that can cache state across multiple
16/// operations.
17///
18/// Since different row ranges of the reader may be evaluated by different threads, it is required
19/// to be both `Send` and `Sync`.
20pub trait LayoutReader: 'static + ExprEvaluator {
21    /// Returns the [`Layout`] of this reader.
22    fn layout(&self) -> &Layout;
23
24    /// Returns the row count of the layout.
25    fn row_count(&self) -> u64 {
26        self.layout().row_count()
27    }
28
29    /// Returns the DType of the layout.
30    fn dtype(&self) -> &DType {
31        self.layout().dtype()
32    }
33
34    fn children(&self) -> VortexResult<Vec<Arc<dyn LayoutReader>>>;
35}
36
37pub trait LayoutReaderExt: LayoutReader {
38    /// Box the layout scan.
39    fn into_arc(self) -> Arc<dyn LayoutReader>
40    where
41        Self: Sized + 'static,
42    {
43        Arc::new(self) as _
44    }
45}
46
47impl<L: LayoutReader> LayoutReaderExt for L {}
48
49pub type MaskFuture = Shared<BoxFuture<'static, SharedVortexResult<Mask>>>;
50
51/// Create a resolved [`MaskFuture`] from a [`Mask`].
52pub fn mask_future_ready(mask: Mask) -> MaskFuture {
53    async move { Ok::<_, Arc<VortexError>>(mask) }
54        .boxed()
55        .shared()
56}
57
58/// A trait for evaluating expressions against a [`LayoutReader`].
59pub trait ExprEvaluator: Send + Sync {
60    /// Performs an approximate evaluation of the expression against the layout reader.
61    fn pruning_evaluation(
62        &self,
63        row_range: &Range<u64>,
64        expr: &ExprRef,
65    ) -> VortexResult<Box<dyn PruningEvaluation>>;
66
67    /// Performs an exact evaluation of the expression against the layout reader.
68    fn filter_evaluation(
69        &self,
70        row_range: &Range<u64>,
71        expr: &ExprRef,
72    ) -> VortexResult<Box<dyn MaskEvaluation>>;
73
74    /// Evaluates the expression against the layout.
75    fn projection_evaluation(
76        &self,
77        row_range: &Range<u64>,
78        expr: &ExprRef,
79    ) -> VortexResult<Box<dyn ArrayEvaluation>>;
80}
81
82impl ExprEvaluator for Arc<dyn LayoutReader> {
83    fn pruning_evaluation(
84        &self,
85        row_range: &Range<u64>,
86        expr: &ExprRef,
87    ) -> VortexResult<Box<dyn PruningEvaluation>> {
88        self.as_ref().pruning_evaluation(row_range, expr)
89    }
90
91    fn filter_evaluation(
92        &self,
93        row_range: &Range<u64>,
94        expr: &ExprRef,
95    ) -> VortexResult<Box<dyn MaskEvaluation>> {
96        self.as_ref().filter_evaluation(row_range, expr)
97    }
98
99    fn projection_evaluation(
100        &self,
101        row_range: &Range<u64>,
102        expr: &ExprRef,
103    ) -> VortexResult<Box<dyn ArrayEvaluation>> {
104        self.as_ref().projection_evaluation(row_range, expr)
105    }
106}
107
108#[async_trait]
109pub trait PruningEvaluation: 'static + Send + Sync {
110    async fn invoke(&self, mask: Mask) -> VortexResult<Mask>;
111}
112
113pub struct NoOpPruningEvaluation;
114
115#[async_trait]
116impl PruningEvaluation for NoOpPruningEvaluation {
117    async fn invoke(&self, mask: Mask) -> VortexResult<Mask> {
118        Ok(mask)
119    }
120}
121
122/// Refines the given mask, returning a mask equal in length to the input mask.
123#[async_trait]
124pub trait MaskEvaluation: 'static + Send + Sync {
125    async fn invoke(&self, mask: Mask) -> VortexResult<Mask>;
126}
127
128/// Evaluates an expression against an array, returning an array equal in length to the true count
129/// of the input mask.
130#[async_trait]
131pub trait ArrayEvaluation: 'static + Send + Sync {
132    async fn invoke(&self, mask: Mask) -> VortexResult<ArrayRef>;
133}
vortex_layout/reader.rs

vortex_layout/
reader.rs
