From ffa6451b9e09f65fc3cef0a3e3179db1cbea3da5 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Thu, 27 Jun 2024 19:12:10 -0400 Subject: [PATCH] Add to_pyarrow method --- arro3-compute/src/concat.rs | 2 +- arro3-compute/src/take.rs | 2 +- pyo3-arrow/src/array.rs | 15 +++++++++++++-- pyo3-arrow/src/chunked.rs | 13 ++++++++++++- pyo3-arrow/src/field.rs | 13 ++++++++++++- pyo3-arrow/src/record_batch.rs | 13 ++++++++++++- pyo3-arrow/src/record_batch_reader.rs | 15 ++++++++++++++- pyo3-arrow/src/schema.rs | 13 ++++++++++++- pyo3-arrow/src/table.rs | 13 ++++++++++++- 9 files changed, 89 insertions(+), 10 deletions(-) diff --git a/arro3-compute/src/concat.rs b/arro3-compute/src/concat.rs index f830916..377ea41 100644 --- a/arro3-compute/src/concat.rs +++ b/arro3-compute/src/concat.rs @@ -7,5 +7,5 @@ pub fn concat(input: PyChunkedArray) -> PyArrowResult { let (chunks, field) = input.into_inner(); let array_refs = chunks.iter().map(|arr| arr.as_ref()).collect::>(); 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 index 12fb71b..09bf6b5 100644 --- a/arro3-compute/src/take.rs +++ b/arro3-compute/src/take.rs @@ -7,5 +7,5 @@ pub fn take(values: PyArray, indices: PyArray) -> PyArrowResult { 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 index 3e9f728..a06094f 100644 --- a/pyo3-arrow/src/array.rs +++ b/pyo3-arrow/src/array.rs @@ -42,10 +42,10 @@ 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 { + pub fn to_arro3(&self, py: Python) -> PyArrowResult { 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"), @@ -53,6 +53,17 @@ impl PyArray { )?; Ok(core_obj.to_object(py)) } + + /// Export to a pyarrow.Array + /// + /// Requires pyarrow >=14 + pub fn to_pyarrow(self, py: Python) -> PyArrowResult { + 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 index 5baf268..6214a34 100644 --- 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 { + pub fn to_arro3(&self, py: Python) -> PyArrowResult { 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 { + 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 index d38f55d..9cb2e20 100644 --- a/pyo3-arrow/src/field.rs +++ b/pyo3-arrow/src/field.rs @@ -23,7 +23,7 @@ impl PyField { } /// Export this to a Python `arro3.core.Field`. - pub fn to_python(&self, py: Python) -> PyArrowResult { + pub fn to_arro3(&self, py: Python) -> PyArrowResult { 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"), @@ -31,6 +31,17 @@ impl PyField { )?; Ok(core_obj.to_object(py)) } + + /// Export to a pyarrow.Field + /// + /// Requires pyarrow >=14 + pub fn to_pyarrow(self, py: Python) -> PyArrowResult { + 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 for FieldRef { diff --git a/pyo3-arrow/src/record_batch.rs b/pyo3-arrow/src/record_batch.rs index f928749..24dfac5 100644 --- 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 { + pub fn to_arro3(&self, py: Python) -> PyArrowResult { 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 { + 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 for PyRecordBatch { diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index 0aee31e..0622d2e 100644 --- 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 { + pub fn to_arro3(&mut self, py: Python) -> PyArrowResult { 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 { + 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 index 8174406..ef4c43f 100644 --- a/pyo3-arrow/src/schema.rs +++ b/pyo3-arrow/src/schema.rs @@ -23,7 +23,7 @@ impl PySchema { } /// Export this to a Python `arro3.core.Schema`. - pub fn to_python(&self, py: Python) -> PyArrowResult { + pub fn to_arro3(&self, py: Python) -> PyArrowResult { 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"), @@ -31,6 +31,17 @@ impl PySchema { )?; Ok(core_obj.to_object(py)) } + + /// Export to a pyarrow.Schema + /// + /// Requires pyarrow >=14 + pub fn to_pyarrow(self, py: Python) -> PyArrowResult { + 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 for SchemaRef { diff --git a/pyo3-arrow/src/table.rs b/pyo3-arrow/src/table.rs index 09a7690..cc8419c 100644 --- a/pyo3-arrow/src/table.rs +++ b/pyo3-arrow/src/table.rs @@ -37,7 +37,7 @@ impl PyTable { } /// Export this to a Python `arro3.core.Table`. - pub fn to_python(&self, py: Python) -> PyArrowResult { + pub fn to_arro3(&self, py: Python) -> PyArrowResult { 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"), @@ -45,6 +45,17 @@ impl PyTable { )?; Ok(core_obj.to_object(py)) } + + /// Export to a pyarrow.Table + /// + /// Requires pyarrow >=14 + pub fn to_pyarrow(self, py: Python) -> PyArrowResult { + 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]