Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

In Polars, how do you generate a column of lists, where each list is a range() defined by another column of type Int?

Given a sample DataFrame df

>>> df = pl.DataFrame({'l': [3,5,8]})
>>> df
shape: (3, 1)
┌─────┐
│ l   │
│ --- │
│ i64 │
╞═════╡
│ 3   │
│ 5   │
│ 8   │
└─────┘

How do I make a new column that looks like this?

>>> df
shape: (3, 2)
┌─────┬─────────────┐
│ l   ┆ column_0    │
│ --- ┆ ---         │
│ i64 ┆ list[i64]   │
╞═════╪═════════════╡
│ 3   ┆ [0, 1, 2]   │
│ 5   ┆ [0, 1, … 4] │
│ 8   ┆ [0, 1, … 7] │
└─────┴─────────────┘

This is my best way of doing it but it is using map_rows, not so efficient when things scale up

>>> temp = df.map_rows(lambda x: (list(range(x[0])),))
>>> df = df.hstack(temp)
like image 806
Horace Avatar asked Oct 21 '25 11:10

Horace


1 Answers

For this, pl.int_ranges can be used.

import polars as pl

df = pl.DataFrame({
    "length": [3,5,8]
})

df.with_columns(
    pl.int_ranges("length").alias("list")
)
shape: (3, 2)
┌────────┬─────────────┐
│ length ┆ list        │
│ ---    ┆ ---         │
│ i64    ┆ list[i64]   │
╞════════╪═════════════╡
│ 3      ┆ [0, 1, 2]   │
│ 5      ┆ [0, 1, … 4] │
│ 8      ┆ [0, 1, … 7] │
└────────┴─────────────┘
like image 113
Hericks Avatar answered Oct 23 '25 01:10

Hericks



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!