diff --git a/py-polars/polars/functions/range/int_range.py b/py-polars/polars/functions/range/int_range.py index 5928bf435450..d58d45a90a30 100644 --- a/py-polars/polars/functions/range/int_range.py +++ b/py-polars/polars/functions/range/int_range.py @@ -20,8 +20,8 @@ @overload def arange( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = ..., + end: int | IntoExprColumn | None = ..., step: int = ..., *, dtype: PolarsIntegerType = ..., @@ -32,8 +32,8 @@ def arange( @overload def arange( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = ..., + end: int | IntoExprColumn | None = ..., step: int = ..., *, dtype: PolarsIntegerType = ..., @@ -44,8 +44,8 @@ def arange( @overload def arange( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = ..., + end: int | IntoExprColumn | None = ..., step: int = ..., *, dtype: PolarsIntegerType = ..., @@ -55,8 +55,8 @@ def arange( def arange( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = 1, + end: int | IntoExprColumn | None = None, step: int = 1, *, dtype: PolarsIntegerType = Int64, @@ -107,8 +107,8 @@ def arange( @overload def int_range( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = ..., + end: int | IntoExprColumn | None = ..., step: int = ..., *, dtype: PolarsIntegerType = ..., @@ -119,8 +119,8 @@ def int_range( @overload def int_range( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = ..., + end: int | IntoExprColumn | None = ..., step: int = ..., *, dtype: PolarsIntegerType = ..., @@ -131,8 +131,8 @@ def int_range( @overload def int_range( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = ..., + end: int | IntoExprColumn | None = ..., step: int = ..., *, dtype: PolarsIntegerType = ..., @@ -142,8 +142,8 @@ def int_range( def int_range( - start: int | IntoExprColumn, - end: int | IntoExprColumn, + start: int | IntoExprColumn = 0, + end: int | IntoExprColumn | None = None, step: int = 1, *, dtype: PolarsIntegerType = Int64, @@ -155,9 +155,10 @@ def int_range( Parameters ---------- start - Lower bound of the range (inclusive). + Start of the range (inclusive). Defaults to 0. end - Upper bound of the range (exclusive). + End of the range (exclusive). If set to `None` (default), + the value of `start` is used and `start` is set to `0`. step Step size of the range. dtype @@ -186,12 +187,23 @@ def int_range( 2 ] + `end` can be omitted for a shorter syntax. + + >>> pl.int_range(3, eager=True).alias("int") + shape: (3,) + Series: 'int' [i64] + [ + 0 + 1 + 2 + ] + `int_range` can be used in conjunction with `count` to generate an index column for a DataFrame. >>> df = pl.DataFrame({"a": [1, 3, 5], "b": [2, 4, 6]}) >>> df.select( - ... pl.int_range(0, pl.count(), dtype=pl.UInt32).alias("index"), + ... pl.int_range(pl.count(), dtype=pl.UInt32).alias("index"), ... pl.all(), ... ) shape: (3, 3) @@ -205,6 +217,10 @@ def int_range( │ 2 ┆ 5 ┆ 6 │ └───────┴─────┴─────┘ """ + if end is None: + end = start + start = 0 + start = parse_as_expression(start) end = parse_as_expression(end) result = wrap_expr(plr.int_range(start, end, step, dtype)) @@ -265,9 +281,9 @@ def int_ranges( Parameters ---------- start - Lower bound of the range (inclusive). + Start of the range (inclusive). end - Upper bound of the range (exclusive). + End of the range (exclusive). step Step size of the range. dtype diff --git a/py-polars/tests/unit/functions/range/test_int_range.py b/py-polars/tests/unit/functions/range/test_int_range.py index 9261b97a37c1..7f6a8b5c23f2 100644 --- a/py-polars/tests/unit/functions/range/test_int_range.py +++ b/py-polars/tests/unit/functions/range/test_int_range.py @@ -37,6 +37,18 @@ def test_int_range() -> None: assert_series_equal(pl.select(int_range=result).to_series(), expected) +def test_int_range_short_syntax() -> None: + result = pl.int_range(3) + expected = pl.Series("int", [0, 1, 2]) + assert_series_equal(pl.select(int=result).to_series(), expected) + + +def test_int_range_start_default() -> None: + result = pl.int_range(end=3) + expected = pl.Series("int", [0, 1, 2]) + assert_series_equal(pl.select(int=result).to_series(), expected) + + def test_int_range_eager() -> None: result = pl.int_range(0, 3, eager=True) expected = pl.Series("literal", [0, 1, 2])