如何在Python中高效将Solr流式API返回的JSON响应加载到Polars中

你现在用外部工具转NDJSON的方案已经很靠谱了，但既然想找纯Python内的快速实现，我整理了几个适配Polars特性的思路，帮你在不跳出Python环境的前提下高效加载数据：

先明确核心需求：要提取Solr返回JSON里result-set.docs的所有元素，唯独排除最后那个标记结束的EOF条目，同时追求加载速度。

方案1：指定Schema直接用Polars解析（最快，推荐） #

你提到手动把JSON改成顶层列表后，用pl.read_json加指定Schema能跑到3秒——其实不用手动改文件，直接在Python里利用Polars的强Schema支持就能实现同样的速度：

核心思路是：先定义好数据的准确Schema，让Polars直接解析原JSON的result-set.docs数组，再过滤掉最后那个EOF条目。

代码示例：

import polars as pl

# 提前定义好你的数据Schema，要和Solr返回的字段完全匹配
custom_schema = pl.Struct({
    "col1": pl.Utf8,
    "col2": pl.Utf8,
    # 可选：添加EOF和RESPONSE_TIME字段，方便后续过滤
    "EOF": pl.Boolean,
    "RESPONSE_TIME": pl.Int64
})

# 读取原JSON，指定schema_overrides让Polars正确识别嵌套结构
df_raw = pl.read_json(
    "result.json",
    schema_overrides={
        "result-set": pl.Struct({
            "docs": pl.List(custom_schema)
        })
    }
)

# 展开docs数组、过滤无效行、清理多余字段
df = (
    df_raw
    .select(pl.col("result-set").struct.field("docs").alias("docs"))
    .explode("docs")
    .unnest("docs")
    .filter(pl.col("EOF").is_null())
    .drop("EOF", "RESPONSE_TIME")  # 按需删除不需要的标记字段
)

这个方案完全依赖Polars的Rust级解析速度，和你手动改JSON后的3秒速度几乎一致，全程在Python内完成，无外部工具依赖。

方案2：全量加载后直接转Polars（内存充裕时用） #

如果你的机器内存足够容纳整个770MiB的JSON文件，这个方法最简单直接：用Python标准库的json模块加载整个文件，提取docs数组并去掉最后一个元素，再转成Polars DataFrame：

import json
import polars as pl

with open("result.json", "r") as f:
    solr_data = json.load(f)
    # 提取docs数组，跳过最后一个EOF标记条目
    valid_docs = solr_data["result-set"]["docs"][:-1]

# 直接转成Polars DataFrame
df = pl.DataFrame(valid_docs)

这个方法速度也很快，但缺点是会把整个JSON文件加载到内存中（转成Python字典后内存占用会比原文件大不少），适合内存充裕的场景。

方案3：流式解析（内存友好，超大文件专用） #

如果要处理几十GB级的超大文件，内存不足以一次性加载，就用ijson库做流式解析，逐行生成NDJSON内容后再用Polars读取：

首先安装依赖：

pip install ijson

代码示例：

import ijson
import polars as pl
import tempfile
import json

# 用临时文件存储流式生成的NDJSON，避免占用过多内存
with tempfile.NamedTemporaryFile(mode="w+", suffix=".ndjson", delete=False) as tmp_file:
    with open("result.json", "r") as solr_file:
        # 流式迭代result-set.docs里的每个条目
        for doc in ijson.items(solr_file, "result-set.docs.item"):
            # 跳过带EOF的结束标记
            if "EOF" not in doc:
                tmp_file.write(json.dumps(doc) + "\n")
    tmp_file.flush()

# 用Polars快速读取NDJSON文件
df = pl.read_ndjson(tmp_file.name)

这个方案内存占用极低，速度接近你用外部工具的8秒，适合处理超大规模的Solr响应文件。

备注：内容来源于stack exchange，提问作者Lars Noschinski