Skip to content

Commit

Permalink
we need GIL, no?
Browse files Browse the repository at this point in the history
  • Loading branch information
aersam committed Mar 13, 2024
1 parent a3d5585 commit 965968c
Showing 1 changed file with 15 additions and 1 deletion.
16 changes: 15 additions & 1 deletion python/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1385,6 +1385,20 @@ impl From<&PyAddAction> for Add {
}
}

struct GilIterator {
reader: ArrowArrayStreamReader,
}

impl<'a> Iterator for GilIterator {
type Item = RecordBatch;

fn next(&mut self) -> Option<Self::Item> {
let reader = &mut self.reader;

Python::with_gil(|_| reader.next().map_or(None, |v| Some(v.unwrap())))
}
}

#[pyfunction]
#[allow(clippy::too_many_arguments)]
fn write_to_deltalake(
Expand All @@ -1405,7 +1419,7 @@ fn write_to_deltalake(
custom_metadata: Option<HashMap<String, String>>,
) -> PyResult<()> {
py.allow_threads(|| {
let batches = data.0.map(|batch| batch.unwrap());
let batches = GilIterator { reader: data.0 }; // The pyarrow reader can be backed by a python object
let save_mode = mode.parse().map_err(PythonError::from)?;

let options = storage_options.clone().unwrap_or_default();
Expand Down

0 comments on commit 965968c

Please sign in to comment.