vortex_layout/layouts/row_idx/
expr.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::{Display, Formatter};
5
6use vortex_array::{ArrayRef, DeserializeMetadata, EmptyMetadata};
7use vortex_dtype::{DType, Nullability, PType};
8use vortex_error::{VortexResult, vortex_bail};
9use vortex_expr::{
10    AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable,
11};
12
13vtable!(RowIdx);
14
15#[derive(Clone, Debug, PartialEq, Eq, Hash)]
16pub struct RowIdxExpr;
17
18impl AnalysisExpr for RowIdxExpr {}
19
20impl Display for RowIdxExpr {
21    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
22        write!(f, "#row_idx")
23    }
24}
25
26#[derive(Clone)]
27pub struct RowIdxExprEncoding;
28
29impl VTable for RowIdxVTable {
30    type Expr = RowIdxExpr;
31    type Encoding = RowIdxExprEncoding;
32    type Metadata = EmptyMetadata;
33
34    fn id(_encoding: &Self::Encoding) -> ExprId {
35        ExprId::new_ref("vortex.row_idx")
36    }
37
38    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
39        ExprEncodingRef::new_ref(RowIdxExprEncoding.as_ref())
40    }
41
42    fn metadata(_expr: &Self::Expr) -> Option<Self::Metadata> {
43        // Serializable, but with no metadata
44        Some(EmptyMetadata)
45    }
46
47    fn children(_expr: &Self::Expr) -> Vec<&ExprRef> {
48        vec![]
49    }
50
51    fn with_children(expr: &Self::Expr, _children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
52        Ok(expr.clone())
53    }
54
55    fn build(
56        _encoding: &Self::Encoding,
57        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
58        children: Vec<ExprRef>,
59    ) -> VortexResult<Self::Expr> {
60        if !children.is_empty() {
61            vortex_bail!(
62                "RowIdxExpr does not expect any children, got {}",
63                children.len()
64            );
65        }
66        Ok(RowIdxExpr)
67    }
68
69    fn evaluate(_expr: &Self::Expr, _scope: &Scope) -> VortexResult<ArrayRef> {
70        vortex_bail!(
71            "RowIdxExpr should not be evaluated directly, use it in the context of a Vortex scan and it will be substituted for a row index array"
72        );
73    }
74
75    fn return_dtype(_expr: &Self::Expr, _scope: &DType) -> VortexResult<DType> {
76        Ok(DType::Primitive(PType::U64, Nullability::NonNullable))
77    }
78}
79
80pub fn row_idx() -> ExprRef {
81    RowIdxExpr.into_expr()
82}