vortex_expr/exprs/
root.rs

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