Polars: Pad list columns to specific size

I think i ran into the XY problem...

Here is what i actually want to do:

To be exact i have a dataframe like:

shape: (3, 3)
┌───────────┬───────┬──────────────────────────┐
│ nrs       ┆ stuff ┆ more_stuff               │
│ ---       ┆ ---   ┆ ---                      │
│ list[i64] ┆ i64   ┆ list[list[i64]]          │
╞═══════════╪═══════╪══════════════════════════╡
│ [1, 2, 3] ┆ 1     ┆ [[1, 1], [2, 2], [3, 3]] │
│ [2, 4]    ┆ 2     ┆ [[4, 4], [5, 5]]         │
│ [1]       ┆ 3     ┆ [[6, 6]]                 │
└───────────┴───────┴──────────────────────────┘

With normal int64 columns, list[int64] columns and one list[list[64]] column.
I want to be able to specify a size and set the length of all the list (also the nested) columns to that size. Either by shortening to that size or by padding them with their last value (list[-1] for normal python lists) for both the nested and the normal lists. The non-list columns should be left unchanged.

So the result for N=2 for the above dataframe should be:

shape: (3, 3)
┌───────────┬───────┬──────────────────┐
│ nrs       ┆ stuff ┆ more_stuff       │
│ ---       ┆ ---   ┆ ---              │
│ list[i64] ┆ i64   ┆ list[list[i64]]  │
╞═══════════╪═══════╪══════════════════╡
│ [1, 2]    ┆ 1     ┆ [[1, 1], [2, 2]] │
│ [2, 4]    ┆ 2     ┆ [[4, 4], [5, 5]] │
│ [1, 1]    ┆ 3     ┆ [[6, 6], [6, 6]] │
└───────────┴───────┴──────────────────┘

Answer via reddit from /u/commandlineluser

df = pl.DataFrame({
   "nrs": [[1, 2, 3], [2, 4], [1]], 
   "stuff": [1, 2, 3], 
   "more_stuff": [[[1, 1], [2, 2], [3, 3]], [[4, 4], [5, 5]], [[6, 6]]]
})

cols = "nrs", "more_stuff"

df.with_columns(
   pl.col(cols).arr.take(
      pl.arange(0, pl.col(cols).arr.lengths().max()), 
      null_on_oob=True
   ).arr.eval(pl.element().forward_fill())
)

# shape: (3, 3)
# ┌───────────┬───────┬──────────────────────────┐
# │ nrs       ┆ stuff ┆ more_stuff               │
# │ ---       ┆ ---   ┆ ---                      │
# │ list[i64] ┆ i64   ┆ list[list[i64]]          │
# ╞═══════════╪═══════╪══════════════════════════╡
# │ [1, 2, 3] ┆ 1     ┆ [[1, 1], [2, 2], [3, 3]] │
# │ [2, 4, 4] ┆ 2     ┆ [[4, 4], [5, 5], [5, 5]] │
# │ [1, 1, 1] ┆ 3     ┆ [[6, 6], [6, 6], [6, 6]] │
# └───────────┴───────┴──────────────────────────┘

You can add a .slice()/.head() to pad them to a smaller length.

EDIT:

Possibly worth noting that .forward_fill() doesn't specifically target the padded nulls, so that could be an issue if there were nulls in the initial data. It's possible to handle this, but requires a bit more code.

See extend_constant.

You can shorten with head / tail