Breaking a 3GB gz file into chunks

Question

I have a 3GB gz file that I am trying to break into chunks of smaller files which are not required to be gz (I tried to make files of 10000000 lines, this is not a requirement, the chunks could be of any size), but my method is taking ages to run (about 12 minutes). I have the requirement to not use any library such as Pandas, Spark, etc. I can only use pure Python. I tried profiling my code and the write seems to be the slowest thing. Here's my code :

import gzip
import os

class FileSplitter:

    def __init__(self):
        self.parse_args(sys.argv)

    @staticmethod
    def run():
        splitter = FileSplitter()
        #run to split the big file into smaller files
        splitter.split()
    

    def split(self):

        file_number = 1
        line_number = 0

        print("Splitting %s into multiple files with %s lines" % (os.path.join(self.working_dir, self.file_base_name), str(self.split_size)))

        out_file = self.get_new_file(file_number)
        a_file = gzip.open(self.in_file, 'rt')
        t = next(a_file)
        with a_file as f: 
            for line in f:
                out_file.write(line)
                line_number += 1
                if line_number == self.split_size:
                    out_file.close()
                    file_number += 1
                    line_number = 1
                    out_file = self.get_new_file(file_number)

        out_file.close()

        print("Created %s files." % (str(file_number)))

    def get_new_file(self,file_number):
        """return a new file object ready to write to"""
        new_file_name = "%s.%s" % (self.file_base_name, str(file_number))
        new_file_path = os.path.join(self.working_dir, 'pychunks', new_file_name)
        print ("creating file %s" % (new_file_path))
        return open(new_file_path, 'w')

    def parse_args(self,argv):
        """parse args and set up instance variables"""
        try:
            self.split_size = 10000000
            if len(argv) > 2:
                self.split_size = int(argv[2])
            self.file_name = argv[1]
            self.in_file = self.file_name
            self.working_dir = os.getcwd()
            self.file_base_name = os.path.basename(self.file_name)
        except:
            print(self.usage())
            sys.exit(1)

    def usage(self):
        return """
        Split a large file into many smaller files with set number of rows.

        Usage:

            $ python file_splitter.py <file_name> [row_count]

        row_count is optional (default is 1000)"""

if __name__ == "__main__":
    FileSplitter.run()

To run my code I call a .py file containing this code with the name of my gz file, example : python myprogfile.py test.gz

Would you do that differently? Would you use threading (i tried but I couldn't succeed yet as I don't have experience with that)? What would you change in this code?

I don't know if it is my laptop that's too weak in terms of memory and CPU ...

Answer 1

Since the split files do not need to be readable text files, I would read & write in chunks of bytes, not in lines. This should be faster than reading and writing line by line. We can go with a chunk size of 4,096 bytes because 4,096 or 8,192 is a typical block size on most file systems.

Although, I should note that when I tested your script on my system with a 9 GB text file which has 81 characters per line, it took about 73 seconds to split it into files of 10,000,000 lines each, or about 772 MB per file. The fact that it takes 12 minutes on your laptop likely confirms @Graipher's suspicion that your disk's write speed is the bottleneck here.

Using a refactored version of your script that reads & writes in byte chunks (provided at the end of this review), when I ran it on the same file with --size 772 (split into files of size 772 MB) it took about 35 seconds. Still a decent improvement, but I'd be curious to know how it runs on your machine.

To implement this, I personally like having a helper method like the following which tells us the chunk sizes to read/write:

from typing import Iterator

def chunk_sizes(total_bytes: int, chunk_size: int) -> Iterator[int]:
    for _ in range(total_bytes // chunk_size):
        yield chunk_size

    if remaining_bytes := total_bytes % chunk_size:
        yield remaining_bytes

For example, list(chunk_sizes(10000, 4096)) which returns [4096, 4096, 1808] would tell us that in order to read 10,000 bytes in 4,096 byte chunks, we should read 4,096 bytes, followed by another 4,096 bytes, followed by 1,808 bytes.

General review

• There's no reason to have run() be a staticmethod, and it's weird how FileSplitter creates an instance of itself inside this method.

  If you change the top-level invocation to FileSplitter().run() then run() could be simplified to

  def run(self):
      self.split()
  

  But if you do that you might as well remove run() and make FileSplitter().split() the top-level call. My recommendation is to move the contents of run() under the __name__ == "__main__" guard and then remove run() entirely as it's no longer needed.

• Always open files with a with block to ensure the file is closed properly.

• I recommend using argparse because you get a lot of stuff for free (input type validations, help/usage message with -h, etc.).

• I would move the command-line argument parsing logic out of FileSplitter, because it should ideally only have one responsibility: splitting files.

• Specify the exact exception you want to handle with a try-except. Handling a bare except is considered a bad practice.

• I wouldn't use sys.exit inside methods you might want to test, because it exits the Python interpreter.

• In Python 3, f-strings are a much more legible way of doing string interpolation.

  Before

  new_file_name = "%s.%s" % (self.file_base_name, str(file_number))
  

  After

  new_file_name = f"{self.file_base_name}.{file_number}"
  

Refactored version

#!/usr/bin/env python3

import argparse
import gzip
import os
import itertools

from typing import BinaryIO, Iterator

CHUNK_SIZE = 4096  # bytes


def megabytes_to_bytes(megabytes: int) -> int:
    return megabytes * 1024 * 1024


def chunk_sizes(total_bytes: int, chunk_size: int) -> Iterator[int]:
    for _ in range(total_bytes // chunk_size):
        yield chunk_size

    if remaining_bytes := total_bytes % chunk_size:
        yield remaining_bytes


class FileSplitter:
    input_file_path: str
    output_directory: str
    output_file_base_name: str
    split_size: int  # size of each partition, in MB

    def __init__(
        self, input_file_path: str, output_directory: str, split_size: int
    ) -> None:
        self.input_file_path = input_file_path
        self.output_directory = output_directory
        name = os.path.basename(input_file_path)
        name_without_file_extension = os.path.splitext(name)[0]
        self.output_file_base_name = name_without_file_extension
        self.split_size = split_size

    def get_new_file(self, file_number: int) -> BinaryIO:
        new_file_name = f"{self.output_file_base_name}.{file_number}"
        new_file_path = os.path.join(self.output_directory, new_file_name)
        return open(new_file_path, "wb")

    def split(self) -> None:
        bytes_per_file = megabytes_to_bytes(self.split_size)
        os.makedirs(self.output_directory, exist_ok=True)

        with gzip.open(self.input_file_path, mode="rb") as in_file:
            for file_number in itertools.count(start=1, step=1):
                with self.get_new_file(file_number) as out_file:
                    for chunk_size in chunk_sizes(bytes_per_file, CHUNK_SIZE):
                        chunk = in_file.read(chunk_size)
                        # chunk == "" means we finished reading the input file
                        if not chunk:
                            return
                        out_file.write(chunk)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Read a gzip-compressed file "
        "and split its uncompressed contents into separate files."
    )
    parser.add_argument(
        "file_name", help="a gzip-compressed file, e.g. foo.gz"
    )
    parser.add_argument(
        "--size",
        type=int,
        default=500,
        help="desired size of each partition, in MB",
    )
    args = parser.parse_args()

    output_directory = os.path.join(os.getcwd(), "pychunks")
    splitter = FileSplitter(args.file_name, output_directory, args.size)
    splitter.split()

Answer 2

What would you change in this code?

I'd simplify the split function. Not counting lines and files myself and opening/closing the output files at different places without a with statement. Just always try to read another line and if that succeeds, put it and the next up to split_size - 1 lines into another output file.

    def split(self):
        print("Splitting %s into multiple files with %s lines" % (os.path.join(self.working_dir, self.file_base_name), str(self.split_size)))

        with gzip.open(self.in_file, 'rt') as fin:
            for file_number, line in enumerate(fin, start=1):
                with self.get_new_file(file_number) as fout:
                    fout.write(line)
                    fout.writelines(itertools.islice(fin, self.split_size - 1))

        print("Created %s files." % (str(file_number)))

While I do frequently use f when I only work on one file, I didn't like the very different names f and out_file. In such cases I use fin and fout (and I think that's rather common), so I changed that as well.

I'd probably also change the inner with to with open(...) as fout and have get_new_file return only the filename, not open the file itself. I'd rather keep the open in the with statement and have get_new_file be a function without such side effects. Or if your only reason for having get_new_file at all was that you had two places for using it, then now you don't need it anymore and could just build the filename inside the split function.

Still unfixed issues in your code:

• While you fixed one line_number = 1 (from your original) to line_number = 0, you forgot the other one. This makes your chunks one line too short. Wouldn't happen if you only had one place :-)
• Lacking import of sys, causing NameError: name 'sys' is not defined.
• Lacking creation of folder pychunks. Without that, the user gets a rather confusing traceback and FileNotFoundError: [Errno 2] No such file or directory: '/home/runner/p/pychunks/main.py.gz.1'.
• t = next(a_file) loses the first line, as that is never used. Might be what you want, but it's not documented and goes against common sense and the description that is there and gets shown when you run the script without arguments. Perhaps if you do want that, make it a command-line argument for skipping the first number of lines, with default 0.