Python3如何根据行内容将大文本文件拆分为较小的文件

Question

I have a file with the data

# FULL_ID BJD MAG UNCERT FLAG

and nearly 12,000 rows. This table contains data for 32 objects, each identified by a unique FULL_ID. So for instance it may say

# FULL_ID   BJD        MAG      UNCERT      FLAG
  2_543     3215.52    19.78    0.02937     OO
  2_543     3215.84    19.42    0.02231     OO
  3_522     3215.52    15.43    0.01122     OO
  3_522     3222.22    16.12    0.01223     OO

What I want is to run this file BigData.dat through the code, and end up with multiple files eg 2_543.dat , 3_522.dat etc, each containing:

# BJD    MAG    UNCERT    FLAG

for all rows of BigData.dat that belonged to that FULL_ID.

Currently I'm doing this:

with open(path, 'r') as BigFile:
    line = BigFile.readline()
    for line in BigFile:
        fields = line.split(None)
        id = fields[0]
        output = open(id+".dat", 'a')
        writeline = str(fields[1])+' '+str(fields[2])+' '+str(fields[3])+' '+str(fields[4])+'\n'
        output.write(writeline)
    output.close()

which does produce the correct outputs but they don't have the header line: # BJD MAG UNCERT FLAG

How can I ensure this line is at the top of each file?

Answer 1

You are overwriting the header line in the for loop, keep it in a separate variable. Additionally you could remember if the header was already written to a file:

path = 'big.dat'
header_written = []

with open(path, 'r') as BigFile:
    header = BigFile.readline()  # keep header separately!
    for line in BigFile:
        fields = line.split(None)
        _id = fields[0]
        output = open(_id+".dat", 'a')
        if _id not in header_written:  # check and save the ID to keep track if header was written
            output.write(header)
            header_written.append(_id)
        writeline = str(fields[1])+' '+str(fields[2])+' '+str(fields[3])+' '+str(fields[4])+'\n'
        output.write(writeline)
        output.close()

File:

# FULL_ID   BJD        MAG      UNCERT      FLAG
3215.52 19.78 0.02937 OO
3215.84 19.42 0.02231 OO

Answer 2

Opening a file is an expensive operation, and repeatedly doing so for each input line is not efficient. I would instead keep a mapping of seen FULL_ID values to a file object. If a FULL_ID is not present, then the file has to be opened in "w" mode and the header should be immediately added. This way:

1. the header is correctly written to the output files
2. if the script is runned more than once, the old values in output files are correctly erased

Code could be:

with open(path) as bigFile:
    outfiles = {}         # mapping FULL_ID -> output file
    header = ' '.join(['#'] + next(bigFile).split()[2:])   # compute output header
    for line in bigFile:
        row = line.split()
        try:
            output = outfiles[row[0]]
        except KeyError:
            output = open(f'{row[0]}.dat', 'w')
            print(header, file=output)
            outfiles[row[0]] = output
        print(' '.join(row[1:]), file=output)
    for output in outfiles.values():               # close all files before exiting
        output.close()

The limit is that you have to keep all files opened until the end of the input file. It should word for 32 objects, but would break for larger numbers. The efficient way would be to change the simple dict into a more sophisticated cache, able of closing the latest file when capacity is exhausted and reopen it (in append mode) if needed.

---

Here is a possible cache implementation:

class FileCache:
    """Caches a number of open files referenced by string Ids.
    (by default the id is the name)"""
    def __init__(self, size, namemapping=None, header=None):
        """Creates a new cache of size size.
        namemapping is a function that gives the filename from an ID
        header is an optional header that will be written at creation
        time
        """
        self.size = size
        self.namemapping = namemapping if namemapping is not None \
            else lambda x: x
        self.header = header
        self.map = {}             # dict id -> slot number
        self.slots = [(None, None)] * size   # list of pairs (id, file object)
        self.curslot = 0          # next slot to be used

    def getFile(self, id):
        """Gets an open file from the cache.
        Directly gets it if it is already present, eventually reopen
        it in append mode. Adds it to the cache if absent and open it
        in truncate mode."""
        try:
            slot = self.map[id]
            if slot != -1:
                return self.slots[slot][1]   # found and active
            mode = 'a'                       # need re-opening
        except:
            mode = 'w'                       # new id: create file
        slot = self.curslot
        self.curslot = (slot + 1) % self.size
        if self.slots[slot][0] is not None:  # eventually close previous
            self.slots[slot][1].close()
            self.map[self.slots[slot][0]] = -1
        fd = open(self.namemapping(id), mode)
        # if file is new, write the optional header
        if (mode == 'w') and self.header is not None:
            print(self.header, file=fd)
        self.slots[slot] = (id, fd)
        self.map[id] = slot
        return fd

    def close(self):
        """Closes any cached file."""
        for i in self.slots:
            i[1].close()
            self.map[i[0]] = -1
        self.slots = [(None, None)] * self.size

Above code would become:

with open(path) as bigFile:
    header = ' '.join(['#'] + next(bigFile).split()[2:])   # compute output header
    outfiles = FileCache(10, lambda x: x+'.dat', header) # cache FULL_ID -> file
    for line in bigFile:
        row = line.split()
        output = outfiles.getFile(row[0])
        print(' '.join(row[1:]), file=output)
    outfiles.close()               # close all files before exiting