Skip to content

Commit

Permalink
Update docs/tests
Browse files Browse the repository at this point in the history
  • Loading branch information
stinodego committed Mar 23, 2024
1 parent 1d6e8e6 commit 695611c
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 31 deletions.
28 changes: 14 additions & 14 deletions py-polars/polars/expr/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -3854,8 +3854,8 @@ def rle(self) -> Self:
Returns
-------
Expr
Expression of data type `Struct` with fields `lengths` of data type `Int32`
and `values` of the original data type.
Expression of data type `Struct` with fields `len` of data type `UInt32`
and `value` of the original data type.
See Also
--------
Expand All @@ -3866,18 +3866,18 @@ def rle(self) -> Self:
>>> df = pl.DataFrame({"a": [1, 1, 2, 1, None, 1, 3, 3]})
>>> df.select(pl.col("a").rle()).unnest("a")
shape: (6, 2)
┌─────────┬────────┐
lengthsvalues
│ --- ┆ ---
i32 ┆ i64
╞═════════╪════════╡
│ 2 ┆ 1
│ 1 ┆ 2
│ 1 ┆ 1
│ 1 ┆ null
│ 1 ┆ 1
│ 2 ┆ 3
└─────────┴────────┘
┌────────────┐
lenvalue
│ --- ┆ --- │
u32 ┆ i64 │
╞════════════╡
│ 2 ┆ 1 │
│ 1 ┆ 2 │
│ 1 ┆ 1 │
│ 1 ┆ null │
│ 1 ┆ 1 │
│ 2 ┆ 3 │
└────────────┘
"""
return self._from_pyexpr(self._pyexpr.rle())

Expand Down
28 changes: 14 additions & 14 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -2565,26 +2565,26 @@ def rle(self) -> Series:
Returns
-------
Series
Series of data type `Struct` with fields `lengths` of data type `Int32`
and `values` of the original data type.
Series of data type `Struct` with fields `len` of data type `UInt32`
and `value` of the original data type.
Examples
--------
>>> s = pl.Series("s", [1, 1, 2, 1, None, 1, 3, 3])
>>> s.rle().struct.unnest()
shape: (6, 2)
┌─────────┬────────┐
lengthsvalues
│ --- ┆ ---
i32 ┆ i64
╞═════════╪════════╡
│ 2 ┆ 1
│ 1 ┆ 2
│ 1 ┆ 1
│ 1 ┆ null
│ 1 ┆ 1
│ 2 ┆ 3
└─────────┴────────┘
┌────────────┐
lenvalue
│ --- ┆ --- │
u32 ┆ i64 │
╞════════════╡
│ 2 ┆ 1 │
│ 1 ┆ 2 │
│ 1 ┆ 1 │
│ 1 ┆ null │
│ 1 ┆ 1 │
│ 2 ┆ 3 │
└────────────┘
"""

def rle_id(self) -> Series:
Expand Down
8 changes: 5 additions & 3 deletions py-polars/tests/unit/operations/test_rle.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ def test_rle() -> None:
lf = pl.LazyFrame({"a": values})

expected = pl.LazyFrame(
{"lengths": [2, 1, 1, 1, 1, 2], "values": [1, 2, 1, None, 1, 3]},
schema_overrides={"lengths": pl.Int32},
{"len": [2, 1, 1, 1, 1, 2], "value": [1, 2, 1, None, 1, 3]},
schema_overrides={"len": pl.get_index_type()},
)

result_expr = lf.select(pl.col("a").rle()).unnest("a")
Expand All @@ -22,7 +22,9 @@ def test_rle_id() -> None:
values = [1, 1, 2, 1, None, 1, 3, 3]
lf = pl.LazyFrame({"a": values})

expected = pl.LazyFrame({"a": [0, 0, 1, 2, 3, 4, 5, 5]}, schema={"a": pl.UInt32})
expected = pl.LazyFrame(
{"a": [0, 0, 1, 2, 3, 4, 5, 5]}, schema={"a": pl.get_index_type()}
)

result_expr = lf.select(pl.col("a").rle_id())
assert_frame_equal(result_expr, expected)
Expand Down

0 comments on commit 695611c

Please sign in to comment.