vortex_expr/
vtable.rs

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

vortex_expr/
vtable.rs
