Skip rechunk in ufunc if n_chunks == 1. (#2296)

pola-rs · Jan 7, 2022 · 6f8737d · 6f8737d
1 parent 98ccee4
commit 6f8737d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -1904,7 +1904,7 @@ def __array_ufunc__(
         """
         Numpy universal functions.
         """
-        if self._s.n_chunks() > 0:
+        if self._s.n_chunks() > 1:
             self._s.rechunk(in_place=True)
 
         if method == "__call__":

diff --git a/py-polars/tests/test_series.py b/py-polars/tests/test_series.py
@@ -319,6 +319,14 @@ def test_ufunc() -> None:
     b = np.exp(a)
     assert b.null_count() == 1
 
+    # test if it works with chunked series.
+    a = pl.Series("a", [1.0, None, 3.0])
+    b = pl.Series("b", [4.0, 5.0, None])
+    a.append(b)
+    assert a.n_chunks() == 2
+    c = np.multiply(a, 3)
+    testing.assert_series_equal(c, pl.Series("a", [3.0, None, 9.0, 12.0, 15.0, None]))
+
 
 def test_get() -> None:
     a = pl.Series("a", [1, 2, 3])