How to convert column to list in expressions in Polars?

Question

I was earlier able to convert column to list but it is not working now after the latest version update.

import polars as pl

df = pl.DataFrame(
    {
        "A": [1,4,4,7,7,10,10,13,16],
        "B": [2,5,5,8,18,11,11,14,17],
        "C": [3,6,6,9,9,12,12,15,18]        
    }
)

I have also referred to polars_list_link but below code is not working

df.with_columns(
    pl.col('B').sum().over('A').alias("sum[B]/A_groups"),
    pl.col("A").list.head(2).alias("list_A"),
    pl.col("A").list.lengths().alias("length"),
)

Error
SchemaError: invalid series dtype: expected List, got i64

Update
Expected Results

┌─────┬─────┬─────┬─────────────────┐
│ A   ┆ B   ┆ C   ┆ sum[B]/A_groups │list_A length
│ --- ┆ --- ┆ --- ┆ ---             │
│ i64 ┆ i64 ┆ i64 ┆ i64             │
╞═════╪═════╪═════╪═════════════════╡
│ 1   ┆ 2   ┆ 3   ┆ 2               │[1,4]   9
│ 4   ┆ 5   ┆ 6   ┆ 10              │[1,4]   9
│ 4   ┆ 5   ┆ 6   ┆ 10              │[1,4]   9
│ 7   ┆ 8   ┆ 9   ┆ 26              │[1,4]   9
│ 7   ┆ 18  ┆ 9   ┆ 26              │[1,4]   9
│ 10  ┆ 11  ┆ 12  ┆ 22              │[1,4]   9
│ 10  ┆ 11  ┆ 12  ┆ 22              │[1,4]   9
│ 13  ┆ 14  ┆ 15  ┆ 14              │[1,4]   9
│ 16  ┆ 17  ┆ 18  ┆ 17              │[1,4]   9
└─────┴─────┴─────┴─────────────────┘

Later I will use this with window partitioning for these to make more sense.

Asked By: ViSa

||

Source

Answer 1

You can pull the data from the Expr directly using the Expr functions head and len. Because the head returns multiple values, you can implode to convert the results to a list.

import polars as pl

df = pl.DataFrame(
    {
        "A": [1,4,4,7,7,10,10,13,16],
        "B": [2,5,5,8,18,11,11,14,17],
        "C": [3,6,6,9,9,12,12,15,18]        
    }
)

print(df.with_columns(
    pl.col('B').sum().over('A').alias("sum[B]/A_groups"),
    pl.col("A").head(2).implode().alias("list_A"),
    pl.col("A").len().alias("length"),
))

Results:

┌─────┬─────┬─────┬─────────────────┬───────────┬────────┐
│ A   ┆ B   ┆ C   ┆ sum[B]/A_groups ┆ list_A    ┆ length │
│ --- ┆ --- ┆ --- ┆ ---             ┆ ---       ┆ ---    │
│ i64 ┆ i64 ┆ i64 ┆ i64             ┆ list[i64] ┆ u32    │
╞═════╪═════╪═════╪═════════════════╪═══════════╪════════╡
│ 1   ┆ 2   ┆ 3   ┆ 2               ┆ [1, 4]    ┆ 9      │
│ 4   ┆ 5   ┆ 6   ┆ 10              ┆ [1, 4]    ┆ 9      │
│ 4   ┆ 5   ┆ 6   ┆ 10              ┆ [1, 4]    ┆ 9      │
│ 7   ┆ 8   ┆ 9   ┆ 26              ┆ [1, 4]    ┆ 9      │
│ 7   ┆ 18  ┆ 9   ┆ 26              ┆ [1, 4]    ┆ 9      │
│ 10  ┆ 11  ┆ 12  ┆ 22              ┆ [1, 4]    ┆ 9      │
│ 10  ┆ 11  ┆ 12  ┆ 22              ┆ [1, 4]    ┆ 9      │
│ 13  ┆ 14  ┆ 15  ┆ 14              ┆ [1, 4]    ┆ 9      │
│ 16  ┆ 17  ┆ 18  ┆ 17              ┆ [1, 4]    ┆ 9      │
└─────┴─────┴─────┴─────────────────┴───────────┴────────┘

Answered By: keraion

How to convert column to list in expressions in Polars?

Question:

Answers: