feat(python): pl.ones, pl.zeros and Series.new_from_index functions (#…

…5260)

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

@@ -34,6 +34,8 @@ Eager/Lazy functions
    date_range
    get_dummies
    repeat
+   ones
+   zeros
 
 Parallelization
 ~~~~~~~~~~~~~~~

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

@@ -182,6 +182,7 @@ Manipulation/ selection
     Series.head
     Series.interpolate
     Series.limit
+    Series.new_from_index
     Series.rechunk
     Series.rename
     Series.reshape

py-polars/polars/__init__.py

@@ -76,6 +76,8 @@ def version() -> str:
     cut,
     date_range,
     get_dummies,
+    ones,
+    zeros,
 )
 from polars.internals.io import read_ipc_schema, read_parquet_schema
 from polars.internals.lazy_functions import _date as date
@@ -228,9 +230,11 @@ def version() -> str:
     "concat",
     "date_range",
     "get_dummies",
+    "ones",
     "repeat",
     "element",
     "cut",
+    "zeros",
     # polars.internals.lazy_functions
     "col",
     "count",

py-polars/polars/internals/functions.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Sequence, overload
 
 from polars import internals as pli
-from polars.datatypes import Categorical, Date, Float64
+from polars.datatypes import Categorical, Date, Float64, PolarsDataType
 from polars.utils import _datetime_to_pl_timestamp, _timedelta_to_pl_duration
 
 try:
@@ -573,3 +573,49 @@ def align_frames(
         aligned_frames = [df.select(select) for df in aligned_frames]
 
     return [df.collect() for df in aligned_frames] if eager else aligned_frames
+
+
+def ones(n: int, dtype: PolarsDataType | None = None) -> pli.Series:
+    """
+    Return a new Series of given length and type, filled with ones.
+
+    Parameters
+    ----------
+    n
+        Number of elements in the ``Series``
+    dtype
+        DataType of the elements, defaults to ``polars.Float64``
+
+    Notes
+    -----
+    In the lazy API you should probably not use this, but use ``lit(1)``
+    instead.
+
+    """
+    s = pli.Series([1.0])
+    if dtype:
+        s = s.cast(dtype)
+    return s.new_from_index(0, n)
+
+
+def zeros(n: int, dtype: PolarsDataType | None = None) -> pli.Series:
+    """
+    Return a new Series of given length and type, filled with zeros.
+
+    Parameters
+    ----------
+    n
+        Number of elements in the ``Series``
+    dtype
+        DataType of the elements, defaults to ``polars.Float64``
+
+    Notes
+    -----
+    In the lazy API you should probably not use this, but use ``lit(0)``
+    instead.
+
+    """
+    s = pli.Series([0.0])
+    if dtype:
+        s = s.cast(dtype)
+    return s.new_from_index(0, n)

py-polars/polars/internals/series/series.py

@@ -4652,6 +4652,10 @@ def set_sorted(self, reverse: bool = False) -> Series:
         """
         return wrap_s(self._s.set_sorted(reverse))
 
+    def new_from_index(self, index: int, length: int) -> pli.Series:
+        """Create a new Series filled with values from the given index."""
+        return wrap_s(self._s.new_from_index(index, length))
+
     # 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/series.rs

@@ -2,7 +2,7 @@ use numpy::PyArray1;
 use polars_core::prelude::QuantileInterpolOptions;
 use polars_core::series::IsSorted;
 use polars_core::utils::CustomIterTools;
-use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::{PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyList, PyTuple};
 use pyo3::Python;
@@ -494,6 +494,14 @@ impl PySeries {
         Ok(())
     }
 
+    pub fn new_from_index(&self, index: usize, length: usize) -> PyResult<Self> {
+        if index >= self.series.len() {
+            Err(PyValueError::new_err("index is out of bounds"))
+        } else {
+            Ok(self.series.new_from_index(index, length).into())
+        }
+    }
+
     pub fn filter(&self, filter: &PySeries) -> PyResult<Self> {
         let filter_series = &filter.series;
         if let Ok(ca) = filter_series.bool() {

py-polars/tests/unit/test_functions.py

@@ -248,3 +248,21 @@ def test_coalesce() -> None:
     assert df.select(pl.coalesce(["a", "b", "c", 10])).to_dict(False) == {
         "a": [1.0, 2.0, 3.0, 10.0]
     }
+
+
+def test_ones_zeros() -> None:
+    ones = pl.ones(5)
+    assert ones.dtype == pl.Float64
+    assert ones.to_list() == [1.0, 1.0, 1.0, 1.0, 1.0]
+
+    ones = pl.ones(3, dtype=pl.UInt8)
+    assert ones.dtype == pl.UInt8
+    assert ones.to_list() == [1, 1, 1]
+
+    zeros = pl.zeros(5)
+    assert zeros.dtype == pl.Float64
+    assert zeros.to_list() == [0.0, 0.0, 0.0, 0.0, 0.0]
+
+    zeros = pl.zeros(3, dtype=pl.UInt8)
+    assert zeros.dtype == pl.UInt8
+    assert zeros.to_list() == [0, 0, 0]