Python CSV 读取器性能优化

有一段 Python 代码，用来读取三个 CSV 文件，并从每个文件中提取特定列，然后将这些列合并成一个新的 CSV 文件。代码如下：

import gf
import csv

cic = gf.ct
cii = gf.cit
li = gf.lt
oc = "Output.csv"

with open(cic, "rb") as input1:
  reader = csv.DictReader(cie,gf.ctih)
  with open(oc,"wb") as outfile:
    writer = csv.DictWriter(outfile,gf.ctoh)
    writer.writerow(dict((h,h) for h in gf.ctoh))
    next(reader)
    for ci in reader:
      row = {}
      row["ci"] = ci["id"]
      row["cyf"] = ci["yf"]
      with open(cii,"rb") as ciif:
        reader2 = csv.DictReader(ciif,gf.citih)
        next(reader2)
        with open(li, "rb") as lif:
          reader3 = csv.DictReader(lif,gf.lih)
          next(reader3)
          for cii in reader2:
            if ci["id"] == cii["id"]:
              row["ci"] = cii["ca"]
          for li in reader3:
            if ci["id"] == li["en_id"]:
              row["cc"] = li["c"]
      writer.writerow(row)

这段代码的功能是：从第一个 CSV 文件中读取两列数据，然后分别从第二个和第三个 CSV 文件中读取一列数据，并将其合并成一个新的 CSV 文件。代码的逻辑是正确的，但是运行速度非常慢。

2. 解决方案

2.1 使用 pandas DataFrame

一种提高代码效率的方法是使用 pandas DataFrame。pandas 是一个功能强大的数据分析库，可以轻松处理大型数据集。我们可以使用 pandas 来读取三个 CSV 文件，然后使用 merge() 方法将它们合并成一个 DataFrame。代码如下：

import pandas as pd

df1 = pd.read_csv(cic)
df2 = pd.read_csv(cii)
df3 = pd.read_csv(li)

df = pd.merge(df1, df2, on="id")
df = pd.merge(df, df3, on="en_id")

df.to_csv("Output.csv", index=False)

这种方法比之前的代码快很多，因为它只需要读取三个 CSV 文件一次，然后就可以对它们进行合并操作。

2.2 使用 csv.Sniffer()

另一种提高代码效率的方法是使用 csv.Sniffer()。csv.Sniffer() 可以自动检测 CSV 文件的分隔符和引用符。我们可以使用它来读取第一个 CSV 文件，然后将分隔符和引用符信息传递给 csv.DictReader()。代码如下：

import csv
import io

with open(cic, "rb") as input1:
  dialect = csv.Sniffer().sniff(input1.read(1024))
  input1.seek(0)
  reader = csv.DictReader(input1, dialect=dialect)

  with open("Output.csv", "wb") as outfile:
    writer = csv.DictWriter(outfile, fieldnames=["ci", "cyf", "ci", "cc"])
    writer.writeheader()

    for row in reader:
      row["ci"] = row["id"]
      row["cyf"] = row["yf"]

      with open(cii, "rb") as ciif:
        reader2 = csv.DictReader(ciif, dialect=dialect)
        for cii in reader2:
          if row["id"] == cii["id"]:
            row["ci"] = cii["ca"]
            break

      with open(li, "rb") as lif:
        reader3 = csv.DictReader(lif, dialect=dialect)
        for li in reader3:
          if row["id"] == li["en_id"]:
            row["cc"] = li["c"]
            break

      writer.writerow(row)

这种方法比之前的代码快一些，因为它只需要读取第一个 CSV 文件一次，然后就可以对它进行迭代。

2.3 使用多进程

如果需要处理非常大的 CSV 文件，还可以使用多进程来提高代码效率。我们可以将读取 CSV 文件和合并数据的任务分配给多个进程，这样可以同时执行这些任务，从而减少总的运行时间。代码如下：

import csv
import io
import multiprocessing

def read_csv(filename):
  with open(filename, "rb") as input1:
    dialect = csv.Sniffer().sniff(input1.read(1024))
    input1.seek(0)
    reader = csv.DictReader(input1, dialect=dialect)
    return list(reader)

def merge_data(data):
  df = pd.DataFrame(data)
  return df.to_csv(StringIO(), index=False, header=False).getvalue().decode("utf-8")

if __name__ == "__main__":
  pool = multiprocessing.Pool()

  data1 = read_csv(cic)
  data2 = read_csv(cii)
  data3 = read_csv(li)

  result = pool.starmap(merge_data, [(data1, data2, data3)])

  with open("Output.csv", "wb") as outfile:
    outfile.write(result[0].encode("utf-8"))

这种方法可以大大提高代码的效率，因为它可以同时处理多个 CSV 文件。