Add to_pyarrow method on FFI structs (#28)

kylebarron · Jun 27, 2024 · f45419a · f45419a
1 parent 70cb725
commit f45419a
Show file tree

Hide file tree

Showing 9 changed files with 89 additions and 10 deletions.
diff --git a/arro3-compute/src/concat.rs b/arro3-compute/src/concat.rs
@@ -7,5 +7,5 @@ pub fn concat(input: PyChunkedArray) -> PyArrowResult<PyObject> {
     let (chunks, field) = input.into_inner();
     let array_refs = chunks.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
     let concatted = arrow_select::concat::concat(array_refs.as_slice())?;
-    Python::with_gil(|py| PyArray::new(concatted, field).to_python(py))
+    Python::with_gil(|py| PyArray::new(concatted, field).to_arro3(py))
 }
diff --git a/arro3-compute/src/take.rs b/arro3-compute/src/take.rs
@@ -7,5 +7,5 @@ pub fn take(values: PyArray, indices: PyArray) -> PyArrowResult<PyObject> {
     let (values_arr, values_field) = values.into_inner();
     let (indices, _) = indices.into_inner();
     let result = arrow_select::take::take(&values_arr, &indices, None)?;
-    Python::with_gil(|py| PyArray::new(result, values_field).to_python(py))
+    Python::with_gil(|py| PyArray::new(result, values_field).to_arro3(py))
 }
diff --git a/pyo3-arrow/src/array.rs b/pyo3-arrow/src/array.rs
@@ -42,17 +42,28 @@ impl PyArray {
         (self.array, self.field)
     }
 
-    /// Export to a Python Array.
+    /// Export to an arro3.core.Array.
     ///
     /// This requires that you depend on arro3-core from your Python package.
-    pub fn to_python(&self, py: Python) -> PyArrowResult<PyObject> {
+    pub fn to_arro3(&self, py: Python) -> PyArrowResult<PyObject> {
         let arro3_mod = py.import_bound(intern!(py, "arro3.core"))?;
         let core_obj = arro3_mod.getattr(intern!(py, "Array"))?.call_method1(
             intern!(py, "from_arrow_pycapsule"),
             self.__arrow_c_array__(py, None)?,
         )?;
         Ok(core_obj.to_object(py))
     }
+
+    /// Export to a pyarrow.Array
+    ///
+    /// Requires pyarrow >=14
+    pub fn to_pyarrow(self, py: Python) -> PyArrowResult<PyObject> {
+        let pyarrow_mod = py.import_bound(intern!(py, "pyarrow"))?;
+        let pyarrow_obj = pyarrow_mod
+            .getattr(intern!(py, "array"))?
+            .call1(PyTuple::new_bound(py, vec![self.into_py(py)]))?;
+        Ok(pyarrow_obj.to_object(py))
+    }
 }
 
 #[pymethods]

diff --git a/pyo3-arrow/src/chunked.rs b/pyo3-arrow/src/chunked.rs
@@ -41,7 +41,7 @@ impl PyChunkedArray {
     }
 
     /// Export this to a Python `arro3.core.ChunkedArray`.
-    pub fn to_python(&self, py: Python) -> PyArrowResult<PyObject> {
+    pub fn to_arro3(&self, py: Python) -> PyArrowResult<PyObject> {
         let arro3_mod = py.import_bound(intern!(py, "arro3.core"))?;
         let core_obj = arro3_mod
             .getattr(intern!(py, "ChunkedArray"))?
@@ -51,6 +51,17 @@ impl PyChunkedArray {
             )?;
         Ok(core_obj.to_object(py))
     }
+
+    /// Export to a pyarrow.ChunkedArray
+    ///
+    /// Requires pyarrow >=14
+    pub fn to_pyarrow(self, py: Python) -> PyArrowResult<PyObject> {
+        let pyarrow_mod = py.import_bound(intern!(py, "pyarrow"))?;
+        let pyarrow_obj = pyarrow_mod
+            .getattr(intern!(py, "chunked_array"))?
+            .call1(PyTuple::new_bound(py, vec![self.into_py(py)]))?;
+        Ok(pyarrow_obj.to_object(py))
+    }
 }
 
 #[pymethods]

diff --git a/pyo3-arrow/src/field.rs b/pyo3-arrow/src/field.rs
@@ -23,14 +23,25 @@ impl PyField {
     }
 
     /// Export this to a Python `arro3.core.Field`.
-    pub fn to_python(&self, py: Python) -> PyArrowResult<PyObject> {
+    pub fn to_arro3(&self, py: Python) -> PyArrowResult<PyObject> {
         let arro3_mod = py.import_bound(intern!(py, "arro3.core"))?;
         let core_obj = arro3_mod.getattr(intern!(py, "Field"))?.call_method1(
             intern!(py, "from_arrow_pycapsule"),
             PyTuple::new_bound(py, vec![self.__arrow_c_schema__(py)?]),
         )?;
         Ok(core_obj.to_object(py))
     }
+
+    /// Export to a pyarrow.Field
+    ///
+    /// Requires pyarrow >=14
+    pub fn to_pyarrow(self, py: Python) -> PyArrowResult<PyObject> {
+        let pyarrow_mod = py.import_bound(intern!(py, "pyarrow"))?;
+        let pyarrow_obj = pyarrow_mod
+            .getattr(intern!(py, "field"))?
+            .call1(PyTuple::new_bound(py, vec![self.into_py(py)]))?;
+        Ok(pyarrow_obj.to_object(py))
+    }
 }
 
 impl From<PyField> for FieldRef {

diff --git a/pyo3-arrow/src/record_batch.rs b/pyo3-arrow/src/record_batch.rs
@@ -26,7 +26,7 @@ impl PyRecordBatch {
     }
 
     /// Export this to a Python `arro3.core.RecordBatch`.
-    pub fn to_python(&self, py: Python) -> PyArrowResult<PyObject> {
+    pub fn to_arro3(&self, py: Python) -> PyArrowResult<PyObject> {
         let arro3_mod = py.import_bound(intern!(py, "arro3.core"))?;
         let core_obj = arro3_mod
             .getattr(intern!(py, "RecordBatch"))?
@@ -36,6 +36,17 @@ impl PyRecordBatch {
             )?;
         Ok(core_obj.to_object(py))
     }
+
+    /// Export to a pyarrow.RecordBatch
+    ///
+    /// Requires pyarrow >=14
+    pub fn to_pyarrow(self, py: Python) -> PyArrowResult<PyObject> {
+        let pyarrow_mod = py.import_bound(intern!(py, "pyarrow"))?;
+        let pyarrow_obj = pyarrow_mod
+            .getattr(intern!(py, "record_batch"))?
+            .call1(PyTuple::new_bound(py, vec![self.into_py(py)]))?;
+        Ok(pyarrow_obj.to_object(py))
+    }
 }
 
 impl From<RecordBatch> for PyRecordBatch {

diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs
@@ -65,7 +65,7 @@ impl PyRecordBatchReader {
     }
 
     /// Export this to a Python `arro3.core.RecordBatchReader`.
-    pub fn to_python(&mut self, py: Python) -> PyArrowResult<PyObject> {
+    pub fn to_arro3(&mut self, py: Python) -> PyArrowResult<PyObject> {
         let arro3_mod = py.import_bound(intern!(py, "arro3.core"))?;
         let core_obj = arro3_mod
             .getattr(intern!(py, "RecordBatchReader"))?
@@ -75,6 +75,19 @@ impl PyRecordBatchReader {
             )?;
         Ok(core_obj.to_object(py))
     }
+
+    /// Export to a pyarrow.RecordBatchReader
+    ///
+    /// Requires pyarrow >=15
+    pub fn to_pyarrow(self, py: Python) -> PyArrowResult<PyObject> {
+        let pyarrow_mod = py.import_bound(intern!(py, "pyarrow"))?;
+        let record_batch_reader_class = pyarrow_mod.getattr(intern!(py, "RecordBatchReader"))?;
+        let pyarrow_obj = record_batch_reader_class.call_method1(
+            intern!(py, "from_stream"),
+            PyTuple::new_bound(py, vec![self.into_py(py)]),
+        )?;
+        Ok(pyarrow_obj.to_object(py))
+    }
 }
 
 #[pymethods]

diff --git a/pyo3-arrow/src/schema.rs b/pyo3-arrow/src/schema.rs
@@ -23,14 +23,25 @@ impl PySchema {
     }
 
     /// Export this to a Python `arro3.core.Schema`.
-    pub fn to_python(&self, py: Python) -> PyArrowResult<PyObject> {
+    pub fn to_arro3(&self, py: Python) -> PyArrowResult<PyObject> {
         let arro3_mod = py.import_bound(intern!(py, "arro3.core"))?;
         let core_obj = arro3_mod.getattr(intern!(py, "Schema"))?.call_method1(
             intern!(py, "from_arrow_pycapsule"),
             PyTuple::new_bound(py, vec![self.__arrow_c_schema__(py)?]),
         )?;
         Ok(core_obj.to_object(py))
     }
+
+    /// Export to a pyarrow.Schema
+    ///
+    /// Requires pyarrow >=14
+    pub fn to_pyarrow(self, py: Python) -> PyArrowResult<PyObject> {
+        let pyarrow_mod = py.import_bound(intern!(py, "pyarrow"))?;
+        let pyarrow_obj = pyarrow_mod
+            .getattr(intern!(py, "schema"))?
+            .call1(PyTuple::new_bound(py, vec![self.into_py(py)]))?;
+        Ok(pyarrow_obj.to_object(py))
+    }
 }
 
 impl From<PySchema> for SchemaRef {

diff --git a/pyo3-arrow/src/table.rs b/pyo3-arrow/src/table.rs
@@ -37,14 +37,25 @@ impl PyTable {
     }
 
     /// Export this to a Python `arro3.core.Table`.
-    pub fn to_python(&self, py: Python) -> PyArrowResult<PyObject> {
+    pub fn to_arro3(&self, py: Python) -> PyArrowResult<PyObject> {
         let arro3_mod = py.import_bound(intern!(py, "arro3.core"))?;
         let core_obj = arro3_mod.getattr(intern!(py, "Table"))?.call_method1(
             intern!(py, "from_arrow_pycapsule"),
             PyTuple::new_bound(py, vec![self.__arrow_c_stream__(py, None)?]),
         )?;
         Ok(core_obj.to_object(py))
     }
+
+    /// Export to a pyarrow.Table
+    ///
+    /// Requires pyarrow >=14
+    pub fn to_pyarrow(self, py: Python) -> PyArrowResult<PyObject> {
+        let pyarrow_mod = py.import_bound(intern!(py, "pyarrow"))?;
+        let pyarrow_obj = pyarrow_mod
+            .getattr(intern!(py, "table"))?
+            .call1(PyTuple::new_bound(py, vec![self.into_py(py)]))?;
+        Ok(pyarrow_obj.to_object(py))
+    }
 }
 
 #[pymethods]