Added serialize and deserialize methods to Expr

Author: ankane

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.21.2 (unreleased)
 
 - Added `serialize` and `deserialize` methods to `LazyFrame`
+- Added `serialize` and `deserialize` methods to `Expr`
 - Added `storage_options` and `retries` options to `sink_ipc` method
 - Added experimental support for Iceberg
 - Added experimental `cast_options` option to `scan_parquet` method

ext/polars/src/expr/mod.rs

@@ -10,6 +10,7 @@ mod meta;
 mod name;
 mod rolling;
 pub mod selector;
+mod serde;
 mod string;
 mod r#struct;
 

ext/polars/src/expr/serde.rs

@@ -0,0 +1,28 @@
+use std::io::{BufReader, BufWriter};
+
+use magnus::Value;
+use polars::lazy::prelude::Expr;
+use polars_utils::pl_serialize;
+
+use crate::exceptions::ComputeError;
+use crate::file::get_file_like;
+use crate::{RbExpr, RbResult};
+
+impl RbExpr {
+    pub fn serialize_binary(&self, rb_f: Value) -> RbResult<()> {
+        let file = get_file_like(rb_f, true)?;
+        let writer = BufWriter::new(file);
+        pl_serialize::SerializeOptions::default()
+            .serialize_into_writer::<_, _, true>(writer, &self.inner)
+            .map_err(|err| ComputeError::new_err(err.to_string()))
+    }
+
+    pub fn deserialize_binary(rb_f: Value) -> RbResult<RbExpr> {
+        let file = get_file_like(rb_f, false)?;
+        let reader = BufReader::new(file);
+        let expr: Expr = pl_serialize::SerializeOptions::default()
+            .deserialize_from_reader::<_, _, true>(reader)
+            .map_err(|err| ComputeError::new_err(err.to_string()))?;
+        Ok(expr.into())
+    }
+}

ext/polars/src/lib.rs

@@ -578,6 +578,11 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("hist", method!(RbExpr::hist, 4))?;
     class.define_method("into_selector", method!(RbExpr::into_selector, 0))?;
     class.define_singleton_method("new_selector", function!(RbExpr::new_selector, 1))?;
+    class.define_method("serialize_binary", method!(RbExpr::serialize_binary, 1))?;
+    class.define_singleton_method(
+        "deserialize_binary",
+        function!(RbExpr::deserialize_binary, 1),
+    )?;
 
     // bitwise
     class.define_method("bitwise_count_ones", method!(RbExpr::bitwise_count_ones, 0))?;

lib/polars/expr.rb

@@ -146,6 +146,38 @@ def -@
       wrap_expr(_rbexpr.neg)
     end
 
+    # Read a serialized expression from a file.
+    #
+    # @param source [Object]
+    #   Path to a file or a file-like object (by file-like object, we refer to
+    #   objects that have a `read` method, such as a file handler or `StringIO`).
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   This function uses  marshaling if the logical plan contains Ruby UDFs,
+    #   and as such inherits the security implications. Deserializing can execute
+    #   arbitrary code, so it should only be attempted on trusted data.
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example
+    #   expr = Polars.col("foo").sum.over("bar")
+    #   bytes = expr.meta.serialize
+    #   Polars::Expr.deserialize(StringIO.new(bytes))
+    #   # => col("foo").sum().over([col("bar")])
+    def self.deserialize(source)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      deserializer = RbExpr.method(:deserialize_binary)
+
+      _from_rbexpr(deserializer.(source))
+    end
+
     # Cast to physical representation of the logical dtype.
     #
     # - `:date` -> `:i32`

lib/polars/lazy_frame.rb

@@ -45,8 +45,7 @@ def self.read_json(file)
     #
     # @param source [Object]
     #   Path to a file or a file-like object (by file-like object, we refer to
-    #   objects that have a `read` method, such as a file handler (e.g.
-    #   via builtin `open` function) or `StringIO`).
+    #   objects that have a `read` method, such as a file handler or `StringIO`).
     #
     # @return [LazyFrame]
     #

lib/polars/meta_expr.rb

@@ -248,6 +248,29 @@ def as_selector
       Selector._from_rbselector(_rbexpr.into_selector)
     end
 
+    # Serialize this expression to a file or string.
+    #
+    # @param file [Object]
+    #   File path to which the result should be written. If set to `nil`
+    #   (default), the output is returned as a string instead.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example Serialize the expression into a binary representation.
+    #   expr = Polars.col("foo").sum.over("bar")
+    #   bytes = expr.meta.serialize
+    #   Polars::Expr.deserialize(StringIO.new(bytes))
+    #   # => col("foo").sum().over([col("bar")])
+    def serialize(file = nil)
+      serializer = _rbexpr.method(:serialize_binary)
+
+      Utils.serialize_polars_object(serializer, file)
+    end
+
     # Format the expression as a tree.
     #
     # @param return_as_string [Boolean]