如何根据多列拆分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

预期输出：

Outfile1：

Header1  Header2             Header3
Alpha    energy              0.1
Alpha    energy              0.34

Outfile2：

Header1  Header2             Header3
Beta     energy_imbalance    0.66

Outfile3：

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

下面是一个按照@Barmar的建议来实现的方法(只是它没有使用f-string来定义csv_files字典键值)：

import csv


infile_name = 'test.csv'

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

    for row in reader:
        key = '{}_{}'.format(row['Header1'], row['Header2'])
        if key not in csv_files:
            # Create the csv file
            outfile_name = '{}.csv'.format(key)
            fileout = open(outfile_name, 'w', newline='')
            writer = csv.DictWriter(fileout, fieldnames=reader.fieldnames)
            writer.writeheader()
            csv_files[key] = writer
            files.append(fileout)  # To close them later.

        # Write the line to corresponding csv writer.
        csv_files[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 2

在这里试试这个：

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()