如何将含Numpy数组的Polars DataFrame转为list[list[i64]]并写入Parquet？

问题核心 #

将含二维列表的Polars DataFrame转成Numpy数组再转回后，列类型变为object（元素为嵌套Numpy数组），Polars无法将这种类型映射为Arrow支持的嵌套列表类型，导致写入Parquet时抛出PanicException cannot convert object to arrow错误。

解决方案：显式转换为原生嵌套列表并指定类型 #

需要把每个嵌套Numpy数组转为Python原生二维列表，同时显式指定Polars的list[list[i64]]类型，确保Arrow能正确识别。

方法1：使用map_elements批量转换 #

import polars as pl
import numpy as np

# 模拟问题场景：生成含嵌套Numpy数组的DataFrame
sample_data = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
df = pl.DataFrame({"nested_data": [np.array(x) for x in sample_data]})

# 将object列转换为list[list[i64]]类型
df_fixed = df.with_columns(
    pl.col("nested_data").map_elements(
        lambda arr: arr.tolist(),
        return_dtype=pl.List(pl.List(pl.Int64))
    )
)

# 验证类型
print(df_fixed.schema)  # 输出: {'nested_data': List(List(Int64))}

# 写入Parquet
df_fixed.write_parquet("output.parquet")

# 读取验证
df_read = pl.read_parquet("output.parquet")
print(df_read)

方法2：直接重构Series（大数据量更高效） #

如果数据量较大，直接遍历列元素生成原生列表再构造Series的性能更好：

# 提取列中所有Numpy数组，转为原生列表
fixed_values = [arr.tolist() for arr in df["nested_data"]]

# 重构为指定类型的列
df_fixed = df.with_columns(
    pl.Series(
        name="nested_data",
        values=fixed_values,
        dtype=pl.List(pl.List(pl.Int64))
    )
)

关键说明 #

• Polars的object类型是通用容器，无法被Arrow自动解析为嵌套列表，必须显式转换为Python原生列表并指定嵌套dtype。
• 转换后的list[list[i64]]类型完全兼容Parquet格式，能完整保留二维数组结构，读写后不会丢失层级。

内容的提问来源于stack exchange，提问作者Sam