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_buffer::buffer;
93    use vortex_dtype::{DType, Nullability, PType};
94    use vortex_expr::{ExprRef, IntoExpr, VTable};
95
96    use crate::layouts::row_idx::{RowIdxExpr, RowIdxExprEncoding, RowIdxVTable, row_idx};
97
98    #[test]
99    fn test_row_idx_expr_creation() {
100        let expr = row_idx();
101        assert!(expr.is::<RowIdxVTable>());
102    }
103
104    #[test]
105    fn test_vtable_id() {
106        let encoding = RowIdxExprEncoding;
107        let id = RowIdxVTable::id(&encoding);
108        assert_eq!(id.as_ref(), "vortex.row_idx");
109    }
110
111    #[test]
112    fn test_vtable_encoding() {
113        let expr = RowIdxExpr;
114        let encoding_ref = RowIdxVTable::encoding(&expr);
115
116        // Check that the encoding ref is the same instance
117        let encoding_ref2 = RowIdxVTable::encoding(&expr);
118        assert!(std::ptr::eq(
119            encoding_ref.as_ref() as *const _,
120            encoding_ref2.as_ref() as *const _
121        ));
122    }
123
124    #[test]
125    fn test_vtable_metadata() {
126        let expr = RowIdxExpr;
127        let metadata = RowIdxVTable::metadata(&expr);
128        assert!(metadata.is_some());
129        // Just check that we get EmptyMetadata back (can't compare directly)
130        let _empty = metadata.unwrap();
131    }
132
133    #[test]
134    fn test_vtable_children() {
135        let expr = RowIdxExpr;
136        let children = RowIdxVTable::children(&expr);
137        assert!(children.is_empty());
138    }
139
140    #[test]
141    fn test_vtable_with_children() {
142        let expr = RowIdxExpr;
143
144        // Should succeed with empty children
145        let result = RowIdxVTable::with_children(&expr, vec![]);
146        assert!(result.is_ok());
147        assert_eq!(result.unwrap(), expr);
148
149        // Should also succeed with non-empty children (it ignores them and returns the same expr)
150        let dummy_expr = row_idx();
151        let result = RowIdxVTable::with_children(&expr, vec![dummy_expr]);
152        assert!(result.is_ok());
153        assert_eq!(result.unwrap(), expr);
154    }
155
156    #[test]
157    fn test_vtable_build_success() {
158        let encoding = RowIdxExprEncoding;
159        let metadata = EmptyMetadata;
160        let result = RowIdxVTable::build(&encoding, &metadata, vec![]);
161        assert!(result.is_ok());
162        assert_eq!(result.unwrap(), RowIdxExpr);
163    }
164
165    #[test]
166    fn test_vtable_build_with_children_fails() {
167        let encoding = RowIdxExprEncoding;
168        let metadata = EmptyMetadata;
169        let dummy_expr = row_idx();
170        let result = RowIdxVTable::build(&encoding, &metadata, vec![dummy_expr]);
171        assert!(result.is_err());
172        let err_msg = format!("{}", result.unwrap_err());
173        assert!(err_msg.contains("does not expect any children"));
174        assert!(err_msg.contains("got 1"));
175    }
176
177    #[test]
178    fn test_vtable_build_with_multiple_children_fails() {
179        let encoding = RowIdxExprEncoding;
180        let metadata = EmptyMetadata;
181        let dummy_expr1 = row_idx();
182        let dummy_expr2 = row_idx();
183        let result = RowIdxVTable::build(&encoding, &metadata, vec![dummy_expr1, dummy_expr2]);
184        assert!(result.is_err());
185        let err_msg = format!("{}", result.unwrap_err());
186        assert!(err_msg.contains("does not expect any children"));
187        assert!(err_msg.contains("got 2"));
188    }
189
190    #[test]
191    fn test_vtable_evaluate_fails() {
192        use vortex_array::IntoArray;
193        use vortex_expr::Scope;
194
195        let expr = RowIdxExpr;
196        // Create a dummy array for the scope
197        let array = buffer![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
