vortex_expr/exprs/
root.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::operator::OperatorRef;
5use vortex_array::stats::Stat;
6use vortex_array::{ArrayRef, DeserializeMetadata, EmptyMetadata};
7use vortex_dtype::{DType, FieldPath};
8use vortex_error::{VortexResult, vortex_bail};
9
10use crate::display::{DisplayAs, DisplayFormat};
11use crate::{
12    AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, StatsCatalog, VTable, vtable,
13};
14
15vtable!(Root);
16
17/// An expression that returns the full scope of the expression evaluation.
18#[derive(Debug, Clone, PartialEq, Eq, Hash)]
19pub struct RootExpr;
20
21pub struct RootExprEncoding;
22
23impl VTable for RootVTable {
24    type Expr = RootExpr;
25    type Encoding = RootExprEncoding;
26    type Metadata = EmptyMetadata;
27
28    fn id(_encoding: &Self::Encoding) -> ExprId {
29        ExprId::new_ref("root")
30    }
31
32    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
33        ExprEncodingRef::new_ref(RootExprEncoding.as_ref())
34    }
35
36    fn metadata(_expr: &Self::Expr) -> Option<Self::Metadata> {
37        Some(EmptyMetadata)
38    }
39
40    fn children(_expr: &Self::Expr) -> Vec<&ExprRef> {
41        vec![]
42    }
43
44    fn with_children(expr: &Self::Expr, _children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
45        Ok(expr.clone())
46    }
47
48    fn build(
49        _encoding: &Self::Encoding,
50        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
51        children: Vec<ExprRef>,
52    ) -> VortexResult<Self::Expr> {
53        if !children.is_empty() {
54            vortex_bail!(
55                "Root expression does not have children, got: {:?}",
56                children
57            );
58        }
59        Ok(RootExpr)
60    }
61
62    fn evaluate(_expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
63        Ok(scope.root().clone())
64    }
65
66    fn return_dtype(_expr: &Self::Expr, scope: &DType) -> VortexResult<DType> {
67        Ok(scope.clone())
68    }
69
70    fn operator(_expr: &Self::Expr, scope: &OperatorRef) -> VortexResult<Option<OperatorRef>> {
71        Ok(Some(scope.clone()))
72    }
73}
74
75impl DisplayAs for RootExpr {
76    fn fmt_as(&self, df: DisplayFormat, f: &mut std::fmt::Formatter) -> std::fmt::Result {
77        match df {
78            DisplayFormat::Compact => {
79                write!(f, "$")
80            }
81            DisplayFormat::Tree => {
82                write!(f, "Root")
83            }
84        }
85    }
86}
87
88impl AnalysisExpr for RootExpr {
89    fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
90        catalog.stats_ref(&self.field_path()?, Stat::Max)
91    }
92
93    fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
94        catalog.stats_ref(&self.field_path()?, Stat::Min)
95    }
96
97    fn nan_count(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
98        catalog.stats_ref(&self.field_path()?, Stat::NaNCount)
99    }
100
101    fn field_path(&self) -> Option<FieldPath> {
102        Some(FieldPath::root())
103    }
104}
105
106/// Creates an expression that references the root scope.
107///
108/// Returns the entire input array as passed to the expression evaluator.
109/// This is commonly used as the starting point for field access and other operations.
110pub fn root() -> ExprRef {
111    RootExpr.into_expr()
112}
113
114/// Return whether the expression is a root expression.
115pub fn is_root(expr: &ExprRef) -> bool {
116    expr.is::<RootVTable>()
117}