Fix error handling in async stream execution for PySessionContext and PyDataFrame

Author: kosiew

src/context.rs

@@ -1031,7 +1031,7 @@ impl PySessionContext {
         let plan = plan.plan.clone();
         let fut: JoinHandle<datafusion::common::Result<SendableRecordBatchStream>> =
             rt.spawn(async move { plan.execute(part, Arc::new(ctx)) });
-        let stream = wait_for_future(py, async { fut.await.map_err(to_datafusion_err) })??;
+        let stream = wait_for_future(py, async { fut.await.map_err(to_datafusion_err) })???;
         Ok(PyRecordBatchStream::new(stream))
     }
 }

src/dataframe.rs

@@ -753,7 +753,7 @@ impl PyDataFrame {
         let df = self.df.as_ref().clone();
         let fut: JoinHandle<datafusion::common::Result<SendableRecordBatchStream>> =
             rt.spawn(async move { df.execute_stream().await });
-        let stream = wait_for_future(py, async { fut.await.map_err(to_datafusion_err) })??;
+        let stream = wait_for_future(py, async { fut.await.map_err(to_datafusion_err) })???;
         Ok(PyRecordBatchStream::new(stream))
     }
 
@@ -764,6 +764,7 @@ impl PyDataFrame {
         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.map_err(to_datafusion_err) })?
+            .map_err(py_datafusion_err)?
             .map_err(py_datafusion_err)?;
 
         Ok(stream.into_iter().map(PyRecordBatchStream::new).collect())