python improve various ergonomics; Series.shift(expr), Series.shift_a…

…nd_fill

polars/polars-lazy/src/dsl.rs

@@ -1,7 +1,7 @@
 //! Domain specific language for the Lazy api.
 use crate::logical_plan::Context;
 use crate::prelude::*;
-use crate::utils::{has_expr, has_wildcard};
+use crate::utils::{expr_to_root_column_name, has_expr, has_wildcard};
 use polars_core::prelude::*;
 
 #[cfg(feature = "temporal")]
@@ -1196,17 +1196,19 @@ impl Expr {
     #[cfg(feature = "is_in")]
     #[cfg_attr(docsrs, doc(cfg(feature = "is_in")))]
     pub fn is_in(self, other: Expr) -> Self {
+        let name = expr_to_root_column_name(&self).unwrap();
+        let output_field = Some(Field::new(&name, DataType::Boolean));
         map_binary(
             self,
             other,
-            |left, other| {
+            move |left, other| {
                 left.is_in(&other).map(|ca| {
                     let mut s = ca.into_series();
-                    s.rename(left.name());
+                    s.rename(&name);
                     s
                 })
             },
-            Some(Field::new("", DataType::Boolean)),
+            output_field,
         )
     }
 

py-polars/docs/source/reference/series.rst

@@ -131,6 +131,7 @@ Manipulation/ selection
     Series.set
     Series.clone
     Series.shift
+    Series.shift_and_fill
     Series.drop_nulls
     Series.rechunk
     Series.cast

py-polars/polars/__init__.py

@@ -21,7 +21,10 @@
 from .eager import *
 from .functions import *
 from .io import *
+
+# explicit imports make mypy happy
 from .lazy import *
+from .lazy import col, lit
 from .lazy import to_list as list
 from .string_cache import *
 

py-polars/polars/eager/series.py

@@ -1840,7 +1840,7 @@ def __copy__(self) -> "Series":  # type: ignore
     def __deepcopy__(self, memodict={}) -> "Series":  # type: ignore
         return self.clone()
 
-    def fill_null(self, strategy: str) -> "Series":
+    def fill_null(self, strategy: Union[str, "pl.Expr"]) -> "Series":
         """
         Fill null values with a filling strategy.
 
@@ -1869,6 +1869,8 @@ def fill_null(self, strategy: str) -> "Series":
         Parameters
         ----------
         strategy
+
+        Fill null strategy or a value
                * "backward"
                * "forward"
                * "min"
@@ -1877,6 +1879,10 @@ def fill_null(self, strategy: str) -> "Series":
                * "one"
                * "zero"
         """
+        if not isinstance(strategy, str):
+            return self.to_frame().select(pl.col(self.name).fill_null(strategy))[
+                self.name
+            ]
         return wrap_s(self._s.fill_null(strategy))
 
     def round(self, decimals: int) -> "Series":
@@ -2013,6 +2019,22 @@ def shift(self, periods: int = 1) -> "Series":
         """
         return wrap_s(self._s.shift(periods))
 
+    def shift_and_fill(self, periods: int, fill_value: "pl.Expr") -> "Series":
+        """
+        Shift the values by a given period and fill the parts that will be empty due to this operation
+        with the result of the `fill_value` expression.
+
+        Parameters
+        ----------
+        periods
+            Number of places to shift (may be negative).
+        fill_value
+            Fill None values with the result of this expression.
+        """
+        return self.to_frame().select(
+            pl.col(self.name).shift_and_fill(periods, fill_value)
+        )[self.name]
+
     def zip_with(self, mask: "Series", other: "Series") -> "Series":
         """
         Where mask evaluates true, take values from self. Where mask evaluates false, take values from other.

py-polars/polars/lazy/expr.py

@@ -1106,7 +1106,10 @@ def is_in(self, other: Union["Expr", tp.List[Any]]) -> "Expr":
         -------
         Expr that evaluates to a Boolean Series.
         """
-        other = expr_to_lit_or_expr(other, str_to_lit=False)
+        if isinstance(other, list):
+            other = pl.lit(pl.Series(other))
+        else:
+            other = expr_to_lit_or_expr(other, str_to_lit=False)
         return wrap_expr(self._pyexpr.is_in(other._pyexpr))
 
     def repeat_by(self, by: "Expr") -> "Expr":

py-polars/tests/test_lazy.py

@@ -562,3 +562,12 @@ def test_expr_bool_cmp():
 
     with pytest.raises(ValueError):
         df[[pl.col("a").gt(pl.col("b")) or pl.col("b").gt(pl.col("b"))]]
+
+
+def test_is_in():
+    df = pl.DataFrame({"a": [1, 2, 3]})
+    assert df.select(pl.col("a").is_in([1, 2]))["a"].to_list() == [
+        True,
+        True,
+        False,
+    ]

py-polars/tests/test_series.py

@@ -242,6 +242,8 @@ def test_fill_null():
     a = pl.Series("a", [1, 2, None], nullable=True)
     b = a.fill_null("forward")
     assert b == [1, 2, 2]
+    b = a.fill_null(14)
+    assert b.to_list() == [1, 2, 14]
 
 
 def test_apply():
@@ -253,7 +255,7 @@ def test_apply():
     b = a.apply(lambda x: x + "py")
     assert b == ["foopy", "barpy", None]
 
-    b = a.apply(lambda x: len(x), return_dtype=Int32)
+    b = a.apply(lambda x: len(x), return_dtype=pl.Int32)
     assert b == [3, 3, None]
 
     b = a.apply(lambda x: len(x))
@@ -268,22 +270,23 @@ def test_apply():
 
 def test_shift():
     a = pl.Series("a", [1, 2, 3])
-    assert a.shift(1) == [None, 1, 2]
-    assert a.shift(-1) == [1, 2, None]
-    assert a.shift(-2) == [1, None, None]
+    assert a.shift(1).to_list() == [None, 1, 2]
+    assert a.shift(-1).to_list() == [2, 3, None]
+    assert a.shift(-2).to_list() == [3, None, None]
+    assert a.shift_and_fill(-1, 10).to_list() == [2, 3, 10]
 
 
 def test_rolling():
     a = pl.Series("a", [1, 2, 3, 2, 1])
-    assert a.rolling_min(2) == [None, 1, 2, 2, 1]
-    assert a.rolling_max(2) == [None, 2, 3, 3, 2]
-    assert a.rolling_sum(2) == [None, 3, 5, 5, 3]
+    assert a.rolling_min(2).to_list() == [None, 1, 2, 2, 1]
+    assert a.rolling_max(2).to_list() == [None, 2, 3, 3, 2]
+    assert a.rolling_sum(2).to_list() == [None, 3, 5, 5, 3]
 
 
 def test_object():
     vals = [[12], "foo", 9]
     a = pl.Series("a", vals)
-    assert a.dtype == Object
+    assert a.dtype == pl.Object
     assert a.to_list() == vals
     assert a[1] == "foo"