如何将多输出Numpy数组转为Polars DataFrame？KMeans场景报错求解

我在做KMeans图像聚类分析时，想把处理后的Numpy数组存入Polars DataFrame，但数组总是被解析为object类型，导致报错：polars.exceptions.ComputeError: ValueError: cannot parse numpy data type dtype('O') into Polars data type。尝试用tolist()转换又耗时极长，处理多张图片时性能问题更严重。

相关代码如下：

import matplotlib.pyplot as plt
import polars as pl
from sklearn.cluster import KMeans


def kmeans_img(filepath: str, n_clusters: int):
    image = plt.imread(filepath)
    X = image.reshape(-1, 3)
    kmeans = KMeans(n_clusters=n_clusters)
    kmeans.fit(X)
    segmented_image = kmeans.cluster_centers_[kmeans.labels_].reshape(image.shape)

    return (
        image,
        segmented_image,
    )

后续调用代码：

butterfly_df = pl.DataFrame({"filepath": butterfly_filepath})

out = (
    butterfly_df.with_columns(
        results=pl.col("filepath").map_elements(lambda x: kmeans_img(x, 3))
    )
    .with_columns(
        og_img=pl.col("results").list.get(1),
        seg_img=pl.col("results").list.get(1),
    )
    .drop("results")
)

核心问题原因 #

Polars无法自动将嵌套的Numpy数组解析为合法列类型，map_elements返回的元组中的数组被默认识别为object类型，后续提取时触发类型解析错误。直接用tolist()会把多维数组转成嵌套列表，带来巨大的性能开销。

解决方案 #

方案1：直接返回结构化数据（推荐） #

修改kmeans_img函数，返回Polars支持的结构化类型（pl.Struct），明确指定数组类型，避免object类型产生：

import matplotlib.pyplot as plt
import polars as pl
from sklearn.cluster import KMeans

def kmeans_img(filepath: str, n_clusters: int):
    image = plt.imread(filepath)
    X = image.reshape(-1, 3)
    kmeans = KMeans(n_clusters=n_clusters)
    kmeans.fit(X)
    segmented_image = kmeans.cluster_centers_[kmeans.labels_].reshape(image.shape)
    
    # 返回Polars结构体，明确数组类型
    return pl.struct(
        og_img=pl.lit(image).cast(pl.Array(pl.Array(pl.Array(pl.UInt8, 3), image.shape[1]), image.shape[0])),
        seg_img=pl.lit(segmented_image).cast(pl.Array(pl.Array(pl.Array(pl.Float64, 3), image.shape[1]), image.shape[0]))
    )

# 使用方式
butterfly_df = pl.DataFrame({"filepath": [butterfly_filepath]})

out = butterfly_df.with_columns(
    pl.col("filepath").map_elements(
        lambda x: kmeans_img(x, 3),
        return_dtype=pl.Struct([
            pl.Field("og_img", pl.Array(pl.Array(pl.Array(pl.UInt8, 3), None), None)),
            pl.Field("seg_img", pl.Array(pl.Array(pl.Array(pl.Float64, 3), None), None))
        ])
    ).alias("results")
).unnest("results")

方案2：保留原函数，指定返回类型 #

如果不想修改kmeans_img的返回值，可以在map_elements中直接指定返回类型为包含数组的结构体：

# 保持原kmeans_img函数不变
out = butterfly_df.with_columns(
    pl.col("filepath").map_elements(
        lambda x: kmeans_img(x, 3),
        return_dtype=pl.Struct([
            pl.Field("og_img", pl.Array(pl.Array(pl.Array(pl.UInt8, 3), None), None)),
            pl.Field("seg_img", pl.Array(pl.Array(pl.Array(pl.Float64, 3), None), None))
        ])
    ).alias("results")
).unnest("results")

性能优化要点 #

• 绝对避免tolist()：直接传递Numpy数组给Polars，利用其对Numpy的原生支持，减少转换开销。
• 提前指定类型：明确return_dtype让Polars提前分配内存，避免运行时类型推断的额外消耗。
• 使用unnest展开结构体：替代先存中间列再提取的方式，减少数据复制步骤。

补充说明 #

如果处理的图片尺寸不固定，可以用pl.List(pl.List(pl.List(pl.UInt8)))替代固定长度的pl.Array，Polars会自动适配可变维度的数组。

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