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        let expr_id = ScalarFnId::new(expr.id.as_str());
43        let children = expr
44            .children
45            .iter()
46            .map(|e| Expression::from_proto(e, session))
47            .collect::<VortexResult<Vec<_>>>()?;
48
49        let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().find(&expr_id) {
50            vtable.deserialize(expr.metadata(), session)?
51        } else if session.allows_unknown() {
52            ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len())
53        } else {
54            return Err(vortex_err!("unknown expression id: {}", expr_id));
55        };
56
57        Expression::try_new(scalar_fn, children)
58    }
59}
60
61/// Deserialize a [`Expression`] from the protobuf representation.
62#[deprecated(note = "Use Expression::from_proto instead")]
63pub fn deserialize_expr_proto(
64    expr: &pb::Expr,
65    session: &VortexSession,
66) -> VortexResult<Expression> {
67    Expression::from_proto(expr, session)
68}
69
70#[cfg(test)]
71mod tests {
72    use prost::Message;
73    use vortex_proto::expr as pb;
74    use vortex_session::VortexSession;
75
76    use super::ExprSerializeProtoExt;
77    use crate::LEGACY_SESSION;
78    use crate::expr::Expression;
79    use crate::expr::and;
80    use crate::expr::between;
81    use crate::expr::eq;
82    use crate::expr::get_item;
83    use crate::expr::lit;
84    use crate::expr::or;
85    use crate::expr::root;
86    use crate::scalar_fn::fns::between::BetweenOptions;
87    use crate::scalar_fn::fns::between::StrictComparison;
88    use crate::scalar_fn::session::ScalarFnSession;
89
90    #[test]
91    fn expression_serde() {
92        let expr: Expression = or(
93            and(
94                between(
95                    lit(1),
96                    root(),
97                    get_item("a", root()),
98                    BetweenOptions {
99                        lower_strict: StrictComparison::Strict,
100                        upper_strict: StrictComparison::Strict,
101                    },
102                ),
103                lit(1),
104            ),
105            eq(lit(1), root()),
106        );
107
108        let s_expr = expr.serialize_proto().unwrap();
109        let buf = s_expr.encode_to_vec();
110        let s_expr = pb::Expr::decode(buf.as_slice()).unwrap();
111        let deser_expr = Expression::from_proto(&s_expr, &LEGACY_SESSION).unwrap();
112
113        assert_eq!(&deser_expr, &expr);
114    }
115
116    #[test]
117    fn unknown_expression_id_allow_unknown() {
118        let session = VortexSession::empty()
119            .with::<ScalarFnSession>()
120            .allow_unknown();
121
122        let expr_proto = pb::Expr {
123            id: "vortex.test.foreign_scalar_fn".to_string(),
124            metadata: Some(vec![1, 2, 3, 4]),
125            children: vec![root().serialize_proto().unwrap()],
126        };
127
128        let expr = Expression::from_proto(&expr_proto, &session).unwrap();
129        assert_eq!(expr.id().as_ref(), "vortex.test.foreign_scalar_fn");
130
131        let roundtrip = expr.serialize_proto().unwrap();
132        assert_eq!(roundtrip.id, expr_proto.id);
133        assert_eq!(roundtrip.metadata(), expr_proto.metadata());
134        assert_eq!(roundtrip.children.len(), 1);
135    }
136}
vortex_array/expr/proto.rs

vortex_array/expr/
proto.rs
