python to_numpy return date/datetime

pola-rs · Dec 9, 2021 · a074b73 · a074b73
1 parent 63f2b73
commit a074b73
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git a/py-polars/docs/source/reference/series.rst b/py-polars/docs/source/reference/series.rst
@@ -76,7 +76,7 @@ Descriptive stats
     Series.is_float
     Series.is_boolean
     Series.is_utf8
-    Series.is_datetime
+    Series.is_datelike
     Series.len
     Series.n_unique
     Series.has_validity

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -714,7 +714,7 @@ def describe(self) -> "pli.DataFrame":
                 "null_count": self.null_count(),
                 "count": self.len(),
             }
-        elif self.is_datetime():
+        elif self.is_datelike():
             # we coerce all to string, because a polars column
             # only has a single dtype and dates: datetime and count: int don't match
             stats = {
@@ -1867,15 +1867,15 @@ def is_numeric(self) -> bool:
             Float64,
         )
 
-    def is_datetime(self) -> bool:
+    def is_datelike(self) -> bool:
         """
-        Check if this Series datatype is a datetime.
+        Check if this Series datatype is datelike.
 
         Examples
         --------
         >>> from datetime import date
         >>> s = pl.Series([date(2021, 1, 1), date(2021, 1, 2), date(2021, 1, 3)])
-        >>> s.is_datetime()
+        >>> s.is_datelike()
         True
 
         """
@@ -2021,13 +2021,25 @@ def to_numpy(
         kwargs
             kwargs will be sent to pyarrow.Array.to_numpy
         """
-        if _PYARROW_AVAILABLE:
+
+        def convert_to_date(arr):  # type: ignore
+            if self.dtype == Date:
+                tp = "datetime64[D]"
+            else:
+                tp = "datetime64[ms]"
+            return arr.astype(tp)
+
+        if _PYARROW_AVAILABLE and not self.is_datelike():
             return self.to_arrow().to_numpy(
                 *args, zero_copy_only=zero_copy_only, **kwargs
             )
         else:
             if not self.has_validity():
+                if self.is_datelike():
+                    return convert_to_date(self.view(ignore_nulls=True))
                 return self.view(ignore_nulls=True)
+            if self.is_datelike():
+                return convert_to_date(self._s.to_numpy())
             return self._s.to_numpy()
 
     def to_arrow(self) -> "pa.Array":

diff --git a/py-polars/tests/test_datelike.py b/py-polars/tests/test_datelike.py
@@ -220,3 +220,13 @@ def test_rows() -> None:
     rows = df.rows()
     assert rows[0][0] == date(2308, 4, 2)
     assert rows[0][1] == datetime(1970, 1, 1, 0, 2, 3, 543000)
+
+
+def test_to_numpy() -> None:
+    s0 = pl.Series("date", [123543, 283478, 1243]).cast(pl.Date)
+    s1 = pl.Series("datetime", [123543, 283478, 1243]).cast(pl.Datetime)
+    assert str(s0.to_numpy()) == "['2308-04-02' '2746-02-20' '1973-05-28']"
+    assert (
+        str(s1.to_numpy()[:2])
+        == "['1970-01-01T00:02:03.543' '1970-01-01T00:04:43.478']"
+    )