Как использовать полярные выражения для индексации и разрезания внешнего списка или массива numpy? ⇐ Python

[Anonymous]

У меня есть кадр данных Polars с тремя столбцами: audio_idx, i_start и i_stop. У меня также есть внешний список массивов numpy. audio_idx относится к индексу внешнего списка, а i_start и i_stop относятся к фрагменту массива numpy.
Вот пример:

Код: Выделить всё

import polars as pl
import numpy as np

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 в каждую строку, поскольку это экономит мне память. Спасибо.

Подробнее здесь: https://stackoverflow.com/questions/784 ... t-or-numpy