如何高效从703个制表符分隔文本生成目标格式pandas DataFrame

实现方案 #

核心思路

不需要加载文件的全部列和重复数据，仅保留判断所需的region和latitude组合，去重后添加文件ID标记，最后合并所有标记生成宽表，全程最小化内存占用。

代码实现

import os
import pandas as pd
import numpy as np

# 配置参数：替换为你的文件所在目录路径
file_dir = "./tsv_files"
file_list = [f for f in os.listdir(file_dir) if f.startswith("ID_") and f.endswith(".txt")]
df_list = []

for file_name in file_list:
    # 提取ID作为标记列名
    id_col = os.path.splitext(file_name)[0]
    file_path = os.path.join(file_dir, file_name)
    
    # 仅读取需要的两列，指定数据类型减少内存占用
    df = pd.read_csv(
        file_path,
        sep="\t",
        usecols=["region", "latitude"],
        dtype={
            "region": "category",
            "latitude": np.int32
        },
        # 如果文件极大可开启分块读取，示例为每100万行读一次
        # chunksize=1000000
    )
    
    # 分块读取的去重逻辑
    if isinstance(df, pd.io.parsers.readers.TextFileReader):
        unique_pairs = set()
        for chunk in df:
            unique_pairs.update(zip(chunk["region"], chunk["latitude"]))
        df = pd.DataFrame(unique_pairs, columns=["region", "latitude"])
    else:
        # 普通读取直接去重
        df = df.drop_duplicates(subset=["region", "latitude"], keep="first")
    
    # 添加当前文件的存在标记
    df[id_col] = 1
    df_list.append(df)

# 合并所有小表，按组合聚合，缺失值填0
final_df = pd.concat(df_list, axis=0, ignore_index=True)\
            .groupby(["region", "latitude"], as_index=False)\
            .sum()\
            .fillna(0)

# 可选：将ID列转为int8类型进一步节省内存
id_cols = [os.path.splitext(f)[0] for f in file_list]
final_df[id_cols] = final_df[id_cols].astype(np.int8)

优化说明

• 读取时仅加载必要的两列，相比读取全4列数据量直接减少一半
• 用category类型存储region、int32存储latitude、最终用int8存储0/1标记，内存占用可降低70%以上
• 读取后立即去重，单文件内重复的(region,latitude)组合不会占用额外内存
• 超大文件支持分块读取，不会出现内存溢出问题
• 如果数据量超过单机能承载的内存上限，可替换为dask.dataframe实现外存计算，API和pandas基本一致，无需大幅修改代码

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