dataframe - (Polars)如何通过在另一列中指定的索引从列表中获取元素

Question

我有一个包含 2 列的数据框，其中第一列包含列表，第二列包含整数索引。如何通过第二列中指定的索引从第一列获取元素？或者更好的是，将该元素放在第 3 列中。因此，例如，如何从这个

a = pl.DataFrame([{'lst': [1, 2, 3], 'ind': 1}, {'lst': [4, 5, 6], 'ind': 2}])
┌───────────┬─────┐
│ lst       ┆ ind │
│ ---       ┆ --- │
│ list[i64] ┆ i64 │
╞═══════════╪═════╡
│ [1, 2, 3] ┆ 1   │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
│ [4, 5, 6] ┆ 2   │
└───────────┴─────┘

你可以得到这个

b = pl.DataFrame([{'lst': [1, 2, 3], 'ind': 1, 'list[ind]': 2}, {'lst': [4, 5, 6], 'ind': 2, 'list[ind]': 6}])
┌───────────┬─────┬───────────┐
│ lst       ┆ ind ┆ list[ind] │
│ ---       ┆ --- ┆ ---       │
│ list[i64] ┆ i64 ┆ i64       │
╞═══════════╪═════╪═══════════╡
│ [1, 2, 3] ┆ 1   ┆ 2         │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ [4, 5, 6] ┆ 2   ┆ 6         │
└───────────┴─────┴───────────┘

谢谢。

Answer 1

编辑

从 python polars 0.14.24 开始，这可以通过以下方式更轻松地完成

df.with_column(pl.col("lst").arr.get(pl.col("ind")).alias("list[ind]"))

原始答案

您可以使用 with_row_count() 添加行计数列以进行分组，然后使用 explode() 列表，这样每个列表元素都在每一行上。然后使用 over() 在行计数列上调用 take() 以从子组中选择元素。

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

df = (
    df.with_row_count()
    .with_column(
        pl.col("lst").explode().take(pl.col("ind")).over(pl.col("row_nr")).alias("list[ind]")
    )
    .drop("row_nr")
)

shape: (2, 3)
┌───────────┬─────┬───────────┐
│ lst       ┆ ind ┆ list[ind] │
│ ---       ┆ --- ┆ ---       │
│ list[i64] ┆ i64 ┆ i64       │
╞═══════════╪═════╪═══════════╡
│ [1, 2, 3] ┆ 1   ┆ 2         │
├╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ [4, 5, 6] ┆ 2   ┆ 6         │
└───────────┴─────┴───────────┘