From dd8ebe3ccf18c794fb085cb431c0949c2c928dc4 Mon Sep 17 00:00:00 2001 From: Jeroen van Zundert Date: Fri, 17 Feb 2023 23:33:20 +0000 Subject: [PATCH] feat(python): Improve Series & Numpy arithmetic --- py-polars/polars/internals/series/series.py | 5 ++- py-polars/tests/unit/test_series.py | 38 +++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/py-polars/polars/internals/series/series.py b/py-polars/polars/internals/series/series.py index 67e19caa61c4..90a0d63d237c 100644 --- a/py-polars/polars/internals/series/series.py +++ b/py-polars/polars/internals/series/series.py @@ -527,7 +527,8 @@ def _arithmetic(self, other: Any, op_s: str, op_ffi: str) -> Series: other = self.to_frame().select(other).to_series() if isinstance(other, Series): return wrap_s(getattr(self._s, op_s)(other._s)) - + if _check_for_numpy(other) and isinstance(other, np.ndarray): + return wrap_s(getattr(self._s, op_s)(Series(other)._s)) # recurse; the 'if' statement above will ensure we return early if isinstance(other, (date, datetime, timedelta, str)): other = Series("", [other]) @@ -653,6 +654,8 @@ def __pow__(self, power: int | float | Series) -> Series: raise ValueError( "first cast to integer before raising datelike dtypes to a power" ) + if _check_for_numpy(power) and isinstance(power, np.ndarray): + power = Series(power) return self.to_frame().select(pli.col(self.name).pow(power)).to_series() def __rpow__(self, other: Any) -> Series: diff --git a/py-polars/tests/unit/test_series.py b/py-polars/tests/unit/test_series.py index 76155de4c177..d9af31a62260 100644 --- a/py-polars/tests/unit/test_series.py +++ b/py-polars/tests/unit/test_series.py @@ -2443,3 +2443,41 @@ def test_upper_lower_bounds( s = pl.Series("s", dtype=dtype) assert s.lower_bound().item() == lower assert s.upper_bound().item() == upper + + +def test_numpy_series_arithmetic() -> None: + sx = pl.Series(values=[1, 2]) + y = np.array([3.0, 4.0]) + + result_add1 = y + sx + result_add2 = sx + y + expected_add = pl.Series([4.0, 6.0], dtype=pl.Float64) + assert_series_equal(result_add1, expected_add) # type: ignore[arg-type] + assert_series_equal(result_add2, expected_add) + + result_sub1 = cast(pl.Series, y - sx) # py37 is different vs py311 on this one + expected = pl.Series([2.0, 2.0], dtype=pl.Float64) + assert_series_equal(result_sub1, expected) + result_sub2 = sx - y + expected = pl.Series([-2.0, -2.0], dtype=pl.Float64) + assert_series_equal(result_sub2, expected) + + result_mul1 = y * sx + result_mul2 = sx * y + expected = pl.Series([3.0, 8.0], dtype=pl.Float64) + assert_series_equal(result_mul1, expected) # type: ignore[arg-type] + assert_series_equal(result_mul2, expected) + + result_div1 = y / sx + expected = pl.Series([3.0, 2.0], dtype=pl.Float64) + assert_series_equal(result_div1, expected) # type: ignore[arg-type] + result_div2 = sx / y + expected = pl.Series([1 / 3, 0.5], dtype=pl.Float64) + assert_series_equal(result_div2, expected) + + result_pow1 = y**sx + expected = pl.Series([3.0, 16.0], dtype=pl.Float64) + assert_series_equal(result_pow1, expected) # type: ignore[arg-type] + result_pow2 = sx**y + expected = pl.Series([1.0, 16.0], dtype=pl.Float64) + assert_series_equal(result_pow2, expected) # type: ignore[arg-type]