vortex_array/expr/
proto.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexResult;
6use vortex_error::vortex_err;
7use vortex_proto::expr as pb;
8use vortex_session::VortexSession;
9
10use crate::expr::Expression;
11use crate::scalar_fn::ForeignScalarFnVTable;
12use crate::scalar_fn::ScalarFnId;
13use crate::scalar_fn::session::ScalarFnSessionExt;
14
15pub trait ExprSerializeProtoExt {
16    /// Serialize the expression to its protobuf representation.
17    fn serialize_proto(&self) -> VortexResult<pb::Expr>;
18}
19
20impl ExprSerializeProtoExt for Expression {
21    fn serialize_proto(&self) -> VortexResult<pb::Expr> {
22        let children = self
23            .children()
24            .iter()
25            .map(|child| child.serialize_proto())
26            .try_collect()?;
27
28        let metadata = self.options().serialize()?.ok_or_else(|| {
29            vortex_err!("Expression '{}' is not serializable: {}", self.id(), self)
30        })?;
31
32        Ok(pb::Expr {
33            id: self.id().to_string(),
34            children,
35            metadata: Some(metadata),
36        })
37    }
38}
39
40impl Expression {
41    pub fn from_proto(expr: &pb::Expr, session: &VortexSession) -> VortexResult<Expression> {
42        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
43        let expr_id = ScalarFnId::new(expr.id.as_str());
44        let children = expr
45            .children
46            .iter()
47            .map(|e| Expression::from_proto(e, session))
48            .collect::<VortexResult<Vec<_>>>()?;
49
50        let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().find(&expr_id) {
51            vtable.deserialize(expr.metadata(), session)?
52        } else if session.allows_unknown() {
53            ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len())
54        } else {
55            return Err(vortex_err!("unknown expression id: {}", expr_id));
56        };
57
58        Expression::try_new(scalar_fn, children)
59    }
60}
61
62/// Deserialize a [`Expression`] from the protobuf representation.
63#[deprecated(note = "Use Expression::from_proto instead")]
64pub fn deserialize_expr_proto(
65    expr: &pb::Expr,
66    session: &VortexSession,
67) -> VortexResult<Expression> {
68    Expression::from_proto(expr, session)
69}
70
71#[cfg(test)]
72mod tests {
73    use prost::Message;
74    use vortex_proto::expr as pb;
75    use vortex_session::VortexSession;
76
77    use super::ExprSerializeProtoExt;
78    use crate::array_session;
79    use crate::expr::Expression;
80    use crate::expr::and;
81    use crate::expr::between;
82    use crate::expr::eq;
83    use crate::expr::get_item;
84    use crate::expr::lit;
85    use crate::expr::or;
86    use crate::expr::root;
87    use crate::scalar_fn::fns::between::BetweenOptions;
88    use crate::scalar_fn::fns::between::StrictComparison;
89    use crate::scalar_fn::session::ScalarFnSession;
90
91    #[test]
92    fn expression_serde() {
93        let expr: Expression = or(
94            and(
95                between(
96                    lit(1),
97                    root(),
98                    get_item("a", root()),
99                    BetweenOptions {
100                        lower_strict: StrictComparison::Strict,
101                        upper_strict: StrictComparison::Strict,
102                    },
103                ),
104                lit(1),
105            ),
106            eq(lit(1), root()),
107        );
108
109        let s_expr = expr.serialize_proto().unwrap();
110        let buf = s_expr.encode_to_vec();
111        let s_expr = pb::Expr::decode(buf.as_slice()).unwrap();
112        let deser_expr = Expression::from_proto(&s_expr, &array_session()).unwrap();
113
114        assert_eq!(&deser_expr, &expr);
115    }
116
117    #[test]
118    fn unknown_expression_id_allow_unknown() {
119        let session = VortexSession::empty()
120            .with::<ScalarFnSession>()
121            .allow_unknown();
122
123        let expr_proto = pb::Expr {
124            id: "vortex.test.foreign_scalar_fn".to_string(),
125            metadata: Some(vec![1, 2, 3, 4]),
126            children: vec![root().serialize_proto().unwrap()],
127        };
128
129        let expr = Expression::from_proto(&expr_proto, &session).unwrap();
130        assert_eq!(expr.id().as_ref(), "vortex.test.foreign_scalar_fn");
131
132        let roundtrip = expr.serialize_proto().unwrap();
133        assert_eq!(roundtrip.id, expr_proto.id);
134        assert_eq!(roundtrip.metadata(), expr_proto.metadata());
135        assert_eq!(roundtrip.children.len(), 1);
136    }
137}
vortex_array/expr/proto.rs

vortex_array/expr/
proto.rs
