如何高效从大体积CSV生成混合稀疏列pandas DataFrame并实现分组统计

一、高效读取含稀疏列的CSV文件 #

你之前的两种方案各有缺陷：逐列读取需要反复扫描整个CSV文件，IO开销极大；直接指定Sparse dtype报错是因为pandas的CSV解析器目前暂不支持直接读取为扩展数组类型。

可通过「单次读取全量列+批量转稀疏」的方式优化，仅扫描一次CSV，内存占用和逐列转稀疏的效果一致，效率提升非常明显：

import pandas as pd

# 提前定义列分类
dense_cols = ["ID", "其他稠密列名"] # 替换为你的2个稠密列名
sparse_cols = names_of_sparse_columns # 你的1400个稀疏列名
all_cols = dense_cols + sparse_cols

# 1. 单次读取全部列，稀疏列先指定为普通数值类型，自动识别空值为NaN
# 可根据精度需求换成float32进一步降低读取时的内存占用
dtype_map = {col: "float32" for col in sparse_cols}
# 如果稠密列有非float类型，单独指定，比如ID是整数的话
dtype_map["ID"] = "int32"

df = pd.read_csv(
    "file.csv",
    usecols=all_cols,
    dtype=dtype_map,
    na_values=["", "NA", "nan"] # 可根据你CSV里的空值标识补充
)

# 2. 批量转稀疏列
for col in sparse_cols:
    df[col] = pd.arrays.SparseArray(df[col])

如果单次读取全量还是触发内存不足，可通过chunksize参数分块处理，每块转完稀疏列再拼接：

chunks = []
for chunk in pd.read_csv(
    "file.csv",
    usecols=all_cols,
    dtype=dtype_map,
    na_values=["", "NA", "nan"],
    chunksize=100000 # 可根据可用内存调整块大小，比如10万行一块
):
    for col in sparse_cols:
        chunk[col] = pd.arrays.SparseArray(chunk[col])
    chunks.append(chunk)
df = pd.concat(chunks, ignore_index=True)

二、按ID分组统计非空值优化 #

groupby('ID').count()的本质是统计每个ID分组下各列非空值的数量，完全可以用掩码求和替代，避开稀疏DataFrame的groupby限制，效率提升明显：

# 1. 对所有稀疏列生成非空掩码（非空为1，空为0）
mask_df = df.copy()
for col in sparse_cols:
    mask_df[col] = (~mask_df[col].isna()).astype("int8")

# 2. 分组求和，结果和count完全等价
count_result = mask_df.groupby("ID").sum()

如果ID列基数较低，还可以进一步用pivot_table或者numpy.bincount做更极致的性能优化，常规场景下上述方案已经足够应对100万行、1400列的计算需求。

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