Add into_arro3 method

kylebarron · Feb 22, 2025 · 3a869bc · 3a869bc
1 parent b3bbab4
commit 3a869bc
Show file tree

Hide file tree

Showing 11 changed files with 133 additions and 17 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pyo3-arrow/src/array.rs b/pyo3-arrow/src/array.rs
@@ -108,6 +108,17 @@ impl PyArray {
         )
     }
 
+    /// Export to an arro3.core.Array.
+    ///
+    /// This requires that you depend on arro3-core from your Python package.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let array_capsules = to_array_pycapsules(py, self.field.clone(), &self.array, None)?;
+        arro3_mod
+            .getattr(intern!(py, "Array"))?
+            .call_method1(intern!(py, "from_arrow_pycapsule"), array_capsules)
+    }
+
     /// Export this to a Python `nanoarrow.Array`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_array(py, &self.__arrow_c_array__(py, None)?)

diff --git a/pyo3-arrow/src/array_reader.rs b/pyo3-arrow/src/array_reader.rs
@@ -78,7 +78,6 @@ impl PyArrayReader {
     }
 
     /// Export this to a Python `arro3.core.ArrayReader`.
-    #[allow(clippy::wrong_self_convention)]
     pub fn to_arro3<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let arro3_mod = py.import(intern!(py, "arro3.core"))?;
         arro3_mod.getattr(intern!(py, "ArrayReader"))?.call_method1(
@@ -87,8 +86,23 @@ impl PyArrayReader {
         )
     }
 
+    /// Export this to a Python `arro3.core.ArrayReader`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let array_reader = self
+            .0
+            .lock()
+            .unwrap()
+            .take()
+            .ok_or(PyIOError::new_err("Cannot read from closed stream"))?;
+        let stream_pycapsule = to_stream_pycapsule(py, array_reader, None)?;
+        arro3_mod.getattr(intern!(py, "ArrayReader"))?.call_method1(
+            intern!(py, "from_arrow_pycapsule"),
+            PyTuple::new(py, vec![stream_pycapsule])?,
+        )
+    }
+
     /// Export this to a Python `nanoarrow.ArrayStream`.
-    #[allow(clippy::wrong_self_convention)]
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_array_stream(py, &self.__arrow_c_stream__(py, None)?)
     }
@@ -135,8 +149,8 @@ impl PyArrayReader {
 
     // Return self
     // https://stackoverflow.com/a/52056290
-    fn __iter__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        self.to_arro3(py)
+    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
+        slf
     }
 
     fn __next__(&self) -> PyArrowResult<Arro3Array> {

diff --git a/pyo3-arrow/src/chunked.rs b/pyo3-arrow/src/chunked.rs
@@ -197,6 +197,17 @@ impl PyChunkedArray {
             )
     }
 
+    /// Export this to a Python `arro3.core.ChunkedArray`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let capsule = Self::to_stream_pycapsule(py, self.chunks.clone(), self.field.clone(), None)?;
+        arro3_mod
+            .getattr(intern!(py, "ChunkedArray"))?
+            .call_method1(
+                intern!(py, "from_arrow_pycapsule"),
+                PyTuple::new(py, vec![capsule])?,
+            )
+    }
     /// Export this to a Python `nanoarrow.ArrayStream`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_array_stream(py, &self.__arrow_c_stream__(py, None)?)

diff --git a/pyo3-arrow/src/datatypes.rs b/pyo3-arrow/src/datatypes.rs
@@ -63,6 +63,16 @@ impl PyDataType {
         )
     }
 
+    /// Export this to a Python `arro3.core.DataType`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let capsule = to_schema_pycapsule(py, &self.0)?;
+        arro3_mod.getattr(intern!(py, "DataType"))?.call_method1(
+            intern!(py, "from_arrow_pycapsule"),
+            PyTuple::new(py, vec![capsule])?,
+        )
+    }
+
     /// Export this to a Python `nanoarrow.Schema`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_schema(py, &self.__arrow_c_schema__(py)?)

diff --git a/pyo3-arrow/src/field.rs b/pyo3-arrow/src/field.rs
@@ -51,6 +51,16 @@ impl PyField {
         )
     }
 
+    /// Export this to a Python `arro3.core.Field`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let capsule = to_schema_pycapsule(py, self.0.as_ref())?;
+        arro3_mod.getattr(intern!(py, "Field"))?.call_method1(
+            intern!(py, "from_arrow_pycapsule"),
+            PyTuple::new(py, vec![capsule])?,
+        )
+    }
+
     /// Export this to a Python `nanoarrow.Schema`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_schema(py, &self.__arrow_c_schema__(py)?)

diff --git a/pyo3-arrow/src/record_batch.rs b/pyo3-arrow/src/record_batch.rs
@@ -90,6 +90,15 @@ impl PyRecordBatch {
         )
     }
 
+    /// Export this to a Python `arro3.core.RecordBatch`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let capsules = Self::to_array_pycapsules(py, self.0.clone(), None)?;
+        arro3_mod
+            .getattr(intern!(py, "RecordBatch"))?
+            .call_method1(intern!(py, "from_arrow_pycapsule"), capsules)
+    }
+
     /// Export this to a Python `nanoarrow.Array`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_array(py, &self.__arrow_c_array__(py, None)?)

diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs
@@ -96,6 +96,24 @@ impl PyRecordBatchReader {
             )
     }
 
+    /// Export this to a Python `arro3.core.RecordBatchReader`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let reader = self
+            .0
+            .lock()
+            .unwrap()
+            .take()
+            .ok_or(PyIOError::new_err("Cannot read from closed stream"))?;
+        let capsule = Self::to_stream_pycapsule(py, reader, None)?;
+        arro3_mod
+            .getattr(intern!(py, "RecordBatchReader"))?
+            .call_method1(
+                intern!(py, "from_arrow_pycapsule"),
+                PyTuple::new(py, vec![capsule])?,
+            )
+    }
+
     /// Export this to a Python `nanoarrow.ArrayStream`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_array_stream(py, &self.__arrow_c_stream__(py, None)?)
@@ -175,8 +193,8 @@ impl PyRecordBatchReader {
 
     // Return self
     // https://stackoverflow.com/a/52056290
-    fn __iter__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        self.to_arro3(py)
+    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
+        slf
     }
 
     fn __next__(&self) -> PyArrowResult<Arro3RecordBatch> {

diff --git a/pyo3-arrow/src/scalar.rs b/pyo3-arrow/src/scalar.rs
@@ -94,6 +94,17 @@ impl PyScalar {
             self.__arrow_c_array__(py, None)?,
         )
     }
+
+    /// Export to an arro3.core.Scalar.
+    ///
+    /// This requires that you depend on arro3-core from your Python package.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let capsules = to_array_pycapsules(py, self.field.clone(), &self.array, None)?;
+        arro3_mod
+            .getattr(intern!(py, "Scalar"))?
+            .call_method1(intern!(py, "from_arrow_pycapsule"), capsules)
+    }
 }
 
 impl Display for PyScalar {

diff --git a/pyo3-arrow/src/schema.rs b/pyo3-arrow/src/schema.rs
@@ -51,6 +51,16 @@ impl PySchema {
         )
     }
 
+    /// Export this to a Python `arro3.core.Schema`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let capsule = to_schema_pycapsule(py, self.0.as_ref())?;
+        arro3_mod.getattr(intern!(py, "Schema"))?.call_method1(
+            intern!(py, "from_arrow_pycapsule"),
+            PyTuple::new(py, vec![capsule])?,
+        )
+    }
+
     /// Export this to a Python `nanoarrow.Schema`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_schema(py, &self.__arrow_c_schema__(py)?)

diff --git a/pyo3-arrow/src/table.rs b/pyo3-arrow/src/table.rs
@@ -87,6 +87,17 @@ impl PyTable {
         )
     }
 
+    /// Export this to a Python `arro3.core.Table`.
+    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
+        let capsule =
+            Self::to_stream_pycapsule(py, self.batches.clone(), self.schema.clone(), None)?;
+        arro3_mod.getattr(intern!(py, "Table"))?.call_method1(
+            intern!(py, "from_arrow_pycapsule"),
+            PyTuple::new(py, vec![capsule])?,
+        )
+    }
+
     /// Export this to a Python `nanoarrow.ArrayStream`.
     pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         to_nanoarrow_array_stream(py, &self.__arrow_c_stream__(py, None)?)