[rust] Rust Data structures pickling (#716)

* Adds pickling methods for DataType, Field, Expr, Series and Table
* adds tests to ensure correctness

Cargo.toml

@@ -10,6 +10,7 @@ version = "0.14.2"
 version = "1.3.3"
 
 [dependencies.indexmap]
+features = ["serde"]
 version = "1.9.2"
 
 [dependencies.num-traits]

daft/series.py

@@ -221,6 +221,9 @@ def str(self) -> SeriesStringNamespace:
     def dt(self) -> SeriesDateNamespace:
         return SeriesDateNamespace.from_series(self)
 
+    def __reduce__(self) -> tuple:
+        return (Series.from_arrow, (self.to_arrow(), self.name()))
+
 
 SomeSeriesNamespace = TypeVar("SomeSeriesNamespace", bound="SeriesNamespace")
 

daft/table/table.py

@@ -248,3 +248,7 @@ def argsort(self, sort_keys: ExpressionsProjection, descending: bool | list[bool
         else:
             raise TypeError(f"Expected a bool, list[bool] or None for `descending` but got {type(descending)}")
         return Series._from_pyseries(self._table.argsort(pyexprs, descending))
+
+    def __reduce__(self) -> tuple:
+        names = self.column_names()
+        return Table.from_pydict, ({name: self.get_column(name) for name in names},)

src/python/datatype.rs

@@ -1,5 +1,9 @@
 use crate::datatypes::DataType;
-use pyo3::prelude::*;
+use pyo3::{
+    exceptions::PyValueError,
+    prelude::*,
+    types::{PyBytes, PyTuple},
+};
 
 #[pyclass]
 #[derive(Clone)]
@@ -9,6 +13,18 @@ pub struct PyDataType {
 
 #[pymethods]
 impl PyDataType {
+    #[new]
+    #[args(args = "*")]
+    fn new(args: &PyTuple) -> PyResult<Self> {
+        match args.len() {
+            0 => Ok(DataType::new_null().into()),
+            _ => Err(PyValueError::new_err(format!(
+                "expected no arguments to make new PyDataType, got : {}",
+                args.len()
+            ))),
+        }
+    }
+
     pub fn __repr__(&self) -> PyResult<String> {
         Ok(format!("{}", self.dtype))
     }
@@ -96,6 +112,20 @@ impl PyDataType {
             Ok(false)
         }
     }
+
+    pub fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> {
+        match state.extract::<&PyBytes>(py) {
+            Ok(s) => {
+                self.dtype = bincode::deserialize(s.as_bytes()).unwrap();
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
+        Ok(PyBytes::new(py, &bincode::serialize(&self.dtype).unwrap()).to_object(py))
+    }
 }
 
 impl From<DataType> for PyDataType {

src/python/field.rs

@@ -1,7 +1,11 @@
-use pyo3::prelude::*;
+use pyo3::{
+    exceptions::PyValueError,
+    prelude::*,
+    types::{PyBytes, PyTuple},
+};
 
 use super::datatype::PyDataType;
-use crate::datatypes;
+use crate::datatypes::{self, DataType, Field};
 
 #[pyclass]
 pub struct PyField {
@@ -10,6 +14,18 @@ pub struct PyField {
 
 #[pymethods]
 impl PyField {
+    #[new]
+    #[args(args = "*")]
+    fn new(args: &PyTuple) -> PyResult<Self> {
+        match args.len() {
+            0 => Ok(Field::new("null", DataType::new_null()).into()),
+            _ => Err(PyValueError::new_err(format!(
+                "expected no arguments to make new PyDataType, got : {}",
+                args.len()
+            ))),
+        }
+    }
+
     pub fn name(&self) -> PyResult<String> {
         Ok(self.field.name.clone())
     }
@@ -21,6 +37,20 @@ impl PyField {
     pub fn eq(&self, other: &PyField) -> PyResult<bool> {
         Ok(self.field.eq(&other.field))
     }
+
+    pub fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> {
+        match state.extract::<&PyBytes>(py) {
+            Ok(s) => {
+                self.field = bincode::deserialize(s.as_bytes()).unwrap();
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
+        Ok(PyBytes::new(py, &bincode::serialize(&self.field).unwrap()).to_object(py))
+    }
 }
 
 impl From<datatypes::Field> for PyField {

src/python/schema.rs

@@ -1,11 +1,15 @@
 use std::sync::Arc;
 
+use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use pyo3::types::PyBytes;
+use pyo3::types::PyTuple;
 
 use super::datatype::PyDataType;
 use super::field::PyField;
 use crate::datatypes;
 use crate::schema;
+use crate::schema::Schema;
 
 #[pyclass]
 pub struct PySchema {
@@ -14,6 +18,20 @@ pub struct PySchema {
 
 #[pymethods]
 impl PySchema {
+    #[new]
+    #[args(args = "*")]
+    fn new(args: &PyTuple) -> PyResult<Self> {
+        match args.len() {
+            0 => Ok(Self {
+                schema: Schema::empty().into(),
+            }),
+            _ => Err(PyValueError::new_err(format!(
+                "expected no arguments to make new PyDataType, got : {}",
+                args.len()
+            ))),
+        }
+    }
+
     pub fn __getitem__(&self, name: &str) -> PyResult<PyField> {
         Ok(self.schema.get_field(name)?.clone().into())
     }
@@ -44,6 +62,20 @@ impl PySchema {
             schema: schema.into(),
         })
     }
+
+    pub fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> {
+        match state.extract::<&PyBytes>(py) {
+            Ok(s) => {
+                self.schema = bincode::deserialize(s.as_bytes()).unwrap();
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
+        Ok(PyBytes::new(py, &bincode::serialize(&self.schema).unwrap()).to_object(py))
+    }
 }
 
 impl From<schema::SchemaRef> for PySchema {

src/schema.rs

@@ -5,6 +5,7 @@ use std::{
 };
 
 use indexmap::IndexMap;
+use serde::{Deserialize, Serialize};
 
 use crate::{
     datatypes::Field,
@@ -13,8 +14,10 @@ use crate::{
 
 pub type SchemaRef = Arc<Schema>;
 
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+#[serde(transparent)]
 pub struct Schema {
+    #[serde(with = "indexmap::serde_seq")]
     pub fields: indexmap::IndexMap<String, Field>,
 }