Skip to content

Commit 5218810

Browse files
mbrobbelmbrobbel
committed
Return Bound values in to_pyarrow and into_pyarrow 
1 parent 43c18c4 commit 5218810

File tree

2 files changed

+25
-24
lines changed
  • arrow-pyarrow-integration-testing/src
  • arrow-pyarrow/src

2 files changed

+25
-24
lines changed

‎arrow-pyarrow-integration-testing/src/lib.rs‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ fn to_py_err(err: ArrowError) -> PyErr {
4141

4242
/// Returns `array + array` of an int64 array.
4343
#[pyfunction]
44-
fn double(array: &Bound<PyAny>, py: Python) -> PyResult<Py<PyAny>> {
44+
fn double<'py>(array: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
45+
let py = array.py();
4546
// import
4647
let array = make_array(ArrayData::from_pyarrow_bound(array)?);
4748

@@ -61,21 +62,24 @@ fn double(array: &Bound<PyAny>, py: Python) -> PyResult<Py<PyAny>> {
6162
/// calls a lambda function that receives and returns an array
6263
/// whose result must be the array multiplied by two
6364
#[pyfunction]
64-
fn double_py(lambda: &Bound<PyAny>, py: Python) -> PyResult<bool> {
65+
fn double_py(lambda: &Bound<PyAny>) -> PyResult<bool> {
6566
// create
6667
let array = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
6768
let expected = Arc::new(Int64Array::from(vec![Some(2), None, Some(6)])) as ArrayRef;
6869

6970
// to py
70-
let pyarray = array.to_data().to_pyarrow(py)?;
71+
let pyarray = array.to_data().to_pyarrow(lambda.py())?;
7172
let pyarray = lambda.call1((pyarray,))?;
7273
let array = make_array(ArrayData::from_pyarrow_bound(&pyarray)?);
7374

7475
Ok(array == expected)
7576
}
7677

7778
#[pyfunction]
78-
fn make_empty_array(datatype: PyArrowType<DataType>, py: Python) -> PyResult<Py<PyAny>> {
79+
fn make_empty_array<'py>(
80+
datatype: PyArrowType<DataType>,
81+
py: Python<'py>,
82+
) -> PyResult<Bound<'py, PyAny>> {
7983
let array = new_empty_array(&datatype.0);
8084

8185
array.to_data().to_pyarrow(py)
@@ -95,7 +99,7 @@ fn substring(array: PyArrowType<ArrayData>, start: i64) -> PyResult<PyArrowType<
9599

96100
/// Returns the concatenate
97101
#[pyfunction]
98-
fn concatenate(array: PyArrowType<ArrayData>, py: Python) -> PyResult<Py<PyAny>> {
102+
fn concatenate<'py>(array: PyArrowType<ArrayData>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
99103
let array = make_array(array.0);
100104

101105
// concat

‎arrow-pyarrow/src/lib.rs‎

Lines changed: 16 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -95,17 +95,17 @@ pub trait FromPyArrow: Sized {
9595
/// Create a new PyArrow object from a arrow-rs type.
9696
pub trait ToPyArrow {
9797
/// Convert the implemented type into a Python object without consuming it.
98-
fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>>;
98+
fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>>;
9999
}
100100

101101
/// Convert an arrow-rs type into a PyArrow object.
102102
pub trait IntoPyArrow {
103103
/// Convert the implemented type into a Python object while consuming it.
104-
fn into_pyarrow(self, py: Python) -> PyResult<Py<PyAny>>;
104+
fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>>;
105105
}
106106

107107
impl<T: ToPyArrow> IntoPyArrow for T {
108-
fn into_pyarrow(self, py: Python) -> PyResult<Py<PyAny>> {
108+
fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
109109
self.to_pyarrow(py)
110110
}
111111
}
@@ -172,7 +172,7 @@ impl FromPyArrow for DataType {
172172
}
173173

174174
impl ToPyArrow for DataType {
175-
fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {
175+
fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
176176
let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?;
177177
let c_schema_ptr = &c_schema as *const FFI_ArrowSchema;
178178
let module = py.import("pyarrow")?;
@@ -208,7 +208,7 @@ impl FromPyArrow for Field {
208208
}
209209

210210
impl ToPyArrow for Field {
211-
fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {
211+
fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
212212
let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?;
213213
let c_schema_ptr = &c_schema as *const FFI_ArrowSchema;
214214
let module = py.import("pyarrow")?;
@@ -244,7 +244,7 @@ impl FromPyArrow for Schema {
244244
}
245245

246246
impl ToPyArrow for Schema {
247-
fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {
247+
fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
248248
let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?;
249249
let c_schema_ptr = &c_schema as *const FFI_ArrowSchema;
250250
let module = py.import("pyarrow")?;
@@ -303,7 +303,7 @@ impl FromPyArrow for ArrayData {
303303
}
304304

305305
impl ToPyArrow for ArrayData {
306-
fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {
306+
fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
307307
let array = FFI_ArrowArray::new(self);
308308
let schema = FFI_ArrowSchema::try_from(self.data_type()).map_err(to_py_err)?;
309309

@@ -316,7 +316,7 @@ impl ToPyArrow for ArrayData {
316316
addr_of!(schema) as Py_uintptr_t,
317317
),
318318
)?;
319-
Ok(array.unbind())
319+
Ok(array)
320320
}
321321
}
322322

@@ -328,12 +328,12 @@ impl<T: FromPyArrow> FromPyArrow for Vec<T> {
328328
}
329329

330330
impl<T: ToPyArrow> ToPyArrow for Vec<T> {
331-
fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {
331+
fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
332332
let values = self
333333
.iter()
334334
.map(|v| v.to_pyarrow(py))
335335
.collect::<PyResult<Vec<_>>>()?;
336-
Ok(PyList::new(py, values)?.unbind().into())
336+
Ok(PyList::new(py, values)?.into_any())
337337
}
338338
}
339339

@@ -412,12 +412,12 @@ impl FromPyArrow for RecordBatch {
412412
}
413413

414414
impl ToPyArrow for RecordBatch {
415-
fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {
415+
fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
416416
// Workaround apache/arrow#37669 by returning RecordBatchIterator
417417
let reader = RecordBatchIterator::new(vec![Ok(self.clone())], self.schema());
418418
let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
419419
let py_reader = reader.into_pyarrow(py)?;
420-
py_reader.call_method0(py, "read_next_batch")
420+
py_reader.call_method0("read_next_batch")
421421
}
422422
}
423423

@@ -463,7 +463,7 @@ impl FromPyArrow for ArrowArrayStreamReader {
463463
impl IntoPyArrow for Box<dyn RecordBatchReader + Send> {
464464
// We can't implement `ToPyArrow` for `T: RecordBatchReader + Send` because
465465
// there is already a blanket implementation for `T: ToPyArrow`.
466-
fn into_pyarrow(self, py: Python) -> PyResult<Py<PyAny>> {
466+
fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
467467
let mut stream = FFI_ArrowArrayStream::new(self);
468468

469469
let stream_ptr = (&mut stream) as *mut FFI_ArrowArrayStream;
@@ -472,13 +472,13 @@ impl IntoPyArrow for Box<dyn RecordBatchReader + Send> {
472472
let args = PyTuple::new(py, [stream_ptr as Py_uintptr_t])?;
473473
let reader = class.call_method1("_import_from_c", args)?;
474474

475-
Ok(Py::from(reader))
475+
Ok(reader)
476476
}
477477
}
478478

479479
/// Convert a [`ArrowArrayStreamReader`] into a `pyarrow.RecordBatchReader`.
480480
impl IntoPyArrow for ArrowArrayStreamReader {
481-
fn into_pyarrow(self, py: Python) -> PyResult<Py<PyAny>> {
481+
fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
482482
let boxed: Box<dyn RecordBatchReader + Send> = Box::new(self);
483483
boxed.into_pyarrow(py)
484484
}
@@ -506,10 +506,7 @@ impl<'py, T: IntoPyArrow> IntoPyObject<'py> for PyArrowType<T> {
506506
type Error = PyErr;
507507

508508
fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, PyErr> {
509-
match self.0.into_pyarrow(py) {
510-
Ok(obj) => Result::Ok(obj.into_bound(py)),
511-
Err(err) => Result::Err(err),
512-
}
509+
self.0.into_pyarrow(py)
513510
}
514511
}
515512

0 commit comments

Comments
 (0)