寻求从S3读取多行JSON到Pandas DataFrame的最优高效方法

问题背景 #

我需要从S3读取多行JSON文件到Pandas DataFrame，当前使用的方法处理1.8GB文件时在Kubernetes上会占用约20GB内存，表现很差：

obj = self.s3.Object(bucket, key)
df = pd.read_json(obj.get()["Body"].read(), orient='index', lines=True)

我自己修改后的方法速度更快且不会出现OOM，但不确定是否有更优方案：

obj = s3.Object(bucket, key)
file_content = obj.get()['Body'].read()
df = pd.DataFrame(
    [json.loads(line) for line in obj.get()['Body'].read().splitlines()]
)

注：上述代码存在冗余——两次调用obj.get()['Body'].read()会重复从S3拉取文件，优化后应复用已读取的file_content：

obj = s3.Object(bucket, key)
file_content = obj.get()['Body'].read()
df = pd.DataFrame([json.loads(line) for line in file_content.splitlines()])

额外需求：方案需支持所有多行JSON，且能无缝处理压缩格式文件（只需对读取逻辑做封装）。

更优实现方案 #

1. 利用Pandas流式读取（推荐首选） #

Pandas的read_json可以直接接受流式文件对象，无需先把整个文件读进内存。S3对象的Body本身就是一个可迭代的流式对象，直接传入即可实现逐行读取，大幅降低内存占用：

obj = s3.Object(bucket, key)
df = pd.read_json(obj.get()["Body"], orient='index', lines=True)

这个方案比手动遍历行更简洁，且Pandas内部的解析逻辑经过优化，性能优于手动json.loads循环。

2. 分块读取超大文件（Dask方案） #

如果文件体积远超内存容量，可使用Dask进行分块并行读取，最后再转换为Pandas DataFrame（适合TB级文件）：

import dask.dataframe as dd

# 分块读取S3上的多行JSON
dask_df = dd.read_json(
    f"s3://{bucket}/{key}",
    orient='index',
    lines=True,
    blocksize='100MB'  # 按100MB分块，可根据内存调整
)
# 转换为Pandas DataFrame（仅当结果能放进内存时使用）
df = dask_df.compute()

Dask会自动处理分块和并行解析，避免一次性加载整个文件。

3. 兼容压缩文件的封装函数 #

针对压缩格式（.gz/.bz2/.xz等），可以封装一个流式解压函数，保持和读取普通文件一致的接口：

import gzip
import bz2
import lzma
from botocore.response import StreamingBody

def get_compatible_stream(obj: s3.Object) -> StreamingBody:
    key = obj.key
    body = obj.get()["Body"]
    if key.endswith('.gz'):
        return gzip.GzipFile(fileobj=body)
    elif key.endswith('.bz2'):
        return bz2.BZ2File(body)
    elif key.endswith('.xz'):
        return lzma.LZMAFile(body)
    else:
        return body

# 使用示例
obj = s3.Object(bucket, key)
stream = get_compatible_stream(obj)
df = pd.read_json(stream, orient='index', lines=True)

这个函数会根据文件后缀自动选择解压方式，返回的流式对象可直接传给pd.read_json，无需修改核心读取逻辑。

方案对比 #

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