Skip to content

Commit

Permalink
python output date/datetime in to_list and to_dict
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Dec 4, 2021
1 parent b830827 commit 10afcf6
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 3 deletions.
11 changes: 11 additions & 0 deletions py-polars/polars/internals/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -1735,6 +1735,17 @@ def cast(
pl_dtype = py_type_to_dtype(dtype)
return wrap_s(self._s.cast(str(pl_dtype), strict))

def to_physical(self) -> "Series":
"""
Cast to physical representation of the logical dtype.
Date -> Int32
Datetime -> Int64
Time -> Int64
other -> other
"""
return wrap_s(self._s.to_physical)

def to_list(self, use_pyarrow: bool = False) -> tp.List[Optional[Any]]:
"""
Convert this Series to a Python List. This operation clones data.
Expand Down
34 changes: 33 additions & 1 deletion py-polars/src/conversion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ use pyo3::basic::CompareOp;
use pyo3::conversion::{FromPyObject, IntoPy};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PySequence};
use pyo3::types::{PyDict, PyList, PySequence};
use pyo3::{PyAny, PyResult};
use std::fmt::{Display, Formatter};
use std::hash::{Hash, Hasher};
Expand Down Expand Up @@ -196,6 +196,38 @@ impl ToPyObject for Wrap<AnyValue<'_>> {
}
}

impl ToPyObject for Wrap<&DatetimeChunked> {
fn to_object(&self, py: Python) -> PyObject {
let pl = PyModule::import(py, "polars").unwrap();
let pli = pl.getattr("internals").unwrap();
let m_series = pli.getattr("series").unwrap();
let convert = m_series.getattr("_to_python_datetime").unwrap();
let py_date_dtype = pl.getattr("Datetime").unwrap();

let iter = self
.0
.into_iter()
.map(|opt_v| opt_v.map(|v| convert.call1((v, py_date_dtype)).unwrap()));
PyList::new(py, iter).into_py(py)
}
}

impl ToPyObject for Wrap<&DateChunked> {
fn to_object(&self, py: Python) -> PyObject {
let pl = PyModule::import(py, "polars").unwrap();
let pli = pl.getattr("internals").unwrap();
let m_series = pli.getattr("series").unwrap();
let convert = m_series.getattr("_to_python_datetime").unwrap();
let py_date_dtype = pl.getattr("Date").unwrap();

let iter = self
.0
.into_iter()
.map(|opt_v| opt_v.map(|v| convert.call1((v, py_date_dtype)).unwrap()));
PyList::new(py, iter).into_py(py)
}
}

impl<'s> FromPyObject<'s> for Wrap<AnyValue<'s>> {
fn extract(ob: &'s PyAny) -> PyResult<Self> {
if let Ok(v) = ob.extract::<i64>() {
Expand Down
15 changes: 13 additions & 2 deletions py-polars/src/series.rs
Original file line number Diff line number Diff line change
Expand Up @@ -632,6 +632,11 @@ impl PySeries {
self.series.len()
}

pub fn to_physical(&self) -> Self {
let s = self.series.to_physical_repr().into_owned();
s.into()
}

pub fn to_list(&self) -> PyObject {
let gil = Python::acquire_gil();
let python = gil.python();
Expand Down Expand Up @@ -682,8 +687,14 @@ impl PySeries {
}
v
}
DataType::Date => PyList::new(python, &series.date().unwrap().0),
DataType::Datetime => PyList::new(python, &series.datetime().unwrap().0),
DataType::Date => {
let ca = series.date().unwrap();
return Wrap(ca).to_object(python);
}
DataType::Datetime => {
let ca = series.datetime().unwrap();
return Wrap(ca).to_object(python);
}
dt => primitive_to_list(dt, series),
};
pylist.to_object(python)
Expand Down
12 changes: 12 additions & 0 deletions py-polars/tests/test_datelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,3 +199,15 @@ def test_timezone() -> None:

# timezones have no effect, i.e. `s` equals `tz_s`
assert s.series_equal(tz_s)


def test_to_list() -> None:

s = pl.Series("date", [123543, 283478, 1243]).cast(pl.Date)

out = s.to_list()
assert out[0] == date(2308, 4, 2)

s = pl.Series("datetime", [123543, 283478, 1243]).cast(pl.Datetime)
out = s.to_list()
assert out[0] == datetime(1970, 1, 1, 0, 2, 3, 543000)

0 comments on commit 10afcf6

Please sign in to comment.