feat(python,rust): allow extend_constant to work with date literals (…

…#6114)
pola-rs · Jan 8, 2023 · 9565988 · 9565988
1 parent c8d45ae
commit 9565988
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 13 deletions.
diff --git a/polars/polars-core/src/series/ops/extend.rs b/polars/polars-core/src/series/ops/extend.rs
@@ -13,6 +13,7 @@ impl Series {
             Int64(v) => Series::new("", vec![v]),
             Utf8(v) => Series::new("", vec![v]),
             Boolean(v) => Series::new("", vec![v]),
+            Date(v) => Series::new("", vec![v]),
             Null => BooleanChunked::full_null("", 1).into_series(),
             dt => panic!("{dt:?} not supported"),
         };

diff --git a/py-polars/polars/internals/expr/expr.py b/py-polars/polars/internals/expr/expr.py
@@ -5649,17 +5649,19 @@ def ewm_var(
         alpha = _prepare_alpha(com, span, half_life, alpha)
         return wrap_expr(self._pyexpr.ewm_var(alpha, adjust, bias, min_periods))
 
-    def extend_constant(self, value: int | float | str | bool | None, n: int) -> Expr:
+    def extend_constant(
+        self, value: int | float | str | bool | date | None, n: int
+    ) -> Expr:
         """
         Extend the Series with given number of values.
 
         Parameters
         ----------
         value
-            The value to extend the Series with. This value may be None to fill with
-            nulls.
+            A constant literal value (not an expression) with which to extend the
+            Series; can pass None to fill the Series with nulls.
         n
-            The number of values to extend.
+            The number of additional values that will be added into the Series.
 
         Examples
         --------
@@ -5679,6 +5681,8 @@ def extend_constant(self, value: int | float | str | bool | None, n: int) -> Exp
         └────────┘
 
         """
+        if isinstance(value, Expr):
+            raise TypeError(f"'value' must be a supported literal; found {value!r}")
         return wrap_expr(self._pyexpr.extend_constant(value, n))
 
     def value_counts(self, multithreaded: bool = False, sort: bool = False) -> Expr:

diff --git a/py-polars/polars/internals/series/series.py b/py-polars/polars/internals/series/series.py
@@ -4804,17 +4804,19 @@ def ewm_var(
 
         """
 
-    def extend_constant(self, value: int | float | str | bool | None, n: int) -> Series:
+    def extend_constant(
+        self, value: int | float | str | bool | date | None, n: int
+    ) -> Series:
         """
         Extend the Series with given number of values.
 
         Parameters
         ----------
         value
-            The value to extend the Series with. This value may be None to fill with
-            nulls.
+            A constant literal value (not an expression) with which to extend the
+            Series; can pass None to fill the Series with nulls.
         n
-            The number of values to extend.
+            The number of additional values that will be added into the Series.
 
         Examples
         --------

diff --git a/py-polars/tests/unit/test_series.py b/py-polars/tests/unit/test_series.py
@@ -2240,12 +2240,17 @@ def test_ewm_param_validation() -> None:
 
 
 def test_extend_constant() -> None:
-    a = pl.Series("a", [1, 2, 3])
-    expected = pl.Series("a", [1, 2, 3, 1, 1, 1])
-    verify_series_and_expr_api(a, expected, "extend_constant", 1, 3)
+    today = date.today()
 
-    expected = pl.Series("a", [1, 2, 3, None, None, None])
-    verify_series_and_expr_api(a, expected, "extend_constant", None, 3)
+    for const, dtype in (
+        (1, pl.Int8),
+        (today, pl.Date),
+        ("xyz", pl.Utf8),
+        (None, pl.Float64),
+    ):
+        s = pl.Series("s", [None], dtype=dtype)
+        expected = pl.Series("s", [None, const, const, const], dtype=dtype)
+        verify_series_and_expr_api(s, expected, "extend_constant", const, 3)
 
 
 def test_any_all() -> None: