如何根据多列拆分csv

Question

我正在尝试根据两列值将 csv 拆分为多个文件。例如，

源文件：

Header1  Header2             Header3
Alpha    energy              0.1
Alpha    energy              0.34
Beta     energy_imbalance    0.66
Beta     energy              0.7
Beta     energy              0.1
Gamma    energy_imbalance    0.3

预期输出：

输出文件1：

Header1  Header2             Header3
Alpha    energy              0.1
Alpha    energy              0.34

输出文件2：

Header1  Header2             Header3
Beta     energy_imbalance    0.66

输出文件3：

Header1  Header2             Header3
Beta     energy              0.7
Beta     energy              0.1

Outfile4：

Header1  Header2             Header3
Gamma    energy_imbalance    0.3

以下是我开始的：

filein = open('test.csv')
csvin = csv.DictReader(filein)

outputs = {}
for row in csvin:
    primaryValue = row['Header1']
    secondaryValue = row['Header2']
    if primaryValue not in outputs:
        fileout = open('{}_{}.csv'.format(primaryValue,secondaryValue),'w')
        dw = csv.DictWriter(fileout, fieldnames=csvin.fieldnames)
        dw.writeheader()
        outputs[primaryValue] = fileout, dw
    outputs[primaryValue][1].writerow(row)

for fileout, _ in outputs.values():
    fileout.close()

我能够根据 column = Header1 拆分文件，但我不确定如何进一步进行。

Answer 1

这里试试这个：

csvin = csv.DictReader(filein)
csv_files = {}
files = []

for row in csvin:
    key = (row['Header1'], row['Header2'])
    if key not in csv_files:
        # create the csv file
        fileout = open('{}_{}.csv'.format(*key), 'w')
        dw = csv.DictWriter(fileout, fieldnames=csvin.fieldnames)
        dw.writeheader()
        csv_files[key] = dw
        files.append(fileout)  # to close them later

    # write the line into to corresponding csv writer
    csv_files[key].writerow(row)

# close all files
for f in files: f.close()

Answer 2

以下是按照@Barmar 建议的方式实现的方法（即使用两个列值作为字典键）。如图所示，该键用于查找用于写入行的 csv.DictWriter 实例 — 根据需要创建新实例。它还使用一个单独的列表来关闭所有在最后打开的关联文件，以跟踪这些文件。

import csv

infile_name = 'multicol_test.csv'

with open(infile_name, newline='') as infile:
    csv_writers = {}
    files = []
    reader = csv.DictReader(infile)

    for row in reader:
        if (key := f"{row['Header1']}_{row['Header2']}") not in csv_writers:
            # Create the csv file and a corresponding DictWriter.
            outfile_name = f'{key}.csv'
            fileout = open(outfile_name, 'w', newline='')
            files.append(fileout)  # To have it closed later.
            writer = csv.DictWriter(fileout, fieldnames=reader.fieldnames)
            writer.writeheader()
            csv_writers[key] = writer

        # Write the line to corresponding csv writer.
        csv_writers[key].writerow(row)

    # Close all CSV output files.
    for f in files:
        f.close()

应用于示例输入文件，这将产生以下 csv 输出文件：

Alpha_energy.csv
Beta_energy.csv
Beta_energy_imbalance.csv
Gamma_energy_imbalance.csv

其中包含您期望的数据。

Answer 3

使用 pandas df.groupby() 是基于多列值拆分 csv 的另一种选择。

工作示例：

import pandas as pd
df = pd.read_csv('test.csv')
def df_to_grouped_csv(df):
    df_group = df.groupby(['Header1', 'Header2'])
    for name, group in df_group:
        outfile = '_'.join(name) + '.csv'
        group.to_csv(outfile, index=False)

输出：

Alpha_energy.csv
  Header1 Header2  Header3
0   Alpha  energy     0.10
1   Alpha  energy     0.34
Beta_energy.csv
  Header1 Header2  Header3
3    Beta  energy      0.7
4    Beta  energy      0.1
Beta_energy_imbalance.csv
  Header1           Header2  Header3
2    Beta  energy_imbalance     0.66
Gamma_energy_imbalance.csv
  Header1           Header2  Header3
5   Gamma  energy_imbalance      0.3

就性能而言，与the csv.DictWriter approach 相比，这应该会有所改进（特别是对于大文件）。但它确实需要导入 pandas。

性能：

Larger file (500,000 rows)
In [1]: %timeit df_to_grouped_csv()
865 ms ± 36.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [2]: %timeit csv_DictWriter_approach()
2.71 s ± 40.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)