diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py index 666d88e82214..3bd94374359e 100644 --- a/py-polars/polars/dataframe/frame.py +++ b/py-polars/polars/dataframe/frame.py @@ -5251,6 +5251,24 @@ def with_row_index(self, name: str = "index", offset: int = 0) -> Self: │ 1001 ┆ 3 ┆ 4 │ │ 1002 ┆ 5 ┆ 6 │ └──────┴─────┴─────┘ + + An index column can also be created using the expressions :func:`int_range` + and :func:`count`. + + >>> df.select( + ... pl.int_range(pl.count(), dtype=pl.UInt32).alias("index"), + ... pl.all(), + ... ) + shape: (3, 3) + ┌───────┬─────┬─────┐ + │ index ┆ a ┆ b │ + │ --- ┆ --- ┆ --- │ + │ u32 ┆ i64 ┆ i64 │ + ╞═══════╪═════╪═════╡ + │ 0 ┆ 1 ┆ 2 │ + │ 1 ┆ 3 ┆ 4 │ + │ 2 ┆ 5 ┆ 6 │ + └───────┴─────┴─────┘ """ try: return self._from_pydf(self._df.with_row_index(name, offset)) diff --git a/py-polars/polars/functions/lazy.py b/py-polars/polars/functions/lazy.py index 2076e1611558..df8560c98c43 100644 --- a/py-polars/polars/functions/lazy.py +++ b/py-polars/polars/functions/lazy.py @@ -126,6 +126,23 @@ def count(column: str | None = None) -> Expr: ╞═════╡ │ 2 │ └─────┘ + + Generate an index column using `count` in conjunction with :func:`int_range`. + + >>> df.select( + ... pl.int_range(pl.count(), dtype=pl.UInt32).alias("index"), + ... pl.all(), + ... ) + shape: (3, 3) + ┌───────┬──────┬──────┐ + │ index ┆ a ┆ b │ + │ --- ┆ --- ┆ --- │ + │ u32 ┆ i64 ┆ i64 │ + ╞═══════╪══════╪══════╡ + │ 0 ┆ 1 ┆ 3 │ + │ 1 ┆ 2 ┆ null │ + │ 2 ┆ null ┆ null │ + └───────┴──────┴──────┘ """ # noqa: W505 if column is None: return wrap_expr(plr.count()) diff --git a/py-polars/polars/functions/range/int_range.py b/py-polars/polars/functions/range/int_range.py index d58d45a90a30..91518fd5f816 100644 --- a/py-polars/polars/functions/range/int_range.py +++ b/py-polars/polars/functions/range/int_range.py @@ -198,8 +198,7 @@ def int_range( 2 ] - `int_range` can be used in conjunction with `count` to generate an index column - for a DataFrame. + Generate an index column using `int_range` in conjunction with :func:`count`. >>> df = pl.DataFrame({"a": [1, 3, 5], "b": [2, 4, 6]}) >>> df.select( diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py index d3ec947c4b7b..2baf69836f68 100644 --- a/py-polars/polars/lazyframe/frame.py +++ b/py-polars/polars/lazyframe/frame.py @@ -4614,6 +4614,24 @@ def with_row_index(self, name: str = "index", offset: int = 0) -> Self: │ 1001 ┆ 3 ┆ 4 │ │ 1002 ┆ 5 ┆ 6 │ └──────┴─────┴─────┘ + + An index column can also be created using the expressions :func:`int_range` + and :func:`count`. + + >>> lf.select( + ... pl.int_range(pl.count(), dtype=pl.UInt32).alias("index"), + ... pl.all(), + ... ).collect() + shape: (3, 3) + ┌───────┬─────┬─────┐ + │ index ┆ a ┆ b │ + │ --- ┆ --- ┆ --- │ + │ u32 ┆ i64 ┆ i64 │ + ╞═══════╪═════╪═════╡ + │ 0 ┆ 1 ┆ 2 │ + │ 1 ┆ 3 ┆ 4 │ + │ 2 ┆ 5 ┆ 6 │ + └───────┴─────┴─────┘ """ try: return self._from_pyldf(self._ldf.with_row_index(name, offset))