Create csv file for each id value and write all data rows to csv with same id

Question

I have a text file which is tab delimited:

1    324    2344
1    8372   1234
2    62    12
2    872    12111
2    1211    28736
3    87636    198272

The first "column" contains the id value and the rest containing data. I would like a csv file to be created for each id value so in this example, there would be 3 csv files. And for each csv file, I would like it to contain all rows of data with the same id. So when id = 1, the csv file will contain the 2 rows of data; when id = 2, the csv file will contain 3 rows etc.

I have a script which achieves this but I don't think it's very efficient as it stores all unique id values in a list, iterates through this list, for each iteration it reads the text file line by line and if the first element in the line matches the id value, it creates a csv and writes all rows of data with the same id. But if there's 20 id values in the list, it will read the text file 20 times. For small text files, it's not a problem but for larger files, I'm wondering if it would be more efficient to only read the text file once and then create the relevant csv files according to id.

So could this script be improved?

import csv

id_list = []
mylines = []
source_file = 'path/to/source_file.txt'
target_directory = 'save/to/target_directory/'

with open (source_file, 'rt') as myfile:
    # Get and store ids
    for line in myfile:
        data = line.split()
        if data[0] in id_list:
            pass
        else:
            id_list.append(data[0])

for ids in id_list:
    with open(target_directory + ids + '.csv', 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['id', 'data_column_1', 'data_column_2']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        with open (source_file, 'rt') as myfile:
            for line in myfile:
                data = line.split()
                if data[0] == ids:
                    writer.writerow({'id': data[0], \
                        'data_column_1': data[1], \
                        'data_column_2': data[2]})

Answer 1

Maintain a collection of writers. In a situation like this, the basic strategy to avoid the scaling problem mentioned in your question is the following: (1) process the input file line by line; (2) if we see a new row ID, create a CSV writer; and (3) store that writer in a dict mapping each row ID to its writer.

Closing output files. One slightly tricky/annoying detail is to ensure that the output files are closed at the end, which we can do via a contextlib.ExitStack.

Don't hardcode the dict given to the CSV writer. Instead, zip the field names and cells together to create the dict.

Avoid hardcoded paths. In data processing scripts, it's often a good idea to take input/output paths on the command-line rather than hardcoding them (a variation on that idea is to use default paths if no command-line arguments are supplied). There are many benefits to that approach, but one of the most direct is the ability to run the code against alternative (often smaller) inputs/outputs while you are developing the script. Another is the ability to easily run your code on different computers (for example, your computer and the computer of someone on the internet giving you advice).

Even in small scripts, put code in functions. Again, there are many benefits.

import csv
import sys
import os
from contextlib import ExitStack

FIELD_NAMES = ['id', 'data_column_1', 'data_column_2']

def main(args):
    input_path, output_dir = args
    writers = {}
    with ExitStack() as stack:
        input_file = stack.enter_context(open(input_path, 'rt'))
        for line in input_file:
            cells = line.split()
            row_id = cells[0]
            if row_id not in writers:
                writers[row_id] = get_writer(output_dir, row_id, stack)
            writers[row_id].writerow(dict(zip(FIELD_NAMES, cells)))

def get_writer(output_dir, row_id, stack):
    path = os.path.join(output_dir, row_id + '.csv')
    file = open(path, 'w', newline = '', encoding = 'utf-8')
    stack.enter_context(file)
    writer = csv.DictWriter(file, fieldnames = FIELD_NAMES)
    writer.writeheader()
    return writer

if __name__ == '__main__':
    main(sys.argv[1:])

Answer 2

The csv can work with tab-delimited files. By using the delimiter='\t' argument or the dialect argument. If the delimiter varies, look at the Sniffer() class.

Break up the code into testable functions. I find writing I/O functions that accept open "file-like" objects lets me test them in a REPL/notebook using io.StringIO(). Then it's easy to change the test by editing a string instead of a file.

If source_file isn't too large, it can be loaded into a dict mapping id numbers to rows in the source file. A collections.defaultdict() is a convenient data structure for this use case.

When file names include numbers, such as id numbers or sequence numbers, it can be helpful to zero pad the numbers to the same number of digits. That way when you do list the directory ('dir' or 'ls') they are sorted properly (e.g., '01','02','10' instead of '1','10','2'), which makes it easier to see missing files, etc.

import csv
import collections
import pathlib

SOURCE_FILE = 'path/to/source_file.txt'
TARGET_DIRECTORY = 'save/to/target_directory'
TARGET_TEMPLATE = f"{TARGET_DIRECTORY}/id.csv"

FIELDNAMES = ['id', 'data_column_1', 'data_column_2']

def load_data(source_file):
    data_by_id = collections.defaultdict(list)

    reader = csv.DictReader(source_file, fieldnames=FIELDNAMES, delimiter='\t')
    for row in reader:
        data_by_id[row['id']].append(row)

    return data_by_id


def save_csv(csv_file, rows):
    writer = csv.DictWriter(csvfile, fieldnames=FIELDNAMES)
    writer.writeheader()
    writer.writerows(rows)    
    
    
def main():
    with open(SOURCE_FILE, 'rt') as source_file:
        data_by_id = load_data(source_file)

    max_id_length = max(len(id) for id in data_by_id.keys())
    target = pathlib.Path(TARGET_DIRECTORY)

    for id_, rows in data_by_id.items():
        file_name = target.with_stem(f"{id_:0>{max_id_length}}.csv")
        with open(file_name, 'w', newline='', encoding='utf-8') as csvfile:
            save_csv(csvfile, rows)