paleolimbot commented on code in PR #1187:
URL: https://github.com/apache/sedona-db/pull/1187#discussion_r3906473957
##########
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:
If there's any way to test this in Python instead of PyO3 I think it would
be cleaner.
##########
rust/sedona-datasource/src/format.rs:
##########
@@ -216,6 +226,10 @@ impl FileFormat for ExternalFileFormat {
#[derive(Debug, Clone)]
struct ExternalFileSource {
spec: Arc<dyn ExternalFormatSpec>,
+ /// Shared by all openers cloned from this physical scan. Keeping the lock
+ /// here scopes serialization to one scan instead of coupling independent
+ /// user streams through a process-wide lock.
+ reader_lock: Option<Arc<Mutex<()>>>,
Review Comment:
For the pyogrio case, this may not be sufficient: the serialization probably
has to occur process-wide (so two Python threads doing two scans may still
crash?).
--
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]
