python: to_numpy use first type as supertype (#4116)

pola-rs · Jul 22, 2022 · d394f83 · d394f83
1 parent 91f0ae5
commit d394f83
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/py-polars/src/dataframe.rs b/py-polars/src/dataframe.rs
@@ -524,11 +524,17 @@ impl PyDataFrame {
     }
 
     pub fn to_numpy(&self, py: Python) -> Option<PyObject> {
-        let mut st = DataType::Int8;
+        let mut st = None;
         for s in self.df.iter() {
             let dt_i = s.dtype();
-            st = get_supertype(&st, dt_i).ok()?;
+            match st {
+                None => st = Some(dt_i.clone()),
+                Some(ref mut st) => {
+                    *st = get_supertype(&st, dt_i).ok()?;
+                }
+            }
         }
+        let st = st?;
 
         match st {
             DataType::UInt32 => self

diff --git a/py-polars/tests/test_interop.py b/py-polars/tests/test_interop.py
@@ -423,3 +423,9 @@ def test_from_pyarrow_chunked_array() -> None:
     column = pa.chunked_array([[1], [2]])
     series = pl.Series("column", column)
     assert series.to_list() == [1, 2]
+
+
+def test_numpy_preserve_uint64_4112() -> None:
+    assert pl.DataFrame({"a": [1, 2, 3]}).with_column(
+        pl.col("a").hash()
+    ).to_numpy().dtype == np.dtype("uint64")