转置Polars列中的二维数组元素

嘿，看起来你正在处理一个Polars DataFrame，其中列b的每个元素都是800×3的二维float32数组，想要把这些数组逐个转置成3×800的形状对吧？先把你提供的原始数据构建代码补全一下，方便后续操作：

# /// 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.Array(pl.Float32, 3), 800)
    }
)

下面给你两种实用的处理方法：

方法一：结合Numpy转置（直观易理解） #

这个方法用map_elements遍历列中的每个数组元素，借助numpy的转置功能完成操作，非常直观：

# 转置b列中的每个二维数组
df_transposed = df.with_columns(
    pl.col("b").map_elements(
        lambda arr: np.transpose(arr),
        return_dtype=pl.Array(pl.Array(pl.Float32, 800), 3)
    )
)

# 验证转置结果：查看第一行b元素的形状
print(df_transposed["b"][0].shape)  # 输出 (3, 800)

这里需要注意指定return_dtype，告诉Polars转置后的数组结构——外层是3个元素，每个元素是包含800个float32的数组，对应转置后的形状。

方法二：Polars原生操作（性能更优） #

如果你的数据集规模较大，优先考虑这种纯Polars原生操作的方式，不需要依赖numpy，能更好地利用Polars的向量化优化：

# 用Polars数组API完成转置
df_transposed = df.with_columns(
    pl.col("b")
    .arr.eval(pl.element().arr.to_list())  # 将内层数组转为列表，方便转置
    .arr.transpose()  # 对数组进行转置
    .arr.eval(pl.element().cast(pl.Array(pl.Float32, 800)))  # 转回Polars数组类型
)

转置完成后同样可以通过查看元素形状来确认结果是否符合预期。

备注：内容来源于stack exchange，提问作者DJDuque