vortex_expr/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::hash::Hash;
6use std::ops::Deref;
7
8use vortex_array::operator::OperatorRef;
9use vortex_array::{ArrayRef, DeserializeMetadata, SerializeMetadata};
10use vortex_dtype::DType;
11use vortex_error::VortexResult;
12
13use crate::display::DisplayAs;
14use crate::{
15    AnalysisExpr, ExprEncoding, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VortexExpr,
16};
17
18pub trait VTable: 'static + Sized + Send + Sync + Debug {
19    type Expr: 'static
20        + Send
21        + Sync
22        + Clone
23        + Debug
24        + DisplayAs
25        + PartialEq
26        + Eq
27        + Hash
28        + Deref<Target = dyn VortexExpr>
29        + IntoExpr
30        + AnalysisExpr;
31    type Encoding: 'static + Send + Sync + Deref<Target = dyn ExprEncoding>;
32    type Metadata: SerializeMetadata + DeserializeMetadata + Debug;
33
34    /// Returns the ID of the expr encoding.
35    fn id(encoding: &Self::Encoding) -> ExprId;
36
37    /// Returns the encoding for the expr.
38    fn encoding(expr: &Self::Expr) -> ExprEncodingRef;
39
40    /// Returns the serialize-able metadata for the expr, or `None` if serialization is not
41    /// supported.
42    fn metadata(expr: &Self::Expr) -> Option<Self::Metadata>;
43
44    /// Returns the children of the expr.
45    fn children(expr: &Self::Expr) -> Vec<&ExprRef>;
46
47    /// Return a new instance of the expression with the children replaced.
48    ///
49    /// ## Preconditions
50    ///
51    /// The number of children will match the current number of children in the expression.
52    fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr>;
53
54    /// Construct a new [`VortexExpr`] from the provided parts.
55    fn build(
56        encoding: &Self::Encoding,
57        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
58        children: Vec<ExprRef>,
59    ) -> VortexResult<Self::Expr>;
60
61    /// Evaluate the expression in the given scope.
62    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef>;
63
64    /// Compute the return [`DType`] of the expression if evaluated in the given scope.
65    fn return_dtype(expr: &Self::Expr, scope: &DType) -> VortexResult<DType>;
66
67    fn operator(_expr: &Self::Expr, _scope: &OperatorRef) -> VortexResult<Option<OperatorRef>> {
68        Ok(None)
69    }
70}
71
72#[macro_export]
73macro_rules! vtable {
74    ($V:ident) => {
75        $crate::aliases::paste::paste! {
76            #[derive(Debug)]
77            pub struct [<$V VTable>];
78
79            impl AsRef<dyn $crate::VortexExpr> for [<$V Expr>] {
80                fn as_ref(&self) -> &dyn $crate::VortexExpr {
81                    // We can unsafe cast ourselves to a ExprAdapter.
82                    unsafe { &*(self as *const [<$V Expr>] as *const $crate::ExprAdapter<[<$V VTable>]>) }
83                }
84            }
85
86            impl std::ops::Deref for [<$V Expr>] {
87                type Target = dyn $crate::VortexExpr;
88
89                fn deref(&self) -> &Self::Target {
90                    // We can unsafe cast ourselves to an ExprAdapter.
91                    unsafe { &*(self as *const [<$V Expr>] as *const $crate::ExprAdapter<[<$V VTable>]>) }
92                }
93            }
94
95            impl $crate::IntoExpr for [<$V Expr>] {
96                fn into_expr(self) -> $crate::ExprRef {
97                    // We can unsafe transmute ourselves to an ExprAdapter.
98                    std::sync::Arc::new(unsafe { std::mem::transmute::<[<$V Expr>], $crate::ExprAdapter::<[<$V VTable>]>>(self) })
99                }
100            }
101
102            impl From<[<$V Expr>]> for $crate::ExprRef {
103                fn from(value: [<$V Expr>]) -> $crate::ExprRef {
104                    use $crate::IntoExpr;
105                    value.into_expr()
106                }
107            }
108
109            impl AsRef<dyn $crate::ExprEncoding> for [<$V ExprEncoding>] {
110                fn as_ref(&self) -> &dyn $crate::ExprEncoding {
111                    // We can unsafe cast ourselves to an ExprEncodingAdapter.
112                    unsafe { &*(self as *const [<$V ExprEncoding>] as *const $crate::ExprEncodingAdapter<[<$V VTable>]>) }
113                }
114            }
115
116            impl std::ops::Deref for [<$V ExprEncoding>] {
117                type Target = dyn $crate::ExprEncoding;
118
119                fn deref(&self) -> &Self::Target {
120                    // We can unsafe cast ourselves to an ExprEncodingAdapter.
121                    unsafe { &*(self as *const [<$V ExprEncoding>] as *const $crate::ExprEncodingAdapter<[<$V VTable>]>) }
122                }
123            }
124        }
125    };
126}
127
128#[cfg(test)]
129mod tests {
130    use rstest::{fixture, rstest};
131
132    use super::*;
133    use crate::proto::{ExprSerializeProtoExt, deserialize_expr_proto};
134    use crate::*;
135
136    #[fixture]
137    #[once]
138    fn registry() -> ExprRegistry {
139        ExprRegistry::default()
140    }
141
142    #[rstest]
143    // Root and selection expressions
144    #[case(root())]
145    #[case(select(["hello", "world"], root()))]
146    #[case(select_exclude(["world", "hello"], root()))]
147    // Literal expressions
148    #[case(lit(42i32))]
149    #[case(lit(std::f64::consts::PI))]
150    #[case(lit(true))]
151    #[case(lit("hello"))]
152    // Column access expressions
153    #[case(col("column_name"))]
154    #[case(get_item("field", root()))]
155    // Binary comparison expressions
156    #[case(eq(col("a"), lit(10)))]
157    #[case(not_eq(col("a"), lit(10)))]
158    #[case(gt(col("a"), lit(10)))]
159    #[case(gt_eq(col("a"), lit(10)))]
160    #[case(lt(col("a"), lit(10)))]
161    #[case(lt_eq(col("a"), lit(10)))]
162    // Logical expressions
163    #[case(and(col("a"), col("b")))]
164    #[case(or(col("a"), col("b")))]
165    #[case(not(col("a")))]
166    // Arithmetic expressions
167    #[case(checked_add(col("a"), lit(5)))]
168    // Null check expressions
169    #[case(is_null(col("nullable_col")))]
170    // Type casting expressions
171    #[case(cast(
172        col("a"),
173        DType::Primitive(vortex_dtype::PType::I64, vortex_dtype::Nullability::NonNullable)
174    ))]
175    // Between expressions
176    #[case(between(col("a"), lit(10), lit(20), vortex_array::compute::BetweenOptions { lower_strict: vortex_array::compute::StrictComparison::NonStrict, upper_strict: vortex_array::compute::StrictComparison::NonStrict }))]
177    // List contains expressions
178    #[case(list_contains(col("list_col"), lit("item")))]
179    // Pack expressions - creating struct from fields
180    #[case(pack([("field1", col("a")), ("field2", col("b"))], vortex_dtype::Nullability::NonNullable))]
181    // Merge expressions - merging struct expressions
182    #[case(merge([col("struct1"), col("struct2")], vortex_dtype::Nullability::NonNullable))]
183    // Complex nested expressions
184    #[case(and(gt(col("a"), lit(0)), lt(col("a"), lit(100))))]
185    #[case(or(is_null(col("a")), eq(col("a"), lit(0))))]
186    #[case(not(and(eq(col("status"), lit("active")), gt(col("age"), lit(18)))))]
187    fn text_expr_serde_round_trip(
188        registry: &ExprRegistry,
189        #[case] expr: ExprRef,
190    ) -> anyhow::Result<()> {
191        let serialized_pb = expr.serialize_proto()?;
192        let deserialized_expr = deserialize_expr_proto(&serialized_pb, registry)?;
193
194        assert_eq!(&expr, &deserialized_expr);
195
196        Ok(())
197    }
198}
vortex_expr/vtable.rs

vortex_expr/
vtable.rs
