如何使用Polars高效提取JSON对象数组中的指定字段子集？

针对你遇到的Polars处理300GB NDJSON文件时嵌套字段提取性能瓶颈问题，核心优化方向是替换低效的list.eval操作，改用Polars原生向量化API，并优化读写流程，具体方案如下：

1. 替换list.eval为原生结构体列表提取 #

list.eval是逐元素执行的行级操作，开销极大，而Polars提供了针对结构体列表的向量化提取方法，性能能接近chdb的水平：

优化前（低效）代码： #

import polars as pl

df = pl.read_ndjson("large_file.ndjson")
result = df.select(
    "id",
    "standing",
    pl.col("baseElements").list.eval(
        pl.struct(url=pl.col("url"), valueContent=pl.col("valueContent"))
    ).alias("baseElements")
)
result.write_parquet("output.parquet")

优化后代码： #

import polars as pl

# 开启Polars快速执行引擎
pl.Config.set_fast_execution(True)

# 用懒加载模式读取文件，避免一次性加载全量数据到内存
df = pl.scan_ndjson("large_file.ndjson", low_memory=False)

# 用向量化的结构体列表提取替代list.eval
result = df.select(
    "id",
    "standing",
    # 如果需要保留结构体格式，用list.struct.select
    pl.col("baseElements").list.struct.select(["url", "valueContent"]).alias("baseElements")
    # 如果只需要单独提取字段，用list.struct.field
    # pl.col("baseElements").list.struct.field("url").alias("base_urls"),
    # pl.col("baseElements").list.struct.field("valueContent").alias("base_values")
)

# 执行查询并写入Parquet，优化写入参数
result.collect(streaming=True).write_parquet(
    "output.parquet",
    compression="snappy",  # 平衡压缩比和速度
    row_group_size=100_000  # 匹配下游处理的行组大小
)

2. 额外优化点 #

• 提前指定数据类型：如果已知字段的 dtype，读取时直接指定，避免Polars自动推断类型的开销：

  dtypes = {
      "id": pl.UInt64,
      "standing": pl.String,
      "baseElements": pl.List(pl.Struct({
          "url": pl.String,
          "valueContent": pl.String
      }))
  }
  df = pl.scan_ndjson("large_file.ndjson", dtype=dtypes, low_memory=False)
  

• 启用流式处理：collect(streaming=True)让Polars分块处理数据，适合超大文件，避免内存溢出同时提升并行效率。
• 调整Parquet写入参数：根据下游需求选择压缩算法（如zstd压缩比更高但速度稍慢，snappy速度最快），并设置合适的row_group_size，减少IO开销。

效果预期 #

通过以上优化，Polars的性能会大幅提升，尤其是替换list.eval为向量化操作后，嵌套字段提取的速度会接近chdb的水平。如果经过优化后仍无法满足需求，可以考虑混合方案：用chdb完成初始的字段提取，再用Polars处理下游逻辑，兼顾团队熟悉度和性能。

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