fix: improve error message for execution failures in PySessionContext

Author: kosiew

src/context.rs

@@ -846,8 +846,9 @@ impl PySessionContext {
             rt.spawn(async move { plan.execute(part, Arc::new(ctx)) });
         let join_result = wait_for_future(py, fut)
             .map_err(|e| PyDataFusionError::Common(format!("Task failed: {}", e)))?;
-        let stream = join_result.map_err(PyDataFusionError::from)?;
-        Ok(PyRecordBatchStream::new(stream))
+        let stream = join_result
+            .map_err(|e| PyDataFusionError::Common(format!("Execution error: {}", e)))?;
+        Ok(PyRecordBatchStream::new(stream?))
     }
 
     pub fn table_exist(&self, name: &str) -> PyDataFusionResult<bool> {

src/dataframe.rs

@@ -756,7 +756,9 @@ impl PyDataFrame {
         let fut: JoinHandle<datafusion::common::Result<SendableRecordBatchStream>> =
             rt.spawn(async move { df.execute_stream().await });
         let stream = wait_for_future(py, fut).map_err(py_datafusion_err)?;
-        Ok(PyRecordBatchStream::new(stream?))
+        Ok(PyRecordBatchStream::new(
+            stream.map_err(PyDataFusionError::from)?,
+        ))
     }
 
     fn execute_stream_partitioned(&self, py: Python) -> PyResult<Vec<PyRecordBatchStream>> {
@@ -817,7 +819,7 @@ impl PyDataFrame {
 
     // Executes this DataFrame to get the total number of rows.
     fn count(&self, py: Python) -> PyDataFusionResult<usize> {
-        Ok(wait_for_future(py, self.df.as_ref().clone().count())?)
+        Ok(wait_for_future(py, self.df.as_ref().clone().count())??)
     }
 
     /// Fill null values with a specified value for specific columns
@@ -844,11 +846,8 @@ impl PyDataFrame {
 fn print_dataframe(py: Python, df: DataFrame) -> PyDataFusionResult<()> {
     // Get string representation of record batches
     let batches = wait_for_future(py, df.collect())?;
-    let batches_as_string = pretty::pretty_format_batches(&batches);
-    let result = match batches_as_string {
-        Ok(batch) => format!("DataFrame()\n{batch}"),
-        Err(err) => format!("Error: {:?}", err.to_string()),
-    };
+    let batches_as_string = pretty::pretty_format_batches(&batches)?;
+    let result = format!("DataFrame()\n{batch}", batch = batches_as_string);
 
     // Import the Python 'builtins' module to access the print function
     // Note that println! does not print to the Python debug console and is not visible in notebooks for instance

src/record_batch.rs

@@ -63,7 +63,7 @@ impl PyRecordBatchStream {
 impl PyRecordBatchStream {
     fn next(&mut self, py: Python) -> PyResult<PyRecordBatch> {
         let stream = self.stream.clone();
-        wait_for_future(py, next_stream(stream, true))
+        wait_for_future(py, next_stream(stream, true))?
     }
 
     fn __next__(&mut self, py: Python) -> PyResult<PyRecordBatch> {