Expand the Python expression API

docs/api/python/expr.rst

@@ -11,25 +11,174 @@ the following expression represents the set of rows for which the `age` column l
    >>> age = vortex.expr.column("age")
    >>> (23 > age) & (age < 55)  # doctest: +SKIP
 
+Expressions are picklable, so a filter built in one process can be sent to another (for example to a
+``multiprocessing`` worker or a Ray task). Pickling uses the same protobuf wire format exposed by
+:meth:`vortex.expr.Expr.serialize` and :func:`vortex.expr.deserialize`.
+
 .. autosummary::
    :nosignatures:
 
-   ~vortex.expr.column
    ~vortex.expr.Expr
+   ~vortex.expr.root
+   ~vortex.expr.column
+   ~vortex.expr.literal
+   ~vortex.expr.get_item
+   ~vortex.expr.not_
+   ~vortex.expr.and_
+   ~vortex.expr.or_
+   ~vortex.expr.and_collect
+   ~vortex.expr.or_collect
+   ~vortex.expr.eq
+   ~vortex.expr.not_eq
+   ~vortex.expr.gt
+   ~vortex.expr.gt_eq
+   ~vortex.expr.lt
+   ~vortex.expr.lt_eq
+   ~vortex.expr.add
+   ~vortex.expr.sub
+   ~vortex.expr.mul
+   ~vortex.expr.div
+   ~vortex.expr.between
+   ~vortex.expr.is_null
+   ~vortex.expr.is_not_null
+   ~vortex.expr.fill_null
+   ~vortex.expr.like
+   ~vortex.expr.ilike
+   ~vortex.expr.not_like
+   ~vortex.expr.not_ilike
+   ~vortex.expr.byte_length
+   ~vortex.expr.select
+   ~vortex.expr.select_exclude
+   ~vortex.expr.pack
+   ~vortex.expr.merge
+   ~vortex.expr.list_contains
+   ~vortex.expr.list_length
+   ~vortex.expr.list_sum
+   ~vortex.expr.case_when
+   ~vortex.expr.zip_
+   ~vortex.expr.mask
+   ~vortex.expr.cast
+   ~vortex.expr.ext_storage
+   ~vortex.expr.variant_get
+   ~vortex.expr.deserialize
 
 .. raw:: html
 
    <hr>
 
+Leaves and scope
+----------------
+
+.. autofunction:: vortex.expr.root
+
 .. autofunction:: vortex.expr.column
 
+.. autofunction:: vortex.expr.literal
+
+.. autofunction:: vortex.expr.get_item
+
+Boolean logic
+-------------
+
 .. autofunction:: vortex.expr.not_
 
 .. autofunction:: vortex.expr.and_
 
-.. autofunction:: vortex.expr.root
+.. autofunction:: vortex.expr.or_
 
-.. autofunction:: vortex.expr.literal
+.. autofunction:: vortex.expr.and_collect
+
+.. autofunction:: vortex.expr.or_collect
+
+Comparisons and arithmetic
+--------------------------
+
+.. autofunction:: vortex.expr.eq
+
+.. autofunction:: vortex.expr.not_eq
+
+.. autofunction:: vortex.expr.gt
+
+.. autofunction:: vortex.expr.gt_eq
+
+.. autofunction:: vortex.expr.lt
+
+.. autofunction:: vortex.expr.lt_eq
+
+.. autofunction:: vortex.expr.add
+
+.. autofunction:: vortex.expr.sub
+
+.. autofunction:: vortex.expr.mul
+
+.. autofunction:: vortex.expr.div
+
+.. autofunction:: vortex.expr.between
+
+Nullability
+-----------
+
+.. autofunction:: vortex.expr.is_null
+
+.. autofunction:: vortex.expr.is_not_null
+
+.. autofunction:: vortex.expr.fill_null
+
+Strings
+-------
+
+.. autofunction:: vortex.expr.like
+
+.. autofunction:: vortex.expr.ilike
+
+.. autofunction:: vortex.expr.not_like
+
+.. autofunction:: vortex.expr.not_ilike
+
+.. autofunction:: vortex.expr.byte_length
+
+Structs
+-------
+
+.. autofunction:: vortex.expr.select
+
+.. autofunction:: vortex.expr.select_exclude
+
+.. autofunction:: vortex.expr.pack
+
+.. autofunction:: vortex.expr.merge
+
+Lists
+-----
+
+.. autofunction:: vortex.expr.list_contains
+
+.. autofunction:: vortex.expr.list_length
+
+.. autofunction:: vortex.expr.list_sum
+
+Conditionals and conversions
+----------------------------
+
+.. autofunction:: vortex.expr.case_when
+
+.. autofunction:: vortex.expr.zip_
+
+.. autofunction:: vortex.expr.mask
+
+.. autofunction:: vortex.expr.cast
+
+.. autofunction:: vortex.expr.ext_storage
+
+.. autofunction:: vortex.expr.variant_get
+
+Serialization
+-------------
+
+.. autofunction:: vortex.expr.deserialize
+
+The expression class
+--------------------
 
 .. autoclass:: vortex.expr.Expr
    :members:

vortex-array/src/expr/proto.rs

@@ -98,6 +98,8 @@ pub fn deserialize_expr_proto(
 #[cfg(test)]
 mod tests {
     use prost::Message;
+    use rstest::rstest;
+    use vortex_error::VortexResult;
     use vortex_proto::expr as pb;
     use vortex_session::VortexSession;
 
@@ -106,11 +108,14 @@ mod tests {
     use crate::expr::Expression;
     use crate::expr::and;
     use crate::expr::between;
+    use crate::expr::byte_length;
     use crate::expr::eq;
     use crate::expr::get_item;
     use crate::expr::lit;
+    use crate::expr::mask;
     use crate::expr::or;
     use crate::expr::root;
+    use crate::expr::zip_expr;
     use crate::scalar_fn::fns::between::BetweenOptions;
     use crate::scalar_fn::fns::between::StrictComparison;
     use crate::scalar_fn::session::ScalarFnSession;
@@ -141,6 +146,20 @@ mod tests {
         assert_eq!(&deser_expr, &expr);
     }
 
+    /// `ByteLength`, `Mask` and `Zip` implement `serialize`/`deserialize` but were once missing from
+    /// `ScalarFnSession::default()`, so they serialized fine and then failed to deserialize with
+    /// "unknown expression id".
+    #[rstest]
+    #[case::byte_length(byte_length(root()))]
+    #[case::mask(mask(root(), lit(true)))]
+    #[case::zip(zip_expr(lit(true), root(), lit(0)))]
+    fn round_trips_through_proto(#[case] expr: Expression) -> VortexResult<()> {
+        let buf = expr.serialize_proto()?.encode_to_vec();
+        let decoded = pb::Expr::decode(buf.as_slice())?;
+        assert_eq!(Expression::from_proto(&decoded, &array_session())?, expr);
+        Ok(())
+    }
+
     #[test]
     fn unknown_expression_id_allow_unknown() {
         let session = VortexSession::empty().with::<ScalarFnSession>();

vortex-array/src/scalar_fn/session.rs

@@ -14,6 +14,7 @@ use crate::scalar_fn::ScalarFnPluginRef;
 use crate::scalar_fn::ScalarFnVTable;
 use crate::scalar_fn::fns::between::Between;
 use crate::scalar_fn::fns::binary::Binary;
+use crate::scalar_fn::fns::byte_length::ByteLength;
 use crate::scalar_fn::fns::cast::Cast;
 use crate::scalar_fn::fns::ext_storage::ExtStorage;
 use crate::scalar_fn::fns::fill_null::FillNull;
@@ -25,12 +26,14 @@ use crate::scalar_fn::fns::list_contains::ListContains;
 use crate::scalar_fn::fns::list_length::ListLength;
 use crate::scalar_fn::fns::list_sum::ListSum;
 use crate::scalar_fn::fns::literal::Literal;
+use crate::scalar_fn::fns::mask::Mask;
 use crate::scalar_fn::fns::merge::Merge;
 use crate::scalar_fn::fns::not::Not;
 use crate::scalar_fn::fns::pack::Pack;
 use crate::scalar_fn::fns::select::Select;
 use crate::scalar_fn::fns::stat::StatFn;
 use crate::scalar_fn::fns::variant_get::VariantGet;
+use crate::scalar_fn::fns::zip::Zip;
 
 /// Registry of scalar function vtables.
 pub type ScalarFnRegistry = ArcSwapMap<Id, ScalarFnPluginRef>;
@@ -62,6 +65,7 @@ impl Default for ScalarFnSession {
         // Register built-in expressions.
         this.register(Between);
         this.register(Binary);
+        this.register(ByteLength);
         this.register(Cast);
         this.register(ExtStorage);
         this.register(FillNull);
@@ -73,12 +77,14 @@ impl Default for ScalarFnSession {
         this.register(ListLength);
         this.register(ListSum);
         this.register(Literal);
+        this.register(Mask);
         this.register(Merge);
         this.register(Not);
         this.register(Pack);
         this.register(Select);
         this.register(StatFn);
         this.register(VariantGet);
+        this.register(Zip);
 
         this
     }

vortex-python/Cargo.toml

@@ -45,6 +45,7 @@ object_store = { workspace = true, features = [
     "http",
 ] }
 parking_lot = { workspace = true }
+prost = { workspace = true }
 pyo3 = { workspace = true, features = ["abi3", "abi3-py311"] }
 pyo3-bytes = { workspace = true }
 pyo3-log = { workspace = true }