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}
83
84#[cfg(test)]
85mod tests {
86    use rstest::rstest;
87    use vortex_array::EmptyMetadata;
88    use vortex_dtype::{DType, Nullability, PType};
89    use vortex_expr::{ExprRef, IntoExpr, VTable};
90
91    use crate::layouts::row_idx::{RowIdxExpr, RowIdxExprEncoding, RowIdxVTable, row_idx};
92
93    #[test]
94    fn test_row_idx_expr_creation() {
95        let expr = row_idx();
96        assert!(expr.is::<RowIdxVTable>());
97    }
98
99    #[test]
100    fn test_vtable_id() {
101        let encoding = RowIdxExprEncoding;
102        let id = RowIdxVTable::id(&encoding);
103        assert_eq!(id.as_ref(), "vortex.row_idx");
104    }
105
106    #[test]
107    fn test_vtable_encoding() {
108        let expr = RowIdxExpr;
109        let encoding_ref = RowIdxVTable::encoding(&expr);
110
111        // Check that the encoding ref is the same instance
112        let encoding_ref2 = RowIdxVTable::encoding(&expr);
113        assert!(std::ptr::eq(
114            encoding_ref.as_ref() as *const _,
115            encoding_ref2.as_ref() as *const _
116        ));
117    }
118
119    #[test]
120    fn test_vtable_metadata() {
121        let expr = RowIdxExpr;
122        let metadata = RowIdxVTable::metadata(&expr);
123        assert!(metadata.is_some());
124        // Just check that we get EmptyMetadata back (can't compare directly)
125        let _empty = metadata.unwrap();
126    }
127
128    #[test]
129    fn test_vtable_children() {
130        let expr = RowIdxExpr;
131        let children = RowIdxVTable::children(&expr);
132        assert!(children.is_empty());
133    }
134
135    #[test]
136    fn test_vtable_with_children() {
137        let expr = RowIdxExpr;
138
139        // Should succeed with empty children
140        let result = RowIdxVTable::with_children(&expr, vec![]);
141        assert!(result.is_ok());
142        assert_eq!(result.unwrap(), expr);
143
144        // Should also succeed with non-empty children (it ignores them and returns the same expr)
145        let dummy_expr = row_idx();
146        let result = RowIdxVTable::with_children(&expr, vec![dummy_expr]);
147        assert!(result.is_ok());
148        assert_eq!(result.unwrap(), expr);
149    }
150
151    #[test]
152    fn test_vtable_build_success() {
153        let encoding = RowIdxExprEncoding;
154        let metadata = EmptyMetadata;
155        let result = RowIdxVTable::build(&encoding, &metadata, vec![]);
156        assert!(result.is_ok());
157        assert_eq!(result.unwrap(), RowIdxExpr);
158    }
159
160    #[test]
161    fn test_vtable_build_with_children_fails() {
162        let encoding = RowIdxExprEncoding;
163        let metadata = EmptyMetadata;
164        let dummy_expr = row_idx();
165        let result = RowIdxVTable::build(&encoding, &metadata, vec![dummy_expr]);
166        assert!(result.is_err());
167        let err_msg = format!("{}", result.unwrap_err());
168        assert!(err_msg.contains("does not expect any children"));
169        assert!(err_msg.contains("got 1"));
170    }
171
172    #[test]
173    fn test_vtable_build_with_multiple_children_fails() {
174        let encoding = RowIdxExprEncoding;
175        let metadata = EmptyMetadata;
176        let dummy_expr1 = row_idx();
177        let dummy_expr2 = row_idx();
178        let result = RowIdxVTable::build(&encoding, &metadata, vec![dummy_expr1, dummy_expr2]);
179        assert!(result.is_err());
180        let err_msg = format!("{}", result.unwrap_err());
181        assert!(err_msg.contains("does not expect any children"));
182        assert!(err_msg.contains("got 2"));
183    }
184
185    #[test]
186    fn test_vtable_evaluate_fails() {
187        use vortex_array::IntoArray;
188        use vortex_array::arrays::PrimitiveArray;
189        use vortex_expr::Scope;
190
191        let expr = RowIdxExpr;
192        // Create a dummy array for the scope
193        let array = PrimitiveArray::from_iter([0u64, 1, 2]).into_array();
194        let scope = Scope::new(array);
195        let result = RowIdxVTable::evaluate(&expr, &scope);
196        assert!(result.is_err());
197        let err_msg = format!("{}", result.unwrap_err());
198        assert!(err_msg.contains("should not be evaluated directly"));
199        assert!(err_msg.contains("context of a Vortex scan"));
200    }
201
202    #[test]
203    fn test_vtable_return_dtype() {
204        let expr = RowIdxExpr;
205        let scope = DType::Primitive(PType::U32, Nullability::Nullable);
206        let result = RowIdxVTable::return_dtype(&expr, &scope);
207        assert!(result.is_ok());
208        let dtype = result.unwrap();
209        assert_eq!(
210            dtype,
211            DType::Primitive(PType::U64, Nullability::NonNullable)
212        );
213    }
214
215    #[rstest]
216    #[case(DType::Primitive(PType::U8, Nullability::Nullable))]
217    #[case(DType::Primitive(PType::I64, Nullability::NonNullable))]
218    #[case(DType::Primitive(PType::F32, Nullability::Nullable))]
219    #[case(DType::Utf8(Nullability::NonNullable))]
220    #[case(DType::Binary(Nullability::Nullable))]
221    fn test_vtable_return_dtype_with_various_scopes(#[case] scope_dtype: DType) {
222        let expr = RowIdxExpr;
223        let result = RowIdxVTable::return_dtype(&expr, &scope_dtype);
224        assert!(result.is_ok());
225        // Row index dtype should always be U64 NonNullable regardless of scope
226        assert_eq!(
227            result.unwrap(),
228            DType::Primitive(PType::U64, Nullability::NonNullable)
229        );
230    }
231
232    #[test]
233    fn test_row_idx_function_creates_correct_expr() {
234        let expr = row_idx();
235
236        // Verify it's the right type
237        assert!(expr.is::<RowIdxVTable>());
238
239        // Verify it has no children
240        let children = RowIdxVTable::children(&RowIdxExpr);
241        assert_eq!(children.len(), 0);
242
243        // Verify metadata is present
244        let metadata = RowIdxVTable::metadata(&RowIdxExpr);
245        assert!(metadata.is_some());
246    }
247
248    #[test]
249    fn test_expr_into_expr_conversion() {
250        let expr = RowIdxExpr;
251        let expr_ref: ExprRef = expr.into_expr();
252        assert!(expr_ref.is::<RowIdxVTable>());
253    }
254}
vortex_layout/layouts/row_idx/expr.rs

vortex_layout/layouts/row_idx/
expr.rs
