test: add test for dropping stream after close to prevent crashes

Author: kosiew

python/tests/test_dataframe_iter.py

@@ -38,3 +38,11 @@ async def test_async_iter_dataframe(ctx):
     batches = [batch async for batch in df]
     assert len(batches) == 1
     assert batches[0].column(0).to_pylist() == [1, 2]
+
+
+def test_drop_after_close_does_not_crash(ctx):
+    df = ctx.from_pydict({"a": [1, 2]})
+    stream = df.execute_stream()
+    stream.close()
+    # Dropping the stream after close should not crash
+    del stream

src/record_batch.rs

@@ -170,13 +170,12 @@ async fn close_stream(stream: Arc<Mutex<SendableRecordBatchStream>>) {
 
 impl Drop for PyRecordBatchStream {
     fn drop(&mut self) {
-        if let Some(stream) = self.stream.take() {
-            // Drain the stream while the context is still alive
-            Python::with_gil(|py| {
-                let _ = wait_for_future(py, close_stream(stream));
-            });
-        }
-        // Drop the context after the stream has been fully drained
+        // Avoid interacting with Python during interpreter shutdown by
+        // not attempting to drain the stream here.  Any remaining
+        // batches will be dropped and background tasks allowed to
+        // complete naturally.  Users should call `close` explicitly to
+        // drain the stream when Python is still available.
+        self.stream.take();
         self._ctx.take();
     }
 }