Транспонировать элементы 2d Array в столбце Polars ⇐ Python

[Anonymous]

У меня есть такой фрейм данных:

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

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "numpy",
#     "polars",
# ]
# ///

import numpy as np
import polars as pl

n_rows = 5

a = np.random.uniform(size=n_rows).astype(np.float32)
b = np.random.uniform(size=(n_rows, 800, 3)).astype(np.float32)

df = pl.DataFrame(
{
"a": a,
"b": b
},
schema={
"a": pl.Float32,
"b": pl.Array(pl.Float32, (800, 3))
}
)

print(df)
"""
shape: (5, 2)
┌──────────┬─────────────────────────────────┐
│ a        ┆ b                               │
│ ---      ┆ ---                             │
│ f32      ┆ array[f32, (800, 3)]            │
╞══════════╪═════════════════════════════════╡
│ 0.667222 ┆ [[0.958246, 0.358944, 0.221115… │
│ 0.023049 ┆ [[0.514581, 0.48279, 0.998772]… │
│ 0.294279 ┆ [[0.10559, 0.017365, 0.236783]… │
│ 0.024168 ┆ [[0.487084, 0.438834, 0.589524… │
│ 0.259044 ┆ [[0.355924, 0.947524, 0.63777]… │
└──────────┴─────────────────────────────────┘
"""

И вместо этого я хочу транспонировать 2d-массивы в столбце b в массив[f32, (3, 800)].
На самом деле мой столбец — list[array[f32, 3]], и я превращаю его в 2d-массив с помощью .list.to_array(800). На случай, если есть более чистый способ.
Я могу сделать это с помощью map_elements вручную с помощью numpy, но мой реальный фрейм данных имеет > 4 миллионов строк и использует Map_elements (как упоминалось в документации, это очень медленно), и простое транспонирование занимает целую вечность.
Есть ли более быстрый способ сделать это?

Подробнее здесь: https://stackoverflow.com/questions/795 ... ars-column