feat(python): Support Categorical/Enum in Series.to_numpy (#14275)

pola-rs · Feb 5, 2024 · f7e9a20 · f7e9a20
1 parent d959fa1
commit f7e9a20
Show file tree

Hide file tree

Showing 5 changed files with 296 additions and 152 deletions.
diff --git a/py-polars/polars/series/series.py b/py-polars/polars/series/series.py
@@ -4325,7 +4325,7 @@ def to_numpy(
         """
 
         def raise_no_zero_copy() -> None:
-            if zero_copy_only:
+            if zero_copy_only and not self.is_empty():
                 msg = "cannot return a zero-copy array"
                 raise ValueError(msg)
 

diff --git a/py-polars/src/series/export.rs b/py-polars/src/series/export.rs
@@ -195,6 +195,11 @@ impl PySeries {
                 let np_arr = PyArray1::from_iter(py, ca.into_iter().map(|s| s.into_py(py)));
                 np_arr.into_py(py)
             },
+            Categorical(_, _) | Enum(_, _) => {
+                let ca = s.categorical().unwrap();
+                let np_arr = PyArray1::from_iter(py, ca.iter_str().map(|s| s.into_py(py)));
+                np_arr.into_py(py)
+            },
             #[cfg(feature = "object")]
             Object(_, _) => {
                 let ca = s

diff --git a/py-polars/tests/unit/interop/numpy/test_numpy.py b/py-polars/tests/unit/interop/numpy/test_numpy.py
@@ -76,3 +76,14 @@ def test_to_numpy_zero_copy_path() -> None:
     x = df.to_numpy()
     assert x.flags["F_CONTIGUOUS"]
     assert str(x[0, :]) == "[1. 2. 1. 1. 1.]"
+
+
+def test_numpy_disambiguation() -> None:
+    a = np.array([1, 2])
+    df = pl.DataFrame({"a": a})
+    result = df.with_columns(b=a).to_dict(as_series=False)  # type: ignore[arg-type]
+    expected = {
+        "a": [1, 2],
+        "b": [1, 2],
+    }
+    assert result == expected
diff --git a/py-polars/tests/unit/interop/numpy/test_to_numpy_df.py b/py-polars/tests/unit/interop/numpy/test_to_numpy_df.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from decimal import Decimal as D
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -101,3 +102,20 @@ def test__array__() -> None:
     expected_array = np.array([[1, 1], [2, 2], [3, 3]], dtype=np.uint8)
     assert_array_equal(out_array, expected_array)
     assert out_array.flags["F_CONTIGUOUS"] is True
+
+
+def test_numpy_preserve_uint64_4112() -> None:
+    df = pl.DataFrame({"a": [1, 2, 3]}).with_columns(pl.col("a").hash())
+    assert df.to_numpy().dtype == np.dtype("uint64")
+    assert df.to_numpy(structured=True).dtype == np.dtype([("a", "uint64")])
+
+
+@pytest.mark.parametrize("use_pyarrow", [True, False])
+def test_df_to_numpy_decimal(use_pyarrow: bool) -> None:
+    decimal_data = [D("1.234"), D("2.345"), D("-3.456")]
+    df = pl.Series("n", decimal_data).to_frame()
+
+    result = df.to_numpy(use_pyarrow=use_pyarrow)
+
+    expected = np.array(decimal_data).reshape((-1, 1))
+    assert_array_equal(result, expected)