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::Formatter;
5
6use vortex_array::{ArrayRef, DeserializeMetadata, EmptyMetadata};
7use vortex_dtype::{DType, Nullability, PType};
8use vortex_error::{VortexResult, vortex_bail};
9use vortex_expr::display::{DisplayAs, DisplayFormat};
10use vortex_expr::{
11    AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable,
12};
13
14vtable!(RowIdx);
15
16#[derive(Clone, Debug, PartialEq, Eq, Hash)]
17pub struct RowIdxExpr;
18
19impl AnalysisExpr for RowIdxExpr {}
20
21impl DisplayAs for RowIdxExpr {
22    fn fmt_as(&self, df: DisplayFormat, f: &mut Formatter) -> std::fmt::Result {
23        match df {
24            DisplayFormat::Compact => write!(f, "#row_idx"),
25            DisplayFormat::Tree => write!(f, "RowIdx"),
26        }
27    }
28}
29
30#[derive(Clone)]
31pub struct RowIdxExprEncoding;
32
33impl VTable for RowIdxVTable {
34    type Expr = RowIdxExpr;
35    type Encoding = RowIdxExprEncoding;
36    type Metadata = EmptyMetadata;
37
38    fn id(_encoding: &Self::Encoding) -> ExprId {
39        ExprId::new_ref("vortex.row_idx")
40    }
41
42    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
43        ExprEncodingRef::new_ref(RowIdxExprEncoding.as_ref())
44    }
45
46    fn metadata(_expr: &Self::Expr) -> Option<Self::Metadata> {
47        // Serializable, but with no metadata
48        Some(EmptyMetadata)
49    }
50
51    fn children(_expr: &Self::Expr) -> Vec<&ExprRef> {
52        vec![]
53    }
54
55    fn with_children(expr: &Self::Expr, _children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
56        Ok(expr.clone())
57    }
58
59    fn build(
60        _encoding: &Self::Encoding,
61        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
62        children: Vec<ExprRef>,
63    ) -> VortexResult<Self::Expr> {
64        if !children.is_empty() {
65            vortex_bail!(
66                "RowIdxExpr does not expect any children, got {}",
67                children.len()
68            );
69        }
70        Ok(RowIdxExpr)
71    }
72
73    fn evaluate(_expr: &Self::Expr, _scope: &Scope) -> VortexResult<ArrayRef> {
74        vortex_bail!(
75            "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"
76        );
77    }
78
79    fn return_dtype(_expr: &Self::Expr, _scope: &DType) -> VortexResult<DType> {
80        Ok(DType::Primitive(PType::U64, Nullability::NonNullable))
81    }
82}
83
84pub fn row_idx() -> ExprRef {
85    RowIdxExpr.into_expr()
86}
87
88#[cfg(test)]
89mod tests {
90    use rstest::rstest;
91    use vortex_array::EmptyMetadata;
92    use vortex_dtype::{DType, Nullability, PType};
93    use vortex_expr::{ExprRef, IntoExpr, VTable};
94
95    use crate::layouts::row_idx::{RowIdxExpr, RowIdxExprEncoding, RowIdxVTable, row_idx};
96
97    #[test]
98    fn test_row_idx_expr_creation() {
99        let expr = row_idx();
100        assert!(expr.is::<RowIdxVTable>());
101    }
102
103    #[test]
104    fn test_vtable_id() {
105        let encoding = RowIdxExprEncoding;
106        let id = RowIdxVTable::id(&encoding);
107        assert_eq!(id.as_ref(), "vortex.row_idx");
108    }
109
110    #[test]
111    fn test_vtable_encoding() {
112        let expr = RowIdxExpr;
113        let encoding_ref = RowIdxVTable::encoding(&expr);
114
115        // Check that the encoding ref is the same instance
116        let encoding_ref2 = RowIdxVTable::encoding(&expr);
117        assert!(std::ptr::eq(
118            encoding_ref.as_ref() as *const _,
119            encoding_ref2.as_ref() as *const _
120        ));
121    }
122
123    #[test]
124    fn test_vtable_metadata() {
125        let expr = RowIdxExpr;
126        let metadata = RowIdxVTable::metadata(&expr);
127        assert!(metadata.is_some());
128        // Just check that we get EmptyMetadata back (can't compare directly)
129        let _empty = metadata.unwrap();
130    }
131
132    #[test]
133    fn test_vtable_children() {
134        let expr = RowIdxExpr;
135        let children = RowIdxVTable::children(&expr);
136        assert!(children.is_empty());
137    }
138
139    #[test]
140    fn test_vtable_with_children() {
141        let expr = RowIdxExpr;
142
143        // Should succeed with empty children
144        let result = RowIdxVTable::with_children(&expr, vec![]);
145        assert!(result.is_ok());
146        assert_eq!(result.unwrap(), expr);
147
148        // Should also succeed with non-empty children (it ignores them and returns the same expr)
149        let dummy_expr = row_idx();
150        let result = RowIdxVTable::with_children(&expr, vec![dummy_expr]);
151        assert!(result.is_ok());
152        assert_eq!(result.unwrap(), expr);
153    }
154
155    #[test]
156    fn test_vtable_build_success() {
157        let encoding = RowIdxExprEncoding;
158        let metadata = EmptyMetadata;
159        let result = RowIdxVTable::build(&encoding, &metadata, vec![]);
160        assert!(result.is_ok());
161        assert_eq!(result.unwrap(), RowIdxExpr);
162    }
163
164    #[test]
165    fn test_vtable_build_with_children_fails() {
166        let encoding = RowIdxExprEncoding;
167        let metadata = EmptyMetadata;
168        let dummy_expr = row_idx();
169        let result = RowIdxVTable::build(&encoding, &metadata, vec![dummy_expr]);
170        assert!(result.is_err());
171        let err_msg = format!("{}", result.unwrap_err());
172        assert!(err_msg.contains("does not expect any children"));
173        assert!(err_msg.contains("got 1"));
174    }
175
176    #[test]
177    fn test_vtable_build_with_multiple_children_fails() {
178        let encoding = RowIdxExprEncoding;
179        let metadata = EmptyMetadata;
180        let dummy_expr1 = row_idx();
181        let dummy_expr2 = row_idx();
182        let result = RowIdxVTable::build(&encoding, &metadata, vec![dummy_expr1, dummy_expr2]);
183        assert!(result.is_err());
184        let err_msg = format!("{}", result.unwrap_err());
185        assert!(err_msg.contains("does not expect any children"));
186        assert!(err_msg.contains("got 2"));
187    }
188
189    #[test]
190    fn test_vtable_evaluate_fails() {
191        use vortex_array::IntoArray;
192        use vortex_array::arrays::PrimitiveArray;
193        use vortex_expr::Scope;
194
195        let expr = RowIdxExpr;
196        // Create a dummy array for the scope
197        let array = PrimitiveArray::from_iter([0u64, 1, 2]).into_array();
198        let scope = Scope::new(array);
199        let result = RowIdxVTable::evaluate(&expr, &scope);
200        assert!(result.is_err());
201        let err_msg = format!("{}", result.unwrap_err());
202        assert!(err_msg.contains("should not be evaluated directly"));
203        assert!(err_msg.contains("context of a Vortex scan"));
204    }
205
206    #[test]
207    fn test_vtable_return_dtype() {
208        let expr = RowIdxExpr;
209        let scope = DType::Primitive(PType::U32, Nullability::Nullable);
210        let result = RowIdxVTable::return_dtype(&expr, &scope);
211        assert!(result.is_ok());
212        let dtype = result.unwrap();
213        assert_eq!(
214            dtype,
215            DType::Primitive(PType::U64, Nullability::NonNullable)
216        );
217    }
218
219    #[rstest]
220    #[case(DType::Primitive(PType::U8, Nullability::Nullable))]
221    #[case(DType::Primitive(PType::I64, Nullability::NonNullable))]
222    #[case(DType::Primitive(PType::F32, Nullability::Nullable))]
223    #[case(DType::Utf8(Nullability::NonNullable))]
224    #[case(DType::Binary(Nullability::Nullable))]
225    fn test_vtable_return_dtype_with_various_scopes(#[case] scope_dtype: DType) {
226        let expr = RowIdxExpr;
227        let result = RowIdxVTable::return_dtype(&expr, &scope_dtype);
228        assert!(result.is_ok());
229        // Row index dtype should always be U64 NonNullable regardless of scope
230        assert_eq!(
231            result.unwrap(),
232            DType::Primitive(PType::U64, Nullability::NonNullable)
233        );
234    }
235
236    #[test]
237    fn test_row_idx_function_creates_correct_expr() {
238        let expr = row_idx();
239
240        // Verify it's the right type
241        assert!(expr.is::<RowIdxVTable>());
242
243        // Verify it has no children
244        let children = RowIdxVTable::children(&RowIdxExpr);
245        assert_eq!(children.len(), 0);
246
247        // Verify metadata is present
248        let metadata = RowIdxVTable::metadata(&RowIdxExpr);
249        assert!(metadata.is_some());
250    }
251
252    #[test]
253    fn test_expr_into_expr_conversion() {
254        let expr = RowIdxExpr;
255        let expr_ref: ExprRef = expr.into_expr();
256        assert!(expr_ref.is::<RowIdxVTable>());
257    }
258}
vortex_layout/layouts/row_idx/expr.rs

vortex_layout/layouts/row_idx/
expr.rs
