feat: Make serialization code not unwrap and panic on failures (#3546)

# Overview Serialization of subqueries currently fails. Thus, when ray tries to serialize subqueries, the ray-task fails and *indefinitely* hangs. This PR updates the failure mechanism of serialization of subqueries. Instead of unwrapping the `Err` in rust, it is passed through instead. --------- Co-authored-by: Raunak Bhagat <[email protected]>
Eventual-Inc · Dec 11, 2024 · c3088d6 · c3088d6
1 parent dd2dc23
commit c3088d6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/benchmarking/tpcds/ray_entrypoint.py b/benchmarking/tpcds/ray_entrypoint.py
@@ -1,5 +1,4 @@
 import argparse
-import sys
 from pathlib import Path
 
 import helpers
@@ -17,12 +16,9 @@ def run(
     with open(query_file) as f:
         query = f.read()
 
-    try:
-        daft.sql(query, catalog=catalog).explain(show_all=True)
-        if not dry_run:
-            daft.sql(query, catalog=catalog).collect()
-    except Exception as e:
-        print(str(e), file=sys.stderr)
+    daft.sql(query, catalog=catalog).explain(show_all=True)
+    if not dry_run:
+        daft.sql(query, catalog=catalog).collect()
 
 
 if __name__ == "__main__":

diff --git a/src/common/py-serde/src/python.rs b/src/common/py-serde/src/python.rs
@@ -84,16 +84,22 @@ macro_rules! impl_bincode_py_state_serialization {
                 py: Python<'py>,
             ) -> PyResult<(PyObject, (pyo3::Bound<'py, pyo3::types::PyBytes>,))> {
                 use pyo3::{
+                    exceptions::PyRuntimeError,
                     types::{PyAnyMethods, PyBytes},
-                    PyTypeInfo, ToPyObject,
+                    PyErr, PyTypeInfo, ToPyObject,
                 };
                 Ok((
                     Self::type_object_bound(py)
                         .getattr(pyo3::intern!(py, "_from_serialized"))?
                         .into(),
                     (PyBytes::new_bound(
                         py,
-                        &$crate::bincode::serialize(&self).unwrap(),
+                        &$crate::bincode::serialize(&self).map_err(|error| {
+                            PyErr::new::<PyRuntimeError, _>(format!(
+                                "Failed to serialize: {}",
+                                error.to_string()
+                            ))
+                        })?,
                     ),),
                 ))
             }