Polars多进程并行时出现s.__getstate__()未实现崩溃问题

我有一个遍历Polars DataFrame行执行数据处理的函数，虽清楚应避免使用iter_rows()，但当前场景必须保留该写法。尝试将DataFrame拆分后用multiprocessing多进程并行执行函数，通过共享multiprocessing.Queue收集结果时触发崩溃，报错s.__getstate__()未实现。仅当同时执行循环内的d1修改操作与后续的d1赋值时才会出现此问题，单独执行其中一项则正常。

多进程代码 #

import polars as pl
import multiprocessing as mp

import src.polar_utils as pu
from tests import get_test_dict


DICT: pl.DataFrame = get_test_df("test_df_2")["test_df_2"]
PUB_DF: pl.DataFrame = get_test_df("test_df_1")["test_df_1"]

queue = mp.Queue()
n_threads = 4
DICT = DICT.filter(pl.col("IdA").is_null())
dict_splits = pu.split_uniform_rowwise(df=DICT, n_chunks=n_threads)
processes = [
    mp.get_context("spawn").Process(
        target=myfunc1,
        kwargs={
            "d1": PUB_DF,
            "d2": dict_splits[i],
            "queue": queue,
        }
    )
    for i in range(n_threads)
]
[p.start() for p in processes]
results = []
while True:
    running = any(p.is_alive() for p in processes)
    while not queue.empty():
        results.append(queue.get())
    if not running:
        break
[p.join() for p in processes]
while len(results) > 0:
    print(results.pop(0))

处理函数myfunc1代码 #

import polars as pl
import re
from itertools import compress

def myfunc1(
    d1: pl.DataFrame, d2: pl.DataFrame, aff_column="Affiliations", queue=None
) -> pl.DataFrame:
    i = 0
    for (id,) in d2.select("IdE").iter_rows():
        d1 = d1.with_columns(
            pl.col(aff_column)
            .str.extract_all(re.escape(id))
            .list.unique()
            .list.lengths()
            .alias(f"{aff_column}_{i}")
            == 1
        )
        i += 1
        print(i, end="\r")
    d1 = (
        d1.with_columns(matches=pl.concat_list(pl.col("^" + aff_column + "_[0-9]+$")))
        .drop(pl.selectors.matches("^" + aff_column + "_[0-9]+$"))
        .with_columns(
            pl.col("matches").apply(lambda x: list(compress(range(len(x)), x)))
        )
        .with_columns(pl.col("matches").list.lengths().alias("n_matches"))
        .filter(pl.col("n_matches") >= 0)
    )

    if queue:
        queue.put(d1)
        return
    return d1

报错信息 #

thread '<unnamed>' panicked at 'not implemented', /Users/runner/work/polars/polars/crates/polars-core/src/serde/series.rs:66:17
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Exception in thread QueueFeederThread:
Traceback (most recent call last):
  File "/opt/homebrew/Cellar/python@3.11/3.11.4_1/Frameworks/Python.framework/Versions/3.11/lib/python3.11/threading.py", line 1038, in _bootstrap_inner
    self.run()
  File "/opt/homebrew/Cellar/python@3.11/3.11.4_1/Frameworks/Python.framework/Versions/3.11/lib/python3.11/threading.py", line 975, in run
    self._target(*self._args, **self._kwargs)
  File "/opt/homebrew/Cellar/python@3.11/3.11.4_1/Frameworks/Python.framework/Versions/3.11/lib/python3.11/multiprocessing/queues.py", line 244, in _feed
    obj = _ForkingPickler.dumps(obj)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Cellar/python@3.11/3.11.4_1/Frameworks/Python.framework/Versions/3.11/lib/python3.11/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
  File "/Users/malte/Documents/bibinstident/env/lib/python3.11/site-packages/polars/series/series.py", line 435, in __getstate__
    return self._s.__getstate__()

错误原因分析 #

1. Polars对象序列化限制：Polars的DataFrame/Series底层基于Rust实现，默认的__getstate__方法未完全实现跨进程序列化支持。当通过mp.Queue传递处理后的d1时，Python的_ForkingPickler尝试序列化Polars对象，触发了未实现的Rust层方法，导致崩溃。
2. 操作组合触发的隐藏状态：单独执行循环内的修改或后续赋值时，生成的DataFrame未包含特殊序列化的内部状态；但两者结合后，处理过程中产生的临时列或内部结构引入了无法被Pickle序列化的Rust层对象，进而触发报错。

规避方案 #

方案1：转换为可序列化格式传递 #

在进程内将处理后的DataFrame转换为Python原生类型（如字典列表）或Polars支持的序列化格式，传递后再恢复：

• 入队时：用d1.to_dict(as_series=False)将DataFrame转为字典，或用内存字节流保存为Parquet格式
• 出队时：用pl.DataFrame.from_dict()或pl.read_parquet()恢复为DataFrame

方案2：使用Polars内置并行API #

Polars本身提供map_batches、带parallel=True参数的apply等并行化方法，无需手动拆分DataFrame和管理多进程，从根源避免序列化问题。此方案需调整代码逻辑，但能更好利用Polars的底层优化。

方案3：更换进程通信方式 #

使用multiprocessing.Manager创建共享字典，或通过文件系统传递中间结果，避免直接传递Polars对象，适合对性能要求不高的场景。

修改后的代码示例（方案1） #

处理函数修改 #

import polars as pl
import re
from itertools import compress

def myfunc1(
    d1: pl.DataFrame, d2: pl.DataFrame, aff_column="Affiliations", queue=None
) -> pl.DataFrame:
    i = 0
    for (id,) in d2.select("IdE").iter_rows():
        d1 = d1.with_columns(
            pl.col(aff_column)
            .str.extract_all(re.escape(id))
            .list.unique()
            .list.lengths()
            .alias(f"{aff_column}_{i}")
            == 1
        )
        i += 1
        print(i, end="\r")
    d1 = (
        d1.with_columns(matches=pl.concat_list(pl.col("^" + aff_column + "_[0-9]+$")))
        .drop(pl.selectors.matches("^" + aff_column + "_[0-9]+$"))
        .with_columns(
            pl.col("matches").apply(lambda x: list(compress(range(len(x)), x)))
        )
        .with_columns(pl.col("matches").list.lengths().alias("n_matches"))
        .filter(pl.col("n_matches") >= 0)
    )

    if queue:
        # 将DataFrame转为字典格式入队，规避序列化问题
        queue.put(d1.to_dict(as_series=False))
        return
    return d1

主进程结果恢复 #

# 原有多进程启动代码不变
results = []
while True:
    running = any(p.is_alive() for p in processes)
    while not queue.empty():
        # 将字典转回Polars DataFrame
        result_dict = queue.get()
        results.append(pl.DataFrame(result_dict))
    if not running:
        break
# 后续处理代码不变

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