[Python] Transpose 2d Array elements in a column

I have a dataframe like:

# /// 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]… │
└──────────┴─────────────────────────────────┘
"""


And I want to transpose the 2d arrays in column b to be array[f32, (3, 800)] instead.

Actually, my column is list[array[f32, 3]], and I turn it into 2d array with .list.to_array(800). In case there is a cleaner way from this.

I can do it with map_elements by hand with numpy, but my real dataframe has > 4 million rows and using map_elements (as mentioned in the docs is super slow) and it is taking for ever just doing that transpose.

Is there a faster way of doing this?

Continue reading...