Skip to content

Commit

Permalink
refactor(python): Update Series.to_numpy to handle Decimal/Time typ…
Browse files Browse the repository at this point in the history
…es in Rust (#14296)
  • Loading branch information
stinodego authored Feb 7, 2024
1 parent 4827d73 commit f7c8660
Show file tree
Hide file tree
Showing 4 changed files with 58 additions and 39 deletions.
12 changes: 3 additions & 9 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -4358,20 +4358,14 @@ def temporal_dtype_to_numpy(dtype: PolarsDataType) -> Any:
if (
use_pyarrow
and _PYARROW_AVAILABLE
and dtype != Object
and (dtype == Time or not dtype.is_temporal())
and dtype not in (Object, Datetime, Duration, Date)
):
return self.to_arrow().to_numpy(
zero_copy_only=zero_copy_only, writable=writable
)

if dtype in (Time, Decimal):
# There are no native NumPy "time" or "decimal" dtypes
raise_no_zero_copy()
return np.array(self.to_list(), dtype="object", copy=False)

if self.null_count() == 0:
if dtype.is_numeric():
if dtype.is_integer() or dtype.is_float():
np_array = self._view(ignore_nulls=True)
elif dtype == Boolean:
raise_no_zero_copy()
Expand All @@ -4390,7 +4384,7 @@ def temporal_dtype_to_numpy(dtype: PolarsDataType) -> Any:
else:
raise_no_zero_copy()
np_array = self._s.to_numpy()
if dtype.is_temporal():
if dtype in (Datetime, Duration, Date):
np_dtype = temporal_dtype_to_numpy(dtype)
np_array = np_array.view(np_dtype)

Expand Down
70 changes: 41 additions & 29 deletions py-polars/src/conversion/chunked_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -141,16 +141,21 @@ impl ToPyObject for Wrap<&DatetimeChunked> {

impl ToPyObject for Wrap<&TimeChunked> {
fn to_object(&self, py: Python) -> PyObject {
let utils = UTILS.as_ref(py);
let convert = utils.getattr(intern!(py, "_to_python_time")).unwrap();
let iter = self
.0
.into_iter()
.map(|opt_v| opt_v.map(|v| convert.call1((v,)).unwrap()));
let iter = time_to_pyobject_iter(py, self.0);
PyList::new(py, iter).into_py(py)
}
}

pub(crate) fn time_to_pyobject_iter<'a>(
py: Python<'a>,
ca: &'a TimeChunked,
) -> impl ExactSizeIterator<Item = Option<&'a PyAny>> {
let utils = UTILS.as_ref(py);
let convert = utils.getattr(intern!(py, "_to_python_time")).unwrap();
ca.0.into_iter()
.map(|opt_v| opt_v.map(|v| convert.call1((v,)).unwrap()))
}

impl ToPyObject for Wrap<&DateChunked> {
fn to_object(&self, py: Python) -> PyObject {
let utils = UTILS.as_ref(py);
Expand All @@ -165,29 +170,36 @@ impl ToPyObject for Wrap<&DateChunked> {

impl ToPyObject for Wrap<&DecimalChunked> {
fn to_object(&self, py: Python) -> PyObject {
let utils = UTILS.as_ref(py);
let convert = utils.getattr(intern!(py, "_to_python_decimal")).unwrap();
let py_scale = (-(self.0.scale() as i32)).to_object(py);
// if we don't know precision, the only safe bet is to set it to 39
let py_precision = self.0.precision().unwrap_or(39).to_object(py);
let iter = self.0.into_iter().map(|opt_v| {
opt_v.map(|v| {
// TODO! use AnyValue so that we have a single impl.
const N: usize = 3;
let mut buf = [0_u128; N];
let n_digits = decimal_to_digits(v.abs(), &mut buf);
let buf = unsafe {
std::slice::from_raw_parts(
buf.as_slice().as_ptr() as *const u8,
N * std::mem::size_of::<u128>(),
)
};
let digits = PyTuple::new(py, buf.iter().take(n_digits));
convert
.call1((v.is_negative() as u8, digits, &py_precision, &py_scale))
.unwrap()
})
});
let iter = decimal_to_pyobject_iter(py, self.0);
PyList::new(py, iter).into_py(py)
}
}

pub(crate) fn decimal_to_pyobject_iter<'a>(
py: Python<'a>,
ca: &'a DecimalChunked,
) -> impl ExactSizeIterator<Item = Option<&'a PyAny>> {
let utils = UTILS.as_ref(py);
let convert = utils.getattr(intern!(py, "_to_python_decimal")).unwrap();
let py_scale = (-(ca.scale() as i32)).to_object(py);
// if we don't know precision, the only safe bet is to set it to 39
let py_precision = ca.precision().unwrap_or(39).to_object(py);
ca.into_iter().map(move |opt_v| {
opt_v.map(|v| {
// TODO! use AnyValue so that we have a single impl.
const N: usize = 3;
let mut buf = [0_u128; N];
let n_digits = decimal_to_digits(v.abs(), &mut buf);
let buf = unsafe {
std::slice::from_raw_parts(
buf.as_slice().as_ptr() as *const u8,
N * std::mem::size_of::<u128>(),
)
};
let digits = PyTuple::new(py, buf.iter().take(n_digits));
convert
.call1((v.is_negative() as u8, digits, &py_precision, &py_scale))
.unwrap()
})
})
}
2 changes: 1 addition & 1 deletion py-polars/src/conversion/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
pub(crate) mod any_value;
mod chunked_array;
pub(crate) mod chunked_array;

use std::fmt::{Display, Formatter};
use std::hash::{Hash, Hasher};
Expand Down
13 changes: 13 additions & 0 deletions py-polars/src/series/export.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use polars_core::prelude::*;
use pyo3::prelude::*;
use pyo3::types::PyList;

use crate::conversion::chunked_array::{decimal_to_pyobject_iter, time_to_pyobject_iter};
use crate::error::PyPolarsErr;
use crate::prelude::{ObjectValue, *};
use crate::{arrow_interop, raise_err, PySeries};
Expand Down Expand Up @@ -185,6 +186,12 @@ impl PySeries {
},
Date => date_series_to_numpy(py, s),
Datetime(_, _) | Duration(_) => temporal_series_to_numpy(py, s),
Time => {
let ca = s.time().unwrap();
let iter = time_to_pyobject_iter(py, ca);
let np_arr = PyArray1::from_iter(py, iter.map(|v| v.into_py(py)));
np_arr.into_py(py)
},
String => {
let ca = s.str().unwrap();
let np_arr = PyArray1::from_iter(py, ca.into_iter().map(|s| s.into_py(py)));
Expand All @@ -200,6 +207,12 @@ impl PySeries {
let np_arr = PyArray1::from_iter(py, ca.iter_str().map(|s| s.into_py(py)));
np_arr.into_py(py)
},
Decimal(_, _) => {
let ca = s.decimal().unwrap();
let iter = decimal_to_pyobject_iter(py, ca);
let np_arr = PyArray1::from_iter(py, iter.map(|v| v.into_py(py)));
np_arr.into_py(py)
},
#[cfg(feature = "object")]
Object(_, _) => {
let ca = s
Expand Down

0 comments on commit f7c8660

Please sign in to comment.