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;
10
11use crate::expr::Expression;
12use crate::expr::session::ExprRegistry;
13
14pub trait ExprSerializeProtoExt {
15    /// Serialize the expression to its protobuf representation.
16    fn serialize_proto(&self) -> VortexResult<pb::Expr>;
17}
18
19impl ExprSerializeProtoExt for &Expression {
20    fn serialize_proto(&self) -> VortexResult<pb::Expr> {
21        let children = self
22            .children()
23            .iter()
24            .map(|child| child.serialize_proto())
25            .try_collect()?;
26
27        let metadata = self.serialize_metadata()?.ok_or_else(|| {
28            vortex_err!("Expression '{}' is not serializable: {}", self.id(), self)
29        })?;
30
31        Ok(pb::Expr {
32            id: self.id().to_string(),
33            children,
34            metadata: Some(metadata),
35        })
36    }
37}
38
39/// Deserialize a [`Expression`] from the protobuf representation.
40pub fn deserialize_expr_proto(
41    expr: &pb::Expr,
42    registry: &ExprRegistry,
43) -> VortexResult<Expression> {
44    let expr_id = expr.id.as_str();
45    let vtable = registry
46        .find(expr_id)
47        .ok_or_else(|| vortex_err!("unknown expression id: {}", expr_id))?;
48
49    let children = expr
50        .children
51        .iter()
52        .map(|e| deserialize_expr_proto(e, registry))
53        .collect::<VortexResult<Arc<_>>>()?;
54
55    vtable.deserialize(expr.metadata(), children)
56}
57
58#[cfg(test)]
59mod tests {
60    use prost::Message;
61    use vortex_proto::expr as pb;
62
63    use super::ExprSerializeProtoExt;
64    use super::deserialize_expr_proto;
65    use crate::compute::BetweenOptions;
66    use crate::compute::StrictComparison;
67    use crate::expr::Expression;
68    use crate::expr::exprs::between::between;
69    use crate::expr::exprs::binary::and;
70    use crate::expr::exprs::binary::eq;
71    use crate::expr::exprs::binary::or;
72    use crate::expr::exprs::get_item::get_item;
73    use crate::expr::exprs::literal::lit;
74    use crate::expr::exprs::root::root;
75    use crate::expr::session::ExprSession;
76
77    #[test]
78    fn expression_serde() {
79        let registry = ExprSession::default().registry().clone();
80        let expr: Expression = or(
81            and(
82                between(
83                    lit(1),
84                    root(),
85                    get_item("a", root()),
86                    BetweenOptions {
87                        lower_strict: StrictComparison::Strict,
88                        upper_strict: StrictComparison::Strict,
89                    },
90                ),
91                lit(1),
92            ),
93            eq(lit(1), root()),
94        );
95
96        let s_expr = (&expr).serialize_proto().unwrap();
97        let buf = s_expr.encode_to_vec();
98        let s_expr = pb::Expr::decode(buf.as_slice()).unwrap();
99        let deser_expr = deserialize_expr_proto(&s_expr, &registry).unwrap();
100
101        assert_eq!(&deser_expr, &expr);
102    }
103}