Make readers python iterators (#85)

kylebarron · Jul 31, 2024 · b582f3b · b582f3b
1 parent ebadc3d
commit b582f3b
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/arro3-core/python/arro3/core/_core.pyi b/arro3-core/python/arro3/core/_core.pyi
@@ -73,6 +73,8 @@ class Array:
 
 class ArrayReader:
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ...
+    def __iter__(self) -> ArrayReader: ...
+    def __next__(self) -> Array: ...
     def __repr__(self) -> str: ...
     @classmethod
     def from_arrow(
@@ -725,6 +727,8 @@ class RecordBatch:
 
 class RecordBatchReader:
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ...
+    def __iter__(self) -> RecordBatchReader: ...
+    def __next__(self) -> RecordBatch: ...
     def __repr__(self) -> str: ...
     @classmethod
     def from_arrow(

diff --git a/pyo3-arrow/src/array_reader.rs b/pyo3-arrow/src/array_reader.rs
@@ -122,6 +122,16 @@ impl PyArrayReader {
         to_stream_pycapsule(py, array_reader, requested_schema)
     }
 
+    // Return self
+    // https://stackoverflow.com/a/52056290
+    fn __iter__(&mut self, py: Python) -> PyResult<PyObject> {
+        self.to_arro3(py)
+    }
+
+    fn __next__(&mut self, py: Python) -> PyArrowResult<PyObject> {
+        self.read_next_array(py)
+    }
+
     pub fn __repr__(&self) -> String {
         self.to_string()
     }

diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs
@@ -145,6 +145,16 @@ impl PyRecordBatchReader {
         to_stream_pycapsule(py, array_reader, requested_schema)
     }
 
+    // Return self
+    // https://stackoverflow.com/a/52056290
+    fn __iter__(&mut self, py: Python) -> PyResult<PyObject> {
+        self.to_arro3(py)
+    }
+
+    fn __next__(&mut self, py: Python) -> PyArrowResult<PyObject> {
+        self.read_next_batch(py)
+    }
+
     pub fn __repr__(&self) -> String {
         self.to_string()
     }