spiral/scandal/
expression.rs

1//! Bounded serialization of unbound Vortex expressions.
2//!
3//! An [`Expression`] is the contract's base64 protobuf text, so it travels unchanged through
4//! HTTP and the C ABI. Decoding checks size, node and depth bounds before any recursive
5//! function deserialization or type binding.
6
7use std::fmt::Display;
8
9use base64::Engine as _;
10use base64::engine::general_purpose::STANDARD;
11use prost::Message as _;
12use vortex_array::expr::proto::ExprSerializeProtoExt as _;
13use vortex_array::expr::{Expression as VortexExpression, nested_case_when};
14use vortex_array::proto::expr::Expr;
15use vortex_array::scalar_fn::{ForeignScalarFnVTable, ScalarFnId};
16use vortex_error::{VortexResult, vortex_ensure};
17use vortex_session::VortexSession;
18
19pub use super::api::{Expression, NamedExpression};
20use super::error::ScandalError;
21
22const MAX_NODES: usize = 4096;
23const MAX_DEPTH: usize = 64;
24const MAX_METADATA_BYTES: usize = 256 * 1024;
25const MAX_ENCODED_BYTES: usize = 512 * 1024;
26
27/// Build a `CASE` expression from `(condition, value)` branches and an optional `ELSE`.
28///
29/// Returns an error if there are no branches or their child count exceeds the wire format's limit.
30pub fn case_when(
31    branches: Vec<(VortexExpression, VortexExpression)>,
32    otherwise: Option<VortexExpression>,
33) -> VortexResult<VortexExpression> {
34    vortex_ensure!(!branches.is_empty(), "CASE needs at least one WHEN branch");
35    vortex_ensure!(
36        branches.len() <= (u32::MAX / 2) as usize,
37        "CASE has too many branches"
38    );
39    Ok(nested_case_when(branches, otherwise))
40}
41
42/// An unbound SPQL call as a Vortex expression node, for composing with built-in expressions.
43///
44/// Serialize the composed tree with [`Expression::new`]. The client does not link SPQL
45/// implementations, so the node has no return type until the execution host binds it.
46pub fn spql_call(
47    id: &str,
48    arguments: Vec<VortexExpression>,
49) -> Result<VortexExpression, ScandalError> {
50    let call =
51        ForeignScalarFnVTable::make_scalar_fn(ScalarFnId::new(id), Vec::new(), arguments.len());
52    VortexExpression::try_new(call, arguments).map_err(invalid)
53}
54
55impl Expression {
56    /// Construct an unbound SPQL call without linking its private implementation.
57    ///
58    /// Pass literal arguments as literal expressions. The execution host resolves the canonical
59    /// SPQL ID, selects the overload that accepts the argument types and binds it.
60    pub fn call(id: impl Into<String>, arguments: Vec<Self>) -> Result<Self, ScandalError> {
61        let children = arguments
62            .iter()
63            .map(Self::proto)
64            .collect::<Result<Vec<_>, _>>()?;
65        Self::from_proto(Expr {
66            id: id.into(),
67            metadata: Some(Vec::new()),
68            children,
69        })
70    }
71
72    /// The conjunction of this expression and `other`: rows both accept.
73    pub fn and(&self, other: &Self) -> Result<Self, ScandalError> {
74        let mut conjunction =
75            vortex_array::expr::and(vortex_array::expr::lit(true), vortex_array::expr::lit(true))
76                .serialize_proto()
77                .map_err(invalid)?;
78        conjunction.children = vec![self.proto()?, other.proto()?];
79        Self::from_proto(conjunction)
80    }
81
82    /// Serialize a Vortex expression after checking its structural work bounds.
83    pub fn new(expression: VortexExpression) -> Result<Self, ScandalError> {
84        let mut pending = vec![(&expression, 1usize)];
85        let mut nodes = 0;
86        while let Some((node, depth)) = pending.pop() {
87            nodes += 1;
88            validate_shape(nodes, pending.len(), node.children().len(), depth)?;
89            pending.extend(node.children().iter().map(|child| (child, depth + 1)));
90        }
91        let proto = expression.serialize_proto().map_err(invalid)?;
92        Self::from_proto(proto)
93    }
94
95    /// Decode using the execution host's function and data-type registries.
96    ///
97    /// This does not coerce or bind types. Unknown function IDs fail unless the
98    /// supplied session explicitly permits them. Callers must separately check
99    /// execution capabilities and bind against the exact input schema.
100    pub fn to_vortex(&self, session: &VortexSession) -> Result<VortexExpression, ScandalError> {
101        VortexExpression::from_proto(&self.proto()?, session).map_err(invalid)
102    }
103
104    /// The unbound Vortex protobuf, for a binary transport.
105    pub fn to_bytes(&self) -> Result<Vec<u8>, ScandalError> {
106        Ok(self.proto()?.encode_to_vec())
107    }
108
109    /// Size of the encoded protobuf, without decoding it.
110    pub fn encoded_len(&self) -> usize {
111        let padding = self
112            .0
113            .bytes()
114            .rev()
115            .take_while(|byte| *byte == b'=')
116            .count();
117        (self.0.len() / 4 * 3).saturating_sub(padding)
118    }
119
120    /// Wrap Vortex protobuf bytes, rejecting unbounded expression work.
121    pub fn from_bytes(bytes: &[u8]) -> Result<Self, ScandalError> {
122        if bytes.len() > MAX_ENCODED_BYTES {
123            return Err(invalid("expression protobuf exceeds 512 KiB"));
124        }
125        Self::from_proto(Expr::decode(bytes).map_err(invalid)?)
126    }
127
128    /// The bounded protobuf this text encodes.
129    fn proto(&self) -> Result<Expr, ScandalError> {
130        if self.0.len() > MAX_ENCODED_BYTES.div_ceil(3) * 4 {
131            return Err(invalid("expression protobuf exceeds 512 KiB"));
132        }
133        let bytes = STANDARD.decode(&self.0).map_err(invalid)?;
134        let proto = Expr::decode(bytes.as_slice()).map_err(invalid)?;
135        check_bounds(&proto)?;
136        Ok(proto)
137    }
138
139    fn from_proto(proto: Expr) -> Result<Self, ScandalError> {
140        check_bounds(&proto)?;
141        Ok(Self(STANDARD.encode(proto.encode_to_vec())))
142    }
143}
144
145impl NamedExpression {
146    /// Name the output of `expression`.
147    pub fn new(name: impl Into<String>, expression: Expression) -> Self {
148        Self {
149            name: name.into(),
150            expression,
151        }
152    }
153}
154
155fn check_bounds(proto: &Expr) -> Result<(), ScandalError> {
156    let mut pending = vec![(proto, 1usize)];
157    let mut nodes = 0;
158    let mut metadata_bytes = 0usize;
159    while let Some((node, depth)) = pending.pop() {
160        nodes += 1;
161        validate_shape(nodes, pending.len(), node.children.len(), depth)?;
162        metadata_bytes = metadata_bytes
163            .saturating_add(node.id.len())
164            .saturating_add(node.metadata.as_ref().map_or(0, Vec::len));
165        if metadata_bytes > MAX_METADATA_BYTES {
166            return Err(invalid("expression metadata exceeds 256 KiB"));
167        }
168        pending.extend(node.children.iter().map(|child| (child, depth + 1)));
169    }
170    if proto.encoded_len() > MAX_ENCODED_BYTES {
171        return Err(invalid("expression protobuf exceeds 512 KiB"));
172    }
173    Ok(())
174}
175
176fn validate_shape(
177    nodes: usize,
178    pending: usize,
179    children: usize,
180    depth: usize,
181) -> Result<(), ScandalError> {
182    if nodes.saturating_add(pending).saturating_add(children) > MAX_NODES || depth > MAX_DEPTH {
183        return Err(invalid("expression exceeds 4096 nodes or 64 levels"));
184    }
185    Ok(())
186}
187
188fn invalid(error: impl Display) -> ScandalError {
189    ScandalError::InvalidArgument {
190        message: error.to_string(),
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use vortex_array::expr::{and, col, eq, gt, lit, not};
197    use vortex_array::proto::expr::Expr;
198    use vortex_session::VortexSession;
199
200    use super::{Expression, case_when, spql_call};
201
202    #[test]
203    fn case_when_round_trips_with_the_default_vortex_session() {
204        let branches = vec![
205            (eq(col("a"), lit(1i64)), lit("one")),
206            (eq(col("a"), lit(2i64)), lit("two")),
207        ];
208        for otherwise in [None, Some(lit("other"))] {
209            let original = case_when(branches.clone(), otherwise).unwrap();
210            let expression = Expression::new(original.clone()).unwrap();
211            assert_eq!(
212                expression.to_vortex(&VortexSession::empty()).unwrap(),
213                original
214            );
215        }
216        assert!(case_when(Vec::new(), Some(lit(1i64))).is_err());
217    }
218
219    #[test]
220    fn a_conjunction_decodes_to_both_terms() {
221        let left = eq(col("episode"), lit(7u64));
222        let right = gt(col("log_time"), lit(100i64));
223        let conjunction = Expression::new(left.clone())
224            .unwrap()
225            .and(&Expression::new(right.clone()).unwrap())
226            .unwrap();
227        assert_eq!(
228            conjunction.to_vortex(&VortexSession::empty()).unwrap(),
229            and(left, right)
230        );
231    }
232
233    #[test]
234    fn vortex_expression_round_trips_and_rejects_unbounded_or_invalid_wire() {
235        let original = eq(col("value"), lit(u64::MAX));
236        let expression = Expression::new(original.clone()).unwrap();
237        let decoded: Expression =
238            serde_json::from_value(serde_json::to_value(&expression).unwrap()).unwrap();
239        assert_eq!(
240            decoded.to_vortex(&VortexSession::empty()).unwrap(),
241            original
242        );
243        assert_eq!(
244            Expression::from_bytes(&expression.to_bytes().unwrap()).unwrap(),
245            expression
246        );
247        assert!(Expression::from_bytes(&[0xff]).is_err());
248        let unknown = Expression::from_proto(Expr {
249            id: "unregistered.function".into(),
250            ..Default::default()
251        })
252        .unwrap();
253        assert!(unknown.to_vortex(&VortexSession::empty()).is_err());
254        assert!(
255            Expression("not base64".into())
256                .to_vortex(&VortexSession::empty())
257                .is_err()
258        );
259        let mut deep = lit(true);
260        for _ in 0..64 {
261            deep = not(deep);
262        }
263        assert!(Expression::new(deep).is_err());
264        let wide = Expr {
265            id: "vortex.root".into(),
266            children: vec![Expr::default(); 4096],
267            ..Default::default()
268        };
269        assert!(Expression::from_proto(wide).is_err());
270        assert!(Expression::new(lit("x".repeat(256 * 1024))).is_err());
271    }
272
273    #[test]
274    fn composed_spql_call_serializes_like_a_direct_call() {
275        let composed = spql_call("spql.str.upper", vec![col("label")]).unwrap();
276        let direct = Expression::call(
277            "spql.str.upper",
278            vec![Expression::new(col("label")).unwrap()],
279        )
280        .unwrap();
281        assert_eq!(Expression::new(composed.clone()).unwrap(), direct);
282        let nested = Expression::new(not(eq(composed, lit("X")))).unwrap();
283        assert_eq!(nested.to_bytes().unwrap().len(), nested.encoded_len());
284    }
285}
spiral/scandal/expression.rs

spiral/scandal/
expression.rs
