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

vortex_expr/
vtable.rs
