docs(python): Add more doc examples on how to create an index column (p…

…ola-rs#13532)
nameexhaustion · Jan 8, 2024 · 67507ff · 67507ff
1 parent 3478ba8
commit 67507ff
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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))