fix: improve cleanup of streams in StreamWithContext drop implementation

Author: kosiew

src/dataframe.rs

@@ -449,21 +449,72 @@ impl RecordBatchReader for StreamWithContext {
 
 impl Drop for StreamWithContext {
     fn drop(&mut self) {
-        // Explicitly close streams before the context is released
-        self.reader.streams.clear();
+        // Drain streams before the context is released. Ensure cleanup runs on
+        // the Tokio runtime so asynchronous streams are fully consumed before
+        // the `SessionContext` is dropped.
+        let streams = std::mem::take(&mut self.reader.streams);
+        if streams.is_empty() {
+            return;
+        }
+
+        tokio::task::block_in_place(|| {
+            let handle = tokio::runtime::Handle::current();
+            handle.block_on(async move {
+                for stream in streams {
+                    let mut s = stream.lock().await;
+                    loop {
+                        match poll_next_batch(&mut s).await {
+                            Ok(Some(_)) => continue,
+                            Ok(None) => break,
+                            Err(e) => {
+                                #[cfg(debug_assertions)]
+                                {
+                                    log::error!("Stream cleanup error during drop: {e}");
+                                }
+                                break;
+                            }
+                        }
+                    }
+                }
+            });
+        });
+    }
+}
+
+unsafe fn drop_stream_capsule(stream: *mut FFI_ArrowArrayStream) {
+    // Drop the stream within a Tokio runtime context while holding the GIL.
+    // Any errors are logged but ignored to avoid unwinding across FFI
+    // boundaries.
+    let result = std::panic::catch_unwind(|| {
+        Python::with_gil(|_| {
+            get_tokio_runtime().0.block_on(async {
+                drop(Box::from_raw(stream));
+            });
+        });
+    });
+
+    if result.is_err() {
+        #[cfg(debug_assertions)]
+        {
+            log::error!("Failed to drop Arrow stream capsule");
+        }
     }
 }
 
 unsafe extern "C" fn stream_capsule_destructor(capsule: *mut ffi::PyObject) {
     let name = pyo3::ffi::c_str!("arrow_array_stream");
-    unsafe {
-        let ptr = ffi::PyCapsule_GetPointer(capsule, name.as_ptr());
-        if ptr.is_null() {
-            ffi::PyCapsule_SetDestructor(capsule, None);
-            return;
-        }
-        drop(Box::from_raw(ptr.cast::<FFI_ArrowArrayStream>()));
+    let ptr = ffi::PyCapsule_GetPointer(capsule, name.as_ptr());
+    if ptr.is_null() {
+        ffi::PyCapsule_SetDestructor(capsule, None);
+        return;
     }
+
+    // Ensure destructor is idempotent by nulling the pointer and destructor so
+    // subsequent calls become no-ops.
+    ffi::PyCapsule_SetPointer(capsule, std::ptr::null_mut());
+    ffi::PyCapsule_SetDestructor(capsule, None);
+
+    drop_stream_capsule(ptr.cast::<FFI_ArrowArrayStream>());
 }
 
 #[pymethods]