Update src/dataframe.rs

Co-authored-by: Tim Saucer <[email protected]>

Author: kosiew

src/dataframe.rs

@@ -764,7 +764,9 @@ impl PyDataFrame {
         let df = self.df.as_ref().clone();
         let fut: JoinHandle<datafusion::common::Result<Vec<SendableRecordBatchStream>>> =
             rt.spawn(async move { df.execute_stream_partitioned().await });
-        let stream = wait_for_future(py, async { fut.await.expect("Tokio task panicked") })?
+        let stream = wait_for_future(py, async { fut.await })?
+            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?
+            .map_err(py_datafusion_err)?;
             .map_err(py_datafusion_err)?;
 
         Ok(stream.into_iter().map(PyRecordBatchStream::new).collect())