Skip to content

Commit

Permalink
docs(python): Add more doc examples on how to create an index column (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
stinodego authored Jan 8, 2024
1 parent 3478ba8 commit 67507ff
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 2 deletions.
18 changes: 18 additions & 0 deletions py-polars/polars/dataframe/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
17 changes: 17 additions & 0 deletions py-polars/polars/functions/lazy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
3 changes: 1 addition & 2 deletions py-polars/polars/functions/range/int_range.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
18 changes: 18 additions & 0 deletions py-polars/polars/lazyframe/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down

0 comments on commit 67507ff

Please sign in to comment.