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_ensure;
7use vortex_error::vortex_err;
8use vortex_proto::expr as pb;
9use vortex_session::VortexSession;
10
11use crate::expr::Expression;
12use crate::scalar_fn::ForeignScalarFnVTable;
13use crate::scalar_fn::ScalarFnId;
14use crate::scalar_fn::session::ScalarFnSessionExt;
15
16pub trait ExprSerializeProtoExt {
17    /// Serialize the expression to its protobuf representation.
18    fn serialize_proto(&self) -> VortexResult<pb::Expr>;
19}
20
21/// The wire id for [`Expression::Root`], retained from when `Root` was a scalar function so that
22/// already-serialized expressions keep round-tripping.
23pub(crate) const ROOT_ID: &str = "vortex.root";
24
25impl ExprSerializeProtoExt for Expression {
26    fn serialize_proto(&self) -> VortexResult<pb::Expr> {
27        let Some(scalar_fn) = self.as_scalar() else {
28            return Ok(pb::Expr {
29                id: ROOT_ID.to_string(),
30                children: vec![],
31                metadata: Some(vec![]),
32            });
33        };
34
35        let children = self
36            .children()
37            .iter()
38            .map(|child| child.serialize_proto())
39            .try_collect()?;
40
41        let metadata = scalar_fn.options().serialize()?.ok_or_else(|| {
42            vortex_err!(
43                "Expression '{}' is not serializable: {}",
44                scalar_fn.id(),
45                self
46            )
47        })?;
48
49        Ok(pb::Expr {
50            id: scalar_fn.id().to_string(),
51            children,
52            metadata: Some(metadata),
53        })
54    }
55}
56
57impl Expression {
58    pub fn from_proto(expr: &pb::Expr, session: &VortexSession) -> VortexResult<Expression> {
59        // Root is not a registered scalar fn, so it must be resolved before the registry lookup.
60        if expr.id == ROOT_ID {
61            vortex_ensure!(
62                expr.children.is_empty(),
63                "root expression must have no children, got {}",
64                expr.children.len()
65            );
66            return Ok(Expression::Root);
67        }
68
69        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
70        let expr_id = ScalarFnId::new(expr.id.as_str());
71        let children = expr
72            .children
73            .iter()
74            .map(|e| Expression::from_proto(e, session))
75            .collect::<VortexResult<Vec<_>>>()?;
76
77        let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().get(&expr_id) {
78            vtable.deserialize(expr.metadata(), session)?
79        } else if session.allows_unknown() {
80            ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len())
81        } else {
82            return Err(vortex_err!("unknown expression id: {}", expr_id));
83        };
84
85        Expression::try_new(scalar_fn, children)
86    }
87}
88
89/// Deserialize a [`Expression`] from the protobuf representation.
90#[deprecated(note = "Use Expression::from_proto instead")]
91pub fn deserialize_expr_proto(
92    expr: &pb::Expr,
93    session: &VortexSession,
94) -> VortexResult<Expression> {
95    Expression::from_proto(expr, session)
96}
97
98#[cfg(test)]
99mod tests {
100    use prost::Message;
101    use vortex_proto::expr as pb;
102    use vortex_session::VortexSession;
103
104    use super::ExprSerializeProtoExt;
105    use crate::array_session;
106    use crate::expr::Expression;
107    use crate::expr::and;
108    use crate::expr::between;
109    use crate::expr::eq;
110    use crate::expr::get_item;
111    use crate::expr::lit;
112    use crate::expr::or;
113    use crate::expr::root;
114    use crate::scalar_fn::fns::between::BetweenOptions;
115    use crate::scalar_fn::fns::between::StrictComparison;
116    use crate::scalar_fn::session::ScalarFnSession;
117
118    #[test]
119    fn expression_serde() {
120        let expr: Expression = or(
121            and(
122                between(
123                    lit(1),
124                    root(),
125                    get_item("a", root()),
126                    BetweenOptions {
127                        lower_strict: StrictComparison::Strict,
128                        upper_strict: StrictComparison::Strict,
129                    },
130                ),
131                lit(1),
132            ),
133            eq(lit(1), root()),
134        );
135
136        let s_expr = expr.serialize_proto().unwrap();
137        let buf = s_expr.encode_to_vec();
138        let s_expr = pb::Expr::decode(buf.as_slice()).unwrap();
139        let deser_expr = Expression::from_proto(&s_expr, &array_session()).unwrap();
140
141        assert_eq!(&deser_expr, &expr);
142    }
143
144    #[test]
145    fn unknown_expression_id_allow_unknown() {
146        let session = VortexSession::empty().with::<ScalarFnSession>();
147        session.allow_unknown();
148
149        let expr_proto = pb::Expr {
150            id: "vortex.test.foreign_scalar_fn".to_string(),
151            metadata: Some(vec![1, 2, 3, 4]),
152            children: vec![root().serialize_proto().unwrap()],
153        };
154
155        let expr = Expression::from_proto(&expr_proto, &session).unwrap();
156        assert_eq!(
157            expr.as_scalar().map(|f| f.id().as_ref().to_string()),
158            Some("vortex.test.foreign_scalar_fn".to_string())
159        );
160
161        let roundtrip = expr.serialize_proto().unwrap();
162        assert_eq!(roundtrip.id, expr_proto.id);
163        assert_eq!(roundtrip.metadata(), expr_proto.metadata());
164        assert_eq!(roundtrip.children.len(), 1);
165    }
166}
vortex_array/expr/proto.rs

vortex_array/expr/
proto.rs
