vortex_layout/
reader.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use vortex_array::ArrayRef;
5use vortex_dtype::DType;
6use vortex_error::VortexResult;
7use vortex_expr::ExprRef;
8
9use crate::{Layout, RowMask};
10
11/// A [`LayoutReader`] is an instance of a [`Layout`] that can cache state across multiple
12/// operations.
13///
14/// Since different row ranges of the reader may be evaluated by different threads, it is required
15/// to be both `Send` and `Sync`.
16pub trait LayoutReader: 'static + Send + Sync + ExprEvaluator {
17    /// Returns the [`Layout`] of this reader.
18    fn layout(&self) -> &Layout;
19}
20
21impl LayoutReader for Arc<dyn LayoutReader> {
22    fn layout(&self) -> &Layout {
23        self.as_ref().layout()
24    }
25}
26
27/// A trait for evaluating expressions against a [`LayoutReader`].
28///
29/// FIXME(ngates): what if this was evaluating_predicate(mask, expr) -> mask,
30///  evaluate_filter(mask, scan) -> Array, and evaluate_projection(mask, expr) -> Array?
31#[async_trait]
32pub trait ExprEvaluator: Send + Sync {
33    async fn evaluate_expr(&self, row_mask: RowMask, expr: ExprRef) -> VortexResult<ArrayRef>;
34
35    /// Refine the row mask by evaluating any pruning. This should be relatively cheap, statistics
36    /// based evaluation, and returns an approximate result.
37    async fn refine_mask(&self, row_mask: RowMask, _expr: ExprRef) -> VortexResult<RowMask> {
38        Ok(row_mask)
39    }
40}
41
42#[async_trait]
43impl ExprEvaluator for Arc<dyn LayoutReader> {
44    async fn evaluate_expr(&self, row_mask: RowMask, expr: ExprRef) -> VortexResult<ArrayRef> {
45        self.as_ref().evaluate_expr(row_mask, expr).await
46    }
47
48    async fn refine_mask(&self, row_mask: RowMask, expr: ExprRef) -> VortexResult<RowMask> {
49        self.as_ref().refine_mask(row_mask, expr).await
50    }
51}
52
53pub trait LayoutReaderExt: LayoutReader {
54    /// Box the layout scan.
55    fn into_arc(self) -> Arc<dyn LayoutReader>
56    where
57        Self: Sized + 'static,
58    {
59        Arc::new(self) as _
60    }
61
62    /// Returns the row count of the layout.
63    fn row_count(&self) -> u64 {
64        self.layout().row_count()
65    }
66
67    /// Returns the DType of the layout.
68    fn dtype(&self) -> &DType {
69        self.layout().dtype()
70    }
71}
72
73impl<L: LayoutReader> LayoutReaderExt for L {}