Skip to content

Commit

Permalink
perf(python): add future arg to Series.to_arrow
Browse files Browse the repository at this point in the history
  • Loading branch information
ruihe774 committed Jul 2, 2024
1 parent f73937a commit cc6d238
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 4 deletions.
14 changes: 12 additions & 2 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -4342,12 +4342,22 @@ def to_torch(self) -> torch.Tensor:
# tensor.rename(self.name)
return tensor

def to_arrow(self) -> pa.Array:
def to_arrow(self, *, future: bool = False) -> pa.Array:
"""
Return the underlying Arrow array.
If the Series contains only a single chunk this operation is zero copy.
Parameters
----------
future
Setting this to `True` will write Polars' internal data structures that
might not be available by other Arrow implementations.
.. warning::
This functionality is considered **unstable**. It may be changed
at any point without it being considered a breaking change.
Examples
--------
>>> s = pl.Series("a", [1, 2, 3])
Expand All @@ -4360,7 +4370,7 @@ def to_arrow(self) -> pa.Array:
3
]
"""
return self._s.to_arrow()
return self._s.to_arrow(future)

def to_pandas(
self, *, use_pyarrow_extension_array: bool = False, **kwargs: Any
Expand Down
4 changes: 2 additions & 2 deletions py-polars/src/series/export.rs
Original file line number Diff line number Diff line change
Expand Up @@ -145,12 +145,12 @@ impl PySeries {

/// Return the underlying Arrow array.
#[allow(clippy::wrong_self_convention)]
fn to_arrow(&mut self) -> PyResult<PyObject> {
fn to_arrow(&mut self, future: bool) -> PyResult<PyObject> {
self.rechunk(true);
Python::with_gil(|py| {
let pyarrow = py.import_bound("pyarrow")?;

interop::arrow::to_py::to_py_array(self.series.to_arrow(0, false), py, &pyarrow)
interop::arrow::to_py::to_py_array(self.series.to_arrow(0, future), py, &pyarrow)
})
}
}

0 comments on commit cc6d238

Please sign in to comment.