jiayuasu commented on code in PR #1187:
URL: https://github.com/apache/sedona-db/pull/1187#discussion_r3911545652
##########
python/sedonadb/src/datasource.rs:
##########
@@ -400,25 +427,99 @@ impl PyProjectedRecordBatchReader {
/// ArrowArrayStream/RecordBatchReader (e.g., the pyogrio context manager, or
/// an ADBC statement/cursor).
struct WrappedRecordBatchReader {
- pub inner: Box<dyn RecordBatchReader + Send>,
+ pub inner: Option<Box<dyn RecordBatchReader + Send>>,
+ pub schema: SchemaRef,
pub shelter: Option<Py<PyAny>>,
}
+impl WrappedRecordBatchReader {
+ /// Release resources in dependency order: the Arrow FFI stream first and
+ /// its owning Python context second. Explicitly dropping the shelter while
+ /// attached avoids PyO3 deferring its decref until a later Python
callback.
+ fn finish(&mut self) {
+ self.inner = None;
+
+ if let Some(shelter) = self.shelter.take() {
+ if Python::try_attach(|_| drop(shelter)).is_none() {
+ // During interpreter shutdown PyO3 cannot safely decref the
+ // object. Let PyO3's normal deferred-drop path handle it.
+ }
+ }
+ }
+}
+
impl RecordBatchReader for WrappedRecordBatchReader {
fn schema(&self) -> SchemaRef {
- self.inner.schema()
+ self.schema.clone()
}
}
impl Iterator for WrappedRecordBatchReader {
type Item = Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
- if let Some(item) = self.inner.next() {
+ if let Some(item) = self.inner.as_mut().and_then(|inner| inner.next())
{
Some(item)
} else {
- self.shelter = None;
+ self.finish();
None
}
}
}
+
+impl Drop for WrappedRecordBatchReader {
+ fn drop(&mut self) {
+ self.finish();
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::RecordBatchIterator;
+ use pyo3::types::{PyAnyMethods, PyModule};
+
+ #[test]
+ fn wrapped_reader_drops_python_shelter_at_eof() {
+ let (shelter, shelter_class) = Python::attach(|py| {
+ let module = PyModule::from_code(
+ py,
+ cr#"
+class Shelter:
+ destroyed = 0
+
+ def __del__(self):
+ type(self).destroyed += 1
+"#,
Review Comment:
Good point. I added Python integration coverage for the observable behavior.
It now covers serialization within one scan, concurrent multibatch scans, and
paused reader progress for FlatGeoBuf and GeoPackage.
I kept the two narrow Rust tests because they verify an FFI ownership
invariant that Python cannot reliably observe: the imported Arrow stream must
be released before its Python shelter at EOF and early cancellation. The user
visible behavior is covered in Python. The Rust tests only pin that cleanup
order.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
