如何使用极坐标表达式对外部列表进行索引和切片？

Question

我有一个包含三列的极坐标数据框：

audio_idx

、

i_start

和

i_stop

。我还有一个 numpy 数组的外部列表。

audio_idx

指的是外部列表的索引，而

i_start

和

i_stop

指的是numpy数组的切片。

这是一个例子：

audios = [(np.random.rand(10)*10).astype(int) for _ in range(3)]
df = pl.DataFrame({"audio_idx": [0,0,1,2], "i_start": [1,6,6,3], "i_stop": [4,8,9,5]})

音频：

[array([9, 8, 6, 2, 2, 9, 5, 4, 7, 9]),
  array([4, 2, 8, 4, 6, 1, 0, 4, 0, 7]),
  array([3, 2, 6, 9, 6, 8, 8, 8, 0, 1])]

df

 shape: (4, 3)
 ┌───────────┬─────────┬────────┐
 │ audio_idx ┆ i_start ┆ i_stop │
 │ ---       ┆ ---     ┆ ---    │
 │ i64       ┆ i64     ┆ i64    │
 ╞═══════════╪═════════╪════════╡
 │ 0         ┆ 1       ┆ 4      │
 │ 0         ┆ 6       ┆ 8      │
 │ 1         ┆ 6       ┆ 9      │
 │ 2         ┆ 3       ┆ 5      │
 └───────────┴─────────┴────────┘

我想要以下输出：

shape: (4, 4)
┌───────────┬─────────┬────────┬───────────┐
│ audio_idx ┆ i_start ┆ i_stop ┆ audio     │
│ ---       ┆ ---     ┆ ---    ┆ ---       │
│ i64       ┆ i64     ┆ i64    ┆ list[i64] │
╞═══════════╪═════════╪════════╪═══════════╡
│ 0         ┆ 1       ┆ 4      ┆ [8, 6, 2] │
│ 0         ┆ 6       ┆ 8      ┆ [5, 4]    │
│ 1         ┆ 6       ┆ 9      ┆ [0, 4, 0] │
│ 2         ┆ 3       ┆ 5      ┆ [9, 6]    │
└───────────┴─────────┴────────┴───────────┘

我可以通过以下方式使用 map_elements 来实现此目的：

df = df.with_columns(
    audio=pl.struct(["audio_idx", "i_start", "i_stop"]).map_elements(
        lambda x: audios[x["audio_idx"]][x["i_start"] : x["i_stop"]].tolist(),
        return_dtype=pl.List(pl.Int64),
    )
)

但是，有没有办法使用极坐标表达式代替map_elements？我更愿意仍然使用外部列表，而不是将完整的 numpy 数组复制到每一行中，因为这可以节省内存。谢谢你。

Answer 1

如果您从 numpy 数组创建另一个

DataFrame

：

adf = pl.DataFrame({"audio_idx": range(len(audios)), "audio": audios})

┌───────────┬─────────────┐
│ audio_idx ┆ audio       │
│ ---       ┆ ---         │
│ i64       ┆ list[i32]   │
╞═══════════╪═════════════╡
│ 0         ┆ [9, 8, … 9] │
│ 1         ┆ [4, 2, … 7] │
│ 2         ┆ [3, 2, … 1] │
└───────────┴─────────────┘

然后您可以使用

DataFrame.join

和

Expr.list.slice()

来获得结果：

(
    df
    .join(adf, on="audio_idx")
    .with_columns(
        pl.col("audio").list.slice(pl.col("i_start"), pl.col("i_stop") - pl.col("i_start"))
    )
)

┌───────────┬─────────┬────────┬───────────┐
│ audio_idx ┆ i_start ┆ i_stop ┆ audio     │
│ ---       ┆ ---     ┆ ---    ┆ ---       │
│ i64       ┆ i64     ┆ i64    ┆ list[i32] │
╞═══════════╪═════════╪════════╪═══════════╡
│ 0         ┆ 1       ┆ 4      ┆ [8, 6, 2] │
│ 0         ┆ 6       ┆ 8      ┆ [5, 4]    │
│ 1         ┆ 6       ┆ 9      ┆ [0, 4, 0] │
│ 2         ┆ 3       ┆ 5      ┆ [9, 6]    │
└───────────┴─────────┴────────┴───────────┘