Skip to main content

vortex_array/expr/
proto.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use itertools::Itertools;
7use vortex_error::VortexResult;
8use vortex_error::vortex_err;
9use vortex_proto::expr as pb;
10use vortex_session::VortexSession;
11
12use crate::expr::ExprId;
13use crate::expr::Expression;
14use crate::expr::session::ExprSessionExt;
15
16pub trait ExprSerializeProtoExt {
17    /// Serialize the expression to its protobuf representation.
18    fn serialize_proto(&self) -> VortexResult<pb::Expr>;
19}
20
21impl ExprSerializeProtoExt for Expression {
22    fn serialize_proto(&self) -> VortexResult<pb::Expr> {
23        let children = self
24            .children()
25            .iter()
26            .map(|child| child.serialize_proto())
27            .try_collect()?;
28
29        let metadata = self.options().serialize()?.ok_or_else(|| {
30            vortex_err!("Expression '{}' is not serializable: {}", self.id(), self)
31        })?;
32
33        Ok(pb::Expr {
34            id: self.id().to_string(),
35            children,
36            metadata: Some(metadata),
37        })
38    }
39}
40
41impl Expression {
42    pub fn from_proto(expr: &pb::Expr, session: &VortexSession) -> VortexResult<Expression> {
43        let expr_id = ExprId::new_arc(Arc::from(expr.id.to_string()));
44        let vtable = session
45            .expressions()
46            .registry()
47            .find(&expr_id)
48            .ok_or_else(|| vortex_err!("unknown expression id: {}", expr_id))?;
49
50        let children = expr
51            .children
52            .iter()
53            .map(|e| Expression::from_proto(e, session))
54            .collect::<VortexResult<Vec<_>>>()?;
55
56        Expression::try_new(vtable.deserialize(expr.metadata(), session)?, children)
57    }
58}
59
60/// Deserialize a [`Expression`] from the protobuf representation.
61#[deprecated(note = "Use Expression::from_proto instead")]
62pub fn deserialize_expr_proto(
63    expr: &pb::Expr,
64    session: &VortexSession,
65) -> VortexResult<Expression> {
66    Expression::from_proto(expr, session)
67}
68
69#[cfg(test)]
70mod tests {
71    use prost::Message;
72    use vortex_proto::expr as pb;
73
74    use super::ExprSerializeProtoExt;
75    use crate::LEGACY_SESSION;
76    use crate::expr::BetweenOptions;
77    use crate::expr::Expression;
78    use crate::expr::StrictComparison;
79    use crate::expr::exprs::between::between;
80    use crate::expr::exprs::binary::and;
81    use crate::expr::exprs::binary::eq;
82    use crate::expr::exprs::binary::or;
83    use crate::expr::exprs::get_item::get_item;
84    use crate::expr::exprs::literal::lit;
85    use crate::expr::exprs::root::root;
86
87    #[test]
88    fn expression_serde() {
89        let expr: Expression = or(
90            and(
91                between(
92                    lit(1),
93                    root(),
94                    get_item("a", root()),
95                    BetweenOptions {
96                        lower_strict: StrictComparison::Strict,
97                        upper_strict: StrictComparison::Strict,
98                    },
99                ),
100                lit(1),
101            ),
102            eq(lit(1), root()),
103        );
104
105        let s_expr = expr.serialize_proto().unwrap();
106        let buf = s_expr.encode_to_vec();
107        let s_expr = pb::Expr::decode(buf.as_slice()).unwrap();
108        let deser_expr = Expression::from_proto(&s_expr, &LEGACY_SESSION).unwrap();
109
110        assert_eq!(&deser_expr, &expr);
111    }
112}