[BUG] Use Daft Pickle instead of Ray Pickle and use bincode for serializing

src/common/py-serde/src/python.rs

@@ -7,8 +7,7 @@ where
     S: Serializer,
 {
     let bytes = Python::with_gil(|py| {
-        py.import(pyo3::intern!(py, "ray.cloudpickle"))
-            .or_else(|_| py.import(pyo3::intern!(py, "pickle")))
+        py.import(pyo3::intern!(py, "daft.pickle"))
             .and_then(|m| m.getattr(pyo3::intern!(py, "dumps")))
             .and_then(|f| f.call1((obj,)))
             .and_then(|b| b.extract::<Vec<u8>>())
@@ -31,8 +30,7 @@ impl<'de> Visitor<'de> for PyObjectVisitor {
         E: DeError,
     {
         Python::with_gil(|py| {
-            py.import(pyo3::intern!(py, "ray.cloudpickle"))
-                .or_else(|_| py.import(pyo3::intern!(py, "pickle")))
+            py.import(pyo3::intern!(py, "daft.pickle"))
                 .and_then(|m| m.getattr(pyo3::intern!(py, "loads")))
                 .and_then(|f| Ok(f.call1((v,))?.to_object(py)))
                 .map_err(|e| DeError::custom(e.to_string()))
@@ -44,8 +42,7 @@ impl<'de> Visitor<'de> for PyObjectVisitor {
         E: DeError,
     {
         Python::with_gil(|py| {
-            py.import(pyo3::intern!(py, "ray.cloudpickle"))
-                .or_else(|_| py.import(pyo3::intern!(py, "pickle")))
+            py.import(pyo3::intern!(py, "daft.pickle"))
                 .and_then(|m| m.getattr(pyo3::intern!(py, "loads")))
                 .and_then(|f| Ok(f.call1((v,))?.to_object(py)))
                 .map_err(|e| DeError::custom(e.to_string()))

src/daft-dsl/Cargo.toml

@@ -1,4 +1,5 @@
 [dependencies]
+bincode = {workspace = true}
 common-error = {path = "../common/error", default-features = false}
 common-hashable-float-wrapper = {path = "../common/hashable-float-wrapper"}
 common-py-serde = {path = "../common/py-serde", default-features = false}
@@ -11,7 +12,6 @@ itertools = {workspace = true}
 log = {workspace = true}
 pyo3 = {workspace = true, optional = true}
 serde = {workspace = true}
-serde_json = {workspace = true}
 typetag = "0.2.16"
 
 [features]

src/daft-dsl/src/functions/python/mod.rs

@@ -1,5 +1,3 @@
-#[cfg(feature = "python")]
-mod pyobj_serde;
 mod udf;
 
 use std::sync::Arc;
@@ -35,7 +33,7 @@ impl PythonUDF {
 pub struct StatelessPythonUDF {
     pub name: Arc<String>,
     #[cfg(feature = "python")]
-    partial_func: pyobj_serde::PyObjectWrapper,
+    partial_func: crate::pyobj_serde::PyObjectWrapper,
     num_expressions: usize,
     pub return_dtype: DataType,
     pub resource_request: Option<ResourceRequest>,
@@ -46,12 +44,12 @@ pub struct StatelessPythonUDF {
 pub struct StatefulPythonUDF {
     pub name: Arc<String>,
     #[cfg(feature = "python")]
-    pub stateful_partial_func: pyobj_serde::PyObjectWrapper,
+    pub stateful_partial_func: crate::pyobj_serde::PyObjectWrapper,
     pub num_expressions: usize,
     pub return_dtype: DataType,
     pub resource_request: Option<ResourceRequest>,
     #[cfg(feature = "python")]
-    pub init_args: Option<pyobj_serde::PyObjectWrapper>,
+    pub init_args: Option<crate::pyobj_serde::PyObjectWrapper>,
     pub batch_size: Option<usize>,
     pub concurrency: Option<usize>,
 }
@@ -68,7 +66,7 @@ pub fn stateless_udf(
     Ok(Expr::Function {
         func: super::FunctionExpr::Python(PythonUDF::Stateless(StatelessPythonUDF {
             name: name.to_string().into(),
-            partial_func: pyobj_serde::PyObjectWrapper(py_partial_stateless_udf),
+            partial_func: crate::pyobj_serde::PyObjectWrapper(py_partial_stateless_udf),
             num_expressions: expressions.len(),
             return_dtype,
             resource_request,
@@ -113,11 +111,11 @@ pub fn stateful_udf(
     Ok(Expr::Function {
         func: super::FunctionExpr::Python(PythonUDF::Stateful(StatefulPythonUDF {
             name: name.to_string().into(),
-            stateful_partial_func: pyobj_serde::PyObjectWrapper(py_stateful_partial_func),
+            stateful_partial_func: crate::pyobj_serde::PyObjectWrapper(py_stateful_partial_func),
             num_expressions: expressions.len(),
             return_dtype,
             resource_request,
-            init_args: init_args.map(pyobj_serde::PyObjectWrapper),
+            init_args: init_args.map(crate::pyobj_serde::PyObjectWrapper),
             batch_size,
             concurrency,
         })),

src/daft-dsl/src/lib.rs

@@ -8,7 +8,7 @@ pub mod join;
 mod lit;
 pub mod optimization;
 #[cfg(feature = "python")]
-mod pyobject;
+mod pyobj_serde;
 #[cfg(feature = "python")]
 pub mod python;
 mod resolve_expr;

src/daft-dsl/src/lit.rs

@@ -22,7 +22,7 @@ use std::{
 };
 
 #[cfg(feature = "python")]
-use crate::pyobject::DaftPyObject;
+use crate::pyobj_serde::PyObjectWrapper;
 
 /// Stores a literal value for queries and computations.
 /// We only need to support the limited types below since those are the types that we would get from python.
@@ -71,7 +71,7 @@ pub enum LiteralValue {
     Series(Series),
     /// Python object.
     #[cfg(feature = "python")]
-    Python(DaftPyObject),
+    Python(PyObjectWrapper),
 }
 
 impl Eq for LiteralValue {}
@@ -144,12 +144,8 @@ impl Display for LiteralValue {
             #[cfg(feature = "python")]
             Python(pyobj) => write!(f, "PyObject({})", {
                 use pyo3::prelude::*;
-                Python::with_gil(|py| {
-                    pyobj
-                        .pyobject
-                        .call_method0(py, pyo3::intern!(py, "__str__"))
-                })
-                .unwrap()
+                Python::with_gil(|py| pyobj.0.call_method0(py, pyo3::intern!(py, "__str__")))
+                    .unwrap()
             }),
         }
     }
@@ -212,7 +208,7 @@ impl LiteralValue {
             }
             Series(series) => series.clone().rename("literal"),
             #[cfg(feature = "python")]
-            Python(val) => PythonArray::from(("literal", vec![val.pyobject.clone()])).into_series(),
+            Python(val) => PythonArray::from(("literal", vec![val.0.clone()])).into_series(),
         };
         result
     }
@@ -356,7 +352,7 @@ impl Literal for Series {
 #[cfg(feature = "python")]
 impl Literal for pyo3::PyObject {
     fn lit(self) -> ExprRef {
-        Expr::Literal(LiteralValue::Python(DaftPyObject { pyobject: self })).into()
+        Expr::Literal(LiteralValue::Python(PyObjectWrapper(self))).into()
     }
 }
 

src/daft-dsl/src/functions/python/pyobj_serde.rs → src/daft-dsl/src/pyobj_serde.rs

@@ -1,4 +1,7 @@
-use std::hash::{Hash, Hasher};
+use std::{
+    hash::{Hash, Hasher},
+    io::Write,
+};
 
 use common_py_serde::{deserialize_py_object, serialize_py_object};
 use pyo3::{PyObject, Python};
@@ -29,7 +32,25 @@ impl Hash for PyObjectWrapper {
             // If Python object is hashable, hash the Python-side hash.
             Ok(py_obj_hash) => py_obj_hash.hash(state),
             // Fall back to hashing the pickled Python object.
-            Err(_) => serde_json::to_vec(self).unwrap().hash(state),
+            Err(_) => {
+                let hasher = HashWriter { state };
+                bincode::serialize_into(hasher, self)
+                    .expect("Pickling error occurred when computing hash of Pyobject")
+            }
         }
     }
 }
+
+struct HashWriter<'a, H: Hasher> {
+    state: &'a mut H,
+}
+
+impl<'a, H: Hasher> Write for HashWriter<'a, H> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        buf.hash(self.state);
+        Ok(buf.len())
+    }
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}