Suppose I want to take a DataFrame that contains column with a List of numbers, and returns a new List with all the non-zero elements indices.
┌───────────┐
│ dense │
│ --- │
│ list[i64] │
╞═══════════╡
│ [0, 9] │
│ [8, 6] │
│ [3, 3] │
└───────────┘
┌───────────┬────────────────┐
│ dense ┆ sparse_indices │
│ --- ┆ --- │
│ list[i64] ┆ list[i64] │
╞═══════════╪════════════════╡
│ [0, 9] ┆ [1] │
│ [8, 6] ┆ [0, 1] │
│ [3, 3] ┆ [0, 1] │
└───────────┴────────────────┘
df = df.with_columns(
pl.col('dense').map_elements(lambda l: [i for i, v in enumerate(l) if v != 0]).alias('sparse_indices')
)
I know it can be solved with an index and explode, but it's pretty non-trivial code.
Thanks.