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

vortex_array/expr/
proto.rs
