如何将Polars DataFrame转为PySpark DataFrame？数组列转换遇问题

转换包含数组/列表的列时，常规方法会出现类型不兼容报错，以下是复现过程及解决方案：

复现步骤 #

1. 创建PySpark DataFrame #

data = [{"id": 1, "strings": ['A', 'C'], "floats": [0.12, 0.43]},
        {"id": 2, "strings": ['B', 'B'], "floats": [0.01]},
        {"id": 3, "strings": ['C'], "floats": [0.09, 0.01]}
        ]

sparkdf = spark.createDataFrame(data)

2. 转换为Polars DataFrame #

import pyarrow as pa
import polars as pl
pldf = pl.from_arrow(pa.Table.from_batches(sparkdf._collect_as_arrow()))

3. 失败的转换尝试 #

尝试1：直接转Pandas后创建Spark DataFrame

spark.createDataFrame(pldf.to_pandas())

报错信息：

TypeError: Can not infer schema for type: <class 'numpy.ndarray'>
TypeError: Unable to infer the type of the field floats.

尝试2：指定Schema转换

schema = sparkdf.schema
spark.createDataFrame(pldf.to_pandas(), schema)

报错信息：

TypeError: field floats: ArrayType(DoubleType(), True) can not accept object array([0.12, 0.43]) in type <class 'numpy.ndarray'>

解决方案 #

问题根源是Polars数组列转Pandas后会变为numpy.ndarray类型，PySpark无法直接识别。可通过以下两种方式解决：

方法1：将Polars数组列转为Python列表 #

在转Pandas前，把数组列转换为Python原生列表：

# 针对数组列，转换为Python列表
converted_pldf = pldf.with_columns(
    [
        pl.col("strings").map_elements(lambda x: x.tolist(), return_dtype=pl.List(pl.Utf8)),
        pl.col("floats").map_elements(lambda x: x.tolist(), return_dtype=pl.List(pl.Float64))
    ]
)

# 转换为PySpark DataFrame
spark.createDataFrame(converted_pldf.to_pandas(), schema=sparkdf.schema)

方法2：通过Arrow中间层直接转换（高效） #

利用Polars和PySpark对Arrow的原生支持，跳过Pandas转换，性能更优：

# Polars转Arrow Table
arrow_table = pldf.to_arrow()

# Arrow Table转PySpark DataFrame
spark.createDataFrame(arrow_table, schema=sparkdf.schema)

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