Python Polars：优化JSON字符串列的过滤聚合计算方案

直接用map_elements+ast.literal_eval在大数据集下性能拉胯，推荐用Polars内置的JSON解析和矢量化操作，全程避免Python层面的循环，性能提升明显。

步骤1：解析JSON列 #

用pl.col("json").str.json_parse()把JSON字符串转成Polars的结构体类型，这一步是矢量化的，比逐行解析快得多。

步骤2：展开列表并配对x和y值 #

利用pl.struct的arr.eval结合pl.element()，把x和y的列表元素一一配对，再用explode展开成每行一个(x,y)对的形式，方便后续过滤。如果是Polars 0.19+版本，直接用list.zip配对更简洁。

步骤3：过滤条件并计算均值 #

按原行分组，过滤出0 < x < 3的行，然后对y值求平均，最后把结果合并回原DataFrame。

完整代码示例 #

兼容旧版本Polars的写法 #

import polars as pl

# 示例数据
df = pl.DataFrame({
    "json": [
        '{"x":[0,1,2,3], "y":[10,20,30,40]}',
        '{"x":[1,2,4,5], "y":[5,15,25,35]}'
    ]
})

# 高效处理流程
result = df.with_row_index("row_idx").pipe(lambda df: 
    df.with_columns(
        pl.col("json").str.json_parse().alias("parsed")
    ).with_columns(
        # 配对x和y的元素
        pl.col("parsed").arr.eval(
            pl.struct(
                x=pl.element().struct.field("x").list.get(pl.int_range(0, pl.element().struct.field("x").list.len())),
                y=pl.element().struct.field("y").list.get(pl.int_range(0, pl.element().struct.field("y").list.len()))
            ),
            parallel=True
        ).alias("pairs")
    ).explode("pairs").with_columns(
        pl.col("pairs").struct.field("x").alias("x"),
        pl.col("pairs").struct.field("y").alias("y")
    ).filter(
        pl.col("x").is_between(0, 3, closed="neither")
    ).group_by("row_idx").agg(
        pl.col("y").mean().alias("y_avg")
    ).join(df.with_row_index("row_idx"), on="row_idx").drop("row_idx")
)

print(result)

Polars 0.19+简洁写法 #

result = df.with_row_index("row_idx").pipe(lambda df:
    df.with_columns(
        pl.col("json").str.json_parse().alias("parsed")
    ).with_columns(
        # 直接用list.zip配对x和y列表
        pl.col("parsed").struct.field("x").list.zip(pl.col("parsed").struct.field("y")).alias("xy_pairs")
    ).explode("xy_pairs").with_columns(
        x=pl.col("xy_pairs").list.get(0),
        y=pl.col("xy_pairs").list.get(1)
    ).filter(
        (pl.col("x") > 0) & (pl.col("x") < 3)
    ).group_by("row_idx").agg(
        y_avg=pl.col("y").mean()
    ).join(df.with_row_index("row_idx"), on="row_idx").drop("row_idx")

性能优势说明 #

• 全程用Polars的内置矢量化函数，避免了Python循环的开销
• str.json_parse()是批量解析，比逐行ast.literal_eval快几个数量级
• 分组和聚合操作都是Polars的底层优化实现，大数据集下表现碾压纯Python方法

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