写入含Struct内pl.List(枚举类型)列的DataFrame触发Panic问题

问题现象 #

当写入包含pl.Struct内嵌套pl.List(自定义枚举类型)的DataFrame到IPC文件时，会触发Rust Panic，错误核心提示为：

ListArray's child's DataType must match. However, the expected DataType is Struct([Field { name: "alpha", data_type: Dictionary(UInt32, LargeUtf8, false), is_nullable: true, metadata: {"POLARS.CATEGORICAL_TYPE": "ENUM"} }]) while it got Struct([Field { name: "alpha", data_type: Dictionary(UInt32, LargeUtf8, false), is_nullable: true, metadata: {} }])

最小复现代码 #

from pathlib import Path
import polars as pl

Alphas = pl.Enum(["hello"])
MiniStruct = pl.Struct({
  "alpha": Alphas,
})
MiniStructs = pl.List(MiniStruct)
df = pl.DataFrame(
    pl.Series("xs", [[] for _ in range(1)], dtype=MiniStructs)
)

path = Path("xs.arrow")
if path.exists():
  path.unlink()

df.write_ipc(path)

报错详情 #

thread '<unnamed>' panicked at /home/runner/work/polars/polars/crates/polars-arrow/src/array/list/mod.rs:82:61:
called `Result::unwrap()` on an `Err` value: ComputeError(ErrString("ListArray's child's DataType must match. However, the expected DataType is Struct([Field { name: \"alpha\", data_type: Dictionary(UInt32, LargeUtf8, false), is_nullable: true, metadata: {\"POLARS.CATEGORICAL_TYPE\": \"ENUM\"} }]) while it got Struct([Field { name: \"alpha\", data_type: Dictionary(UInt32, LargeUtf8, false), is_nullable: true, metadata: {} }])."))
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

PanicException                            Traceback (most recent call last)
/tmp/ipykernel_264688/3257874384.py in ?()
     13 path = Path("xs.arrow")
     14 if path.exists():
     15   path.unlink()
     16 
---> 17 df.write_ipc(path)

~/bumblebee/.venv/lib/python3.12/site-packages/polars/dataframe/frame.py in ?(self, file, compression, future)
   3251             issue_unstable_warning(
   3252                 "The `future` parameter of `DataFrame.write_ipc` is considered unstable."
   3253             )
   3254 
-> 3255         self._df.write_ipc(file, compression, future)
   3256         return file if return_bytes else None  # type: ignore[return-value]

PanicException: called `Result::unwrap()` on an `Err` value: ComputeError(ErrString("ListArray's child's DataType must match. However, the expected DataType is Struct([Field { name: \"alpha\", data_type: Dictionary(UInt32, LargeUtf8, false), is_nullable: true, metadata: {\"POLARS.CATEGORICAL_TYPE\": \"ENUM\"} }]) while it got Struct([Field { name: \"alpha\", data_type: Dictionary(UInt32, LargeUtf8, false), is_nullable: true, metadata: {} }])."))

原因分析 #

Polars在构造空列表的pl.Series时，未正确保留自定义枚举类型的元数据（即POLARS.CATEGORICAL_TYPE: ENUM标识），导致内部ListArray的子Struct类型与预期的类型元数据不匹配，触发类型校验失败。

解决方案 #

通过先构造包含枚举数据的非空Series，再转为空列表，强制Polars保留枚举类型的元数据：

from pathlib import Path
import polars as pl

Alphas = pl.Enum(["hello"])
MiniStruct = pl.Struct({
  "alpha": Alphas,
})
MiniStructs = pl.List(MiniStruct)

# 先构造带枚举数据的非空Series，确保元数据被正确初始化
df = pl.DataFrame(
    pl.Series("xs", [[MiniStruct({"alpha": Alphas("hello")})]], dtype=MiniStructs)
)
# 将列表清空为目标空状态
df = df.with_columns(pl.col("xs").list.clear())

path = Path("xs.arrow")
if path.exists():
  path.unlink()

df.write_ipc(path)

此方法通过先让Polars识别到枚举类型的元数据，再清空列表，既保留了正确的类型信息，又得到了空列表的DataFrame，可正常写入IPC文件。

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