Added serialize and deserialize methods to DataFrame

Author: ankane

CHANGELOG.md

@@ -1,7 +1,6 @@
 ## 0.21.2 (unreleased)
 
-- Added `serialize` and `deserialize` methods to `LazyFrame`
-- Added `serialize` and `deserialize` methods to `Expr`
+- Added `serialize` and `deserialize` methods to `DataFrame`, `LazyFrame`, and `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/dataframe/serde.rs

@@ -1,15 +1,30 @@
-use crate::exceptions::ComputeError;
+use std::io::{BufReader, BufWriter};
+
+use polars::prelude::*;
+
 use crate::file::get_file_like;
-use crate::{RbDataFrame, RbResult};
+use crate::utils::to_rb_err;
+use crate::{RbDataFrame, RbPolarsErr, RbResult};
 use magnus::Value;
-use std::io::BufWriter;
 
 impl RbDataFrame {
-    // TODO add to Ruby
-    pub fn serialize_json(&self, rb_f: Value) -> RbResult<()> {
+    pub fn serialize_binary(&self, rb_f: Value) -> RbResult<()> {
         let file = get_file_like(rb_f, true)?;
-        let writer = BufWriter::new(file);
-        serde_json::to_writer(writer, &self.df)
-            .map_err(|err| ComputeError::new_err(err.to_string()))
+        let mut writer = BufWriter::new(file);
+
+        Ok(self
+            .df
+            .borrow_mut()
+            .serialize_into_writer(&mut writer)
+            .map_err(RbPolarsErr::from)?)
+    }
+
+    pub fn deserialize_binary(rb_f: Value) -> RbResult<Self> {
+        let file = get_file_like(rb_f, false)?;
+        let mut file = BufReader::new(file);
+
+        DataFrame::deserialize_from_reader(&mut file)
+            .map(|v| v.into())
+            .map_err(to_rb_err)
     }
 }

ext/polars/src/lib.rs

@@ -158,7 +158,14 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("to_struct", method!(RbDataFrame::to_struct, 1))?;
     class.define_method("unnest", method!(RbDataFrame::unnest, 1))?;
     class.define_method("clear", method!(RbDataFrame::clear, 0))?;
-    class.define_method("serialize_json", method!(RbDataFrame::serialize_json, 1))?;
+    class.define_method(
+        "serialize_binary",
+        method!(RbDataFrame::serialize_binary, 1),
+    )?;
+    class.define_singleton_method(
+        "deserialize_binary",
+        function!(RbDataFrame::deserialize_binary, 1),
+    )?;
 
     let class = module.define_class("RbExpr", ruby.class_object())?;
     class.define_method("+", method!(RbExpr::add, 1))?;

lib/polars/data_frame.rb

@@ -72,6 +72,43 @@ def initialize(data = nil, schema: nil, schema_overrides: nil, strict: true, ori
       end
     end
 
+    # Read a serialized DataFrame 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 [DataFrame]
+    #
+    # @note
+    #   Serialization is not stable across Polars versions: a LazyFrame serialized
+    #   in one Polars version may not be deserializable in another Polars version.
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [1, 2, 3], "b" => [4.0, 5.0, 6.0]})
+    #   bytes = df.serialize
+    #   Polars::DataFrame.deserialize(StringIO.new(bytes))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ f64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 4.0 │
+    #   # │ 2   ┆ 5.0 │
+    #   # │ 3   ┆ 6.0 │
+    #   # └─────┴─────┘
+    def self.deserialize(source)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+
+      deserializer = RbDataFrame.method(:deserialize_binary)
+
+      _from_rbdf(deserializer.(source))
+    end
+
     # @private
     def self._from_rbdf(rb_df)
       df = DataFrame.allocate
@@ -627,6 +664,44 @@ def to_series(index = 0)
       Utils.wrap_s(_df.select_at_idx(index))
     end
 
+    # Serialize this DataFrame to a file or string.
+    #
+    # @param file [Object]
+    #   File path or writable file-like object to which the result will be written.
+    #   If set to `None` (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
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "foo" => [1, 2, 3],
+    #       "bar" => [6, 7, 8]
+    #     }
+    #   )
+    #   bytes = df.serialize
+    #   Polars::DataFrame.deserialize(StringIO.new(bytes))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ foo ┆ bar │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 6   │
+    #   # │ 2   ┆ 7   │
+    #   # │ 3   ┆ 8   │
+    #   # └─────┴─────┘
+    def serialize(file = nil)
+      serializer = _df.method(:serialize_binary)
+
+      Utils.serialize_polars_object(serializer, file)
+    end
+
     # Serialize to JSON representation.
     #
     # @param file [String]