add Series::extend (#2172)

polars/polars-core/src/series/ops/extend.rs

@@ -0,0 +1,26 @@
+use crate::prelude::*;
+
+impl Series {
+    /// Extend with a constant value.
+    pub fn extend(&self, value: AnyValue, n: usize) -> Result<Self> {
+        use AnyValue::*;
+        let s = match value {
+            Float32(v) => Series::new("", vec![v]),
+            Float64(v) => Series::new("", vec![v]),
+            UInt32(v) => Series::new("", vec![v]),
+            UInt64(v) => Series::new("", vec![v]),
+            Int32(v) => Series::new("", vec![v]),
+            Int64(v) => Series::new("", vec![v]),
+            Utf8(v) => Series::new("", vec![v]),
+            Boolean(v) => Series::new("", vec![v]),
+            Null => BooleanChunked::full_null("", 1).into_series(),
+            dt => panic!("{:?} not supported", dt),
+        };
+        let s = s.cast(self.dtype())?;
+        let to_append = s.expand_at_index(0, n);
+
+        let mut out = self.clone();
+        out.append(&to_append)?;
+        Ok(out)
+    }
+}

polars/polars-core/src/series/ops/mod.rs

@@ -2,6 +2,7 @@
 pub mod diff;
 #[cfg(feature = "ewma")]
 mod ewm;
+mod extend;
 #[cfg(feature = "moment")]
 pub mod moment;
 mod null;

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

@@ -194,6 +194,7 @@ Manipulation/ selection
     Expr.reshape
     Expr.to_physical
     Expr.shuffle
+    Expr.extend
 
 Column names
 ------------

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

@@ -167,6 +167,7 @@ Manipulation/ selection
     Series.reshape
     Series.to_dummies
     Series.shuffle
+    Series.extend
 
 Various
 --------

py-polars/polars/internals/expr.py

@@ -2190,6 +2190,19 @@ def ewm_var(
         _prepare_alpha(com, span, half_life, alpha)
         return wrap_expr(self._pyexpr.ewm_var(alpha, adjust, min_periods))
 
+    def extend(self, value: Optional[Union[int, float, str, bool]], 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.
+        n
+            The number of values to extend.
+        """
+        return wrap_expr(self._pyexpr.extend(value, n))
+
     # Below are the namespaces defined. Keep these at the end of the definition of Expr, as to not confuse mypy with
     # the type annotation `str` with the namespace "str"
 

py-polars/polars/internals/series.py

@@ -3217,6 +3217,19 @@ def ewm_var(
             .to_series()
         )
 
+    def extend(self, value: Optional[Union[int, float, str, bool]], 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.
+        n
+            The number of values to extend.
+        """
+        return wrap_s(self._s.extend(value, n))
+
     # Below are the namespaces defined. Do not move these up in the definition of Series, as it confuses mypy between the
     # type annotation `str` and the namespace "str
 

py-polars/src/lazy/dsl.rs

@@ -5,6 +5,7 @@ use crate::lazy::utils::py_exprs_to_exprs;
 use crate::prelude::{parse_strategy, str_to_rankmethod};
 use crate::series::PySeries;
 use crate::utils::{reinterpret, str_to_polarstype};
+use crate::PyPolarsEr::Any;
 use polars::lazy::dsl;
 use polars::lazy::dsl::Operator;
 use polars::prelude::*;
@@ -1052,6 +1053,21 @@ impl PyExpr {
         };
         self.inner.clone().ewm_var(options).into()
     }
+    pub fn extend(&self, py: Python, value: Wrap<AnyValue>, n: usize) -> Self {
+        let value = value.into_py(py);
+        self.inner
+            .clone()
+            .apply(
+                move |s| {
+                    let gil = Python::acquire_gil();
+                    let py = gil.python();
+                    let value = value.extract::<Wrap<AnyValue>>(py).unwrap().0;
+                    s.extend(value, n)
+                },
+                GetOutput::same_type(),
+            )
+            .into()
+    }
 }
 
 impl From<dsl::Expr> for PyExpr {

py-polars/src/series.rs

@@ -1380,6 +1380,11 @@ impl PySeries {
     pub fn shuffle(&self, seed: u64) -> Self {
         self.series.shuffle(seed).into()
     }
+    pub fn extend(&self, value: Wrap<AnyValue>, n: usize) -> PyResult<Self> {
+        let value = value.0;
+        let out = self.series.extend(value, n).map_err(PyPolarsEr::from)?;
+        Ok(out.into())
+    }
 }
 
 macro_rules! impl_ufuncs {

py-polars/tests/test_series.py

@@ -1349,3 +1349,12 @@ def test_ewm_std_var() -> None:
     a = pl.Series("a", [2, 5, 3])
 
     assert (a.ewm_std(alpha=0.5) ** 2).to_list() == a.ewm_var(alpha=0.5).to_list()
+
+
+def test_extend() -> 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", 1, 3)
+
+    expected = pl.Series("a", [1, 2, 3, None, None, None])
+    verify_series_and_expr_api(a, expected, "extend", None, 3)