Simple sharding of delimited files to more sophisticated

Question

A while back I came across the issue of having a large delimited file where I wanted to simply parallelize my python code across each line of the file. However, doing so all at once either took up too much RAM, or took too long.

So I made a simple Parallel Parser (ParPar) for sharding such delimited files and auto mapping a function across either each line of the file, or across each file of the shard.

It has served me sufficiently well, but I am quite sure it could be improved.

I am posting the relevant code here, but there may be snippets left out which can be obtained at the repo.

Issues:

• The only main performance issue I have faced is that if whatever I am doing takes most / all cores, uses all ram, and then chews through swap, leaving everything in a laggy frozen state. It would be nice to know how to set a mem limit or something for the processes spawned such that if it hits the limit the processes are paused until some of the ongoing ones finish and release memory.

• This creates shards by making many sub-directories with one file in the base. It may improve by being able to multi-shard... (shard a shard?) for when the first issue is encountered

Design Goals

I wanted to make a sharding library that was user friendly for the following:

• taking a delimited file split it, either by line, or by columns (if columns have n-categories) and putting the subfiles in a directory accordingly
• being able to restore a file sharded by the library (although not necessarily in the same order)
• easily allow one to map a function to all of the sharded files, or to each line in the sharded files
• allow one to have a progress bar for the status of the threaded function

Resources:

• repository
• pypi

Example Usage

from parpar import ParPar
ppf = ParPar()
ppf.shard_by_lines(
    input_file="./short.csv",
    output_dir="./short-lines",
    number_of_lines=3,
) # 3 lines per sharded file
ppf.shard(
    input_file="./short.csv",
    output_dir="./short-cols",
    delim=",",
    columns=[0,5], # assumes column 0 and column 5 are categorical
) 

def foo(file, a, b, **kwargs):
    print(file, kwargs["dest"])

ppf.shard_file_apply(
    './short-lines', # sharded directory
    foo,
    args=[1,2],
    kwargs={
        "output_shard_name": 'short-lines-alt',
        "output_shard_loc": os.path.expanduser('~/Desktop')
        # creates a new shard (with no files) with corresponding name at the specified loction
        # in this case `/Desktop/short-lines-alt/<lines>/
        # using kwargs["dest"] and kwargs["lock"] user can
        # safely create a new shard processing each sharded file at a new location
    }    
)

Code

Notes:

• the following functions are all inside a class called ParPar which takes no arguments for initialization.

• This class houses two key variables:

  • _shared_current = Value('i', 0): shared memory counter
  • _shared_lock = Lock(): for race conditions, etc

• this depends on a package called sil for the progress bar

subnote: while _shared_current could be written as Value('i', 0, lock=True) and then exclude the Lock, since the Lock is not just for the counter, I opted for excluding that option.

Sharding Files

def shard(self,
        input_file: str,
        output_dir: str,
        columns: list,
        delim: str = '\t',
        newline: str = '\n',
        sil_opts:dict = default_sil_options
    )->list:
        '''
        Given the `input_file` and the column indicies, reads each line of the
        `input_file` and dumps the content into a file in the directory:

            output_dir/<columns[0]>/.../<columns[N]>/basename(<input_file>)

        where <columns[i]> is the value found in the current line at position i
        after being split by the specified `delim`.

        WARNINGS:
            everything in specified `output_dir` will be PERMANENTLY REMOVED.

        Arguments:
            input_file (str): full path to the file to be sharded.

            output_dir (str): full path to a directory in which to dump. WARNING:
                everything in specified directory will be PERMANENTLY REMOVED.

            columns (list): the columns (indices) across which to shard. The
                values found in these columns will be used as directory
                names (nested).

            delim (str): How to split each field in a line of the file.
                Defaults to '\t'.

            newline (str): How to split each line of the file. Defaults to '\n'.

            sil_opts (dict): Defaults to {'length': 40, 'every': 1}. See the
                Sil package.

        Returns:
            sharded_files (list): list of all the sharded files
        '''
        lno = filelines(input_file) # number of lines in the input_file
        sts = Sil(lno, **sil_opts)  # status indicator
        basename = os.path.basename(input_file)
        files_made = set({})
        file_objs = {}

        # delete current `output_dir` if it already exists
        if os.path.isdir(output_dir):
            shutil.rmtree(output_dir)

        with open(input_file, 'r') as f:
            for line in f:
                fields = linesplit(line, delim, newline)
                dest = dir_from_cols(fields, columns)
                files_made.add(dest)
                dir_path = os.path.join(output_dir, dest)
                if not os.path.isdir(dir_path):
                    os.makedirs(dir_path)
                    file_objs[dir_path] = open(os.path.join(dir_path, basename), 'a')

                o = file_objs[dir_path]
                o.write(linemend(fields, delim, newline))
                suffix = '\t{} files made'.format(len(files_made))
                sts.tick(suffix=suffix)

        # close all made file objects
        for fo in file_objs.values():
            fo.close()

        return self.shard_files(output_dir)

    def shard_by_lines(self,
        input_file:str,
        output_dir:str,
        number_of_lines:int,
        sil_opts:dict = default_sil_options
    )->list:
        '''
        Given the input_file and the columns, reads each line of the input_file
        into output files in subdirectories labeled by the line numbers
        `'start_stop'` based on the value `number_of_lines`:

            output_dir/<n>_<n+number_of_lines>/basename(<input_file>)

        WARNINGS:
            everything in specified `output_dir` will be PERMANENTLY REMOVED.

        Arguments:
            input_file (str): full path to the file to be sharded.

            output_dir (str): full path to a directory in which to dump. WARNING:
                everything in specified directory will be PERMANENTLY REMOVED.

            number_of_lines (int): the number of lines which should be at most in
                each sharded file.

            sil_opts (dict): Defaults to `{'length': 40, 'every': 1}`. See the
                Sil package.

        Returns:
            sharded_files (list): list of all the sharded files
        '''
        lno = filelines(input_file)
        sts = Sil(lno, **sil_opts)
        basename = os.path.basename(input_file)
        files_made = set({})

        file_objs = {}

        if os.path.isdir(output_dir):
            shutil.rmtree(output_dir)

        with open(input_file, 'r') as f:
            tally = 0
            while tally < lno:
                if tally % number_of_lines == 0:

                    dest = '{}_{}'.format(tally, tally+number_of_lines)
                    files_made.add(dest)
                    dir_path = os.path.join(output_dir, dest)
                    if not os.path.isdir(dir_path):
                        os.makedirs(dir_path)
                        file_objs[dir_path] = open(os.path.join(dir_path, basename), 'a')

                    o = file_objs[dir_path]
                    for i in range(number_of_lines):
                        line = f.readline()
                        if not line:
                            break
                        o.write(line)
                        sts.tick()
                    tally += number_of_lines

        for fo in file_objs.values():
            fo.close()

        return self.shard_files(output_dir)

retrieving the sharded files:

def shard_files(self, directory:str)->list:
        '''
        Arguments:
            directory (str): The top-most directory of a shared file.

        Returns:
            (list): The list of all files under directory (regardless of depth).
        '''
        file_paths = []
        for path, subdirs, files in os.walk(directory):
            if not files: continue
            file_paths += [
                os.path.join(path, f) for f in files
                if 'DS_Store' not in f
            ]
        return file_paths

restoring a shard:

def assemble_shard(self, directory:str, delim:str='\t', newline:str='\n')->list:
        '''
        Arguments:
            directory (str): The top-most directory of a shared file.

        Keyword Arguments:
            delim (str): Defaults to '\t'
            newline (str): Defaults to '\n'

        Returns:
            (list): The list of lists, where each sub-list is a record found
                in one of the sharded leaf files after being split by delim.
                (i.e. all records are returned together)
        '''
        results = []
        files = self.shard_files(directory)


        with Pool(processes=os.cpu_count()) as pool:
            sarg = [(f, delim, newline) for f in files]
            lines = pool.starmap(readlines_split, sarg)

        return list(itertools.chain.from_iterable(lines))

Shard across lines

def shard_line_apply(self,
        directory:str,
        function,
        args:list=[],
        kwargs:dict={},
        processes:int=None,
        sil_opts:dict=default_sil_options
    ):
        '''
        Parallelizes `function` across each _line_ of the sharded files found as
        the leaves of `directory`.

        Notes:
            - if `processes` is None, **all** of them will be used i.e.
                `os.cpu_count()`

            - Several keywords for `kwargs` are reserved:
                1. lock:          a lock, if needed, to prevent race conditions.
                2. full_path:     the full path to the file which was opened.
                3. relative_path: the path under (directory) to the file which
                                  was opened.
                4. output_shard_name: if provided will rename the shard
                5. output_shard_loc:  if provided will move the shard

        Arguments:
            directory (str): The top-most directory of a shared file.

            function (func): The function which will be parallelized. This
                function **MUST** be defined so that it can be called as:
                    `function(line, *args, **kwargs)`

        Keyword Arguments:
            args (list): arguments to be passed to `function` on each thread.

            kwargs (dict): key-word arguments to be passed to `function` on each
                thread.

            processes (int): The number of processes to spawn. Defaults to
                **ALL** availble cpu cores on the calling computer.

            sil_opts (dict): Defaults to `{'length': 40, 'every': 1}`.
                See the Sil package.
        Returns:
            None
        '''
        if processes is None: processes = os.cpu_count()
        sfiles = self.shard_files(directory)
        records = self.sharded_records(sfiles)
        sts = Sil(records, **sil_opts)

        with Pool(processes=processes) as pool:
            self._shared_current.value = -1

            sargs = [
                (directory, file, sts, function, args, kwargs) for file in sfiles
            ]
            results = pool.starmap(self._shard_line_apply, sargs)

            pool.close()
            pool.join()
            pool.terminate()
        return results


    def _shard_line_apply(self,
        directory:str,
        file:str,
        status,
        function,
        args:list,
        kwargs:dict
    ):
        # multiprocessing.Lock
        kwargs['lock']            = self._shared_lock
        # multiprocessing.Value (shared memory counter for progress bar)
        kwargs['shared_current']  = self._shared_current
        # an instance of Sil
        kwargs['status']          = status

        # full path to the current sharded file being processes
        kwargs['full_path']       = file
        kwargs['shard_name']      = shardname(file, directory)
        kwargs['shard_dir']       = sharddir(file, directory)
        kwargs['shard_loc']       = shardloc(file, directory)
        kwargs['relative_path']   = os.path.join(*superdirs(file, directory))
        kwargs['current_process'] = current_process().name


        cp = kwargs['current_process']

        force_overwrite = kwargs['force_overwrite'] if 'force_overwrite' in kwargs else True

        os_name = kwargs['output_shard_name'] if 'output_shard_name' in kwargs else None
        os_loc  = kwargs['output_shard_loc']  if 'output_shard_loc'  in kwargs else kwargs['shard_loc']

        if os_name is not None:
            dest = os.path.join(os_loc, os_name, os.path.dirname(kwargs['relative_path']))
            kwargs['dest'] = dest
            self._shared_lock.acquire()
            try:
                if os.path.isdir(dest):
                    if force_overwrite:
                        shutil.rmtree(dest)
                else:
                    os.makedirs(dest)
            finally:
                self._shared_lock.release()


        with open(file, 'r') as f:
            for line in f:
                function(line, *args, **kwargs)
                self._shared_lock.acquire()
                try:
                    self._shared_current.value += 1
                    suffix = '\tprocess: {}'.format(cp)
                    status.update(current=self._shared_current.value, suffix=suffix)
                finally:
                    self._shared_lock.release()

Shard across files

def shard_file_apply(self,
        directory:str,
        function,
        args:list=[],
        kwargs:dict={},
        processes:int=None,
        sil_opts:dict=default_sil_options
    ):
        '''
        Parallelizes `function` across each of the sharded files found as
        the leaves of `directory`.

        Notes:
            - if `processes` is None, **all** of them will be used i.e.
                `os.cpu_count()`

            - Several keywords for `kwargs` are reserved:
                1. lock:          a lock, if needed, to prevent race conditions.
                2. full_path:     the full path to the file which was opened.
                3. relative_path: the path under (directory) to the file which
                                  was opened.
                4. output_shard_name: if provided will rename the shard
                5. output_shard_loc:  if provided will move the shard

        Arguments:
            directory (str): The top-most directory of a shared file.

            function (func): The function which will be parallelized. This
                function **MUST** be defined so that it can be called as:
                    `function(line, *args, **kwargs)`

        Keyword Arguments:
            args (list): arguments to be passed to `function` on each thread.

            kwargs (dict): key-word arguments to be passed to `function` on each
                thread.

            processes (int): The number of processes to spawn. Defaults to
                **ALL** availble cpu cores on the calling computer.

            sil_opts (dict): Defaults to `{'length': 40, 'every': 1}`.
                See the Sil package.
        Returns:
            None
        '''
        if processes is None: processes = os.cpu_count()
        sfiles = self.shard_files(directory)
        records = self.sharded_records(sfiles)
        sts = Sil(records, **sil_opts)

        with Pool(processes=processes) as pool:
            self._shared_current.value = -1

            sargs = [
                (directory, file, sts, function, args, kwargs) for file in sfiles
            ]
            pool.starmap(self._shard_file_apply, sargs)

            pool.close()
            pool.join()
            pool.terminate()


    def _shard_file_apply(self,
        directory:str,
        file:str,
        status,
        function,
        args:list,
        kwargs:dict
    ):
        # multiprocessing.Lock
        kwargs['lock']            = self._shared_lock
        # multiprocessing.Value (shared memory counter for progress bar)
        kwargs['shared_current']  = self._shared_current
        kwargs['status']          = status
        kwargs['full_path']       = file
        kwargs['shard_name']      = shardname(file, directory)
        kwargs['shard_dir']       = sharddir(file, directory)
        kwargs['shard_loc']       = shardloc(file, directory)
        kwargs['relative_path']   = os.path.join(*superdirs(file, directory))
        kwargs['current_process'] = current_process().name

        force_overwrite = kwargs['force_overwrite'] if 'force_overwrite' in kwargs else True


        os_name = kwargs['output_shard_name'] if 'output_shard_name' in kwargs else None
        os_loc = kwargs['output_shard_loc'] if 'output_shard_loc' in kwargs else kwargs['shard_loc']

        if os_name is not None:
            dest = os.path.join(os_loc, os_name, os.path.dirname(kwargs['relative_path']))
            kwargs['dest'] = dest
            self._shared_lock.acquire()
            try:
                if os.path.isdir(dest):
                    if force_overwrite:
                        shutil.rmtree(dest)
                else:
                    os.makedirs(dest)
            finally:
                self._shared_lock.release()


        function(file, *args, **kwargs)

random utils

import os
def filelines(full_filename):
    """
    Quickly returns the number of lines in a file.
    Returns None if an error is raised.
    """
    with open(full_filename, 'r') as file:
        return sum(1 for line in file)
    return None

def slices(arr, slcs):
    '''
    Args:
        arr (list): a list of elements.
        slcs (list): a list of slice specifications.

    Returns:
        (list): A list of elements extracted from arr given all slcs.
    '''
    return [col for slc in slcs for col in arr[slc]]

def linesplit(line, delim='\t', newline='\n'):
    '''
    Args:
        line (str): a line in a file
    Kwargs:
        delim (str): Defaults to '\t'.
        newline (str): Defaults to '\n'

    Returns:
        (str): the line split and devoid of newline.
    '''
    return line.rstrip(newline).split(delim)

def linemend(line, delim='\t', newline='\n'):
    '''
    Args:
        line (str): a line in a file
    Kwargs:
        delim (str): Defaults to '\t'.
        newline (str): Defaults to '\n'

    Returns:
        (str): the line joined and given a newline. (opposite of linesplit)
    '''
    return delim.join(line)+newline

def readsplit(file_object, delim='\t', newline='\n'):
    '''
    Args:
        file_object (file object): an opened file
    Kwargs:
        delim (str): Defaults to '\t'.
        newline (str): Defaults to '\n'

    Returns:
        (str): Calls linesplit on file_object.readline()
    '''
    return linesplit(file_object.readline(), delim, newline)

def readslice(file_object, slcs, delim='\t', newline='\n', add_newline_q=True):
    '''
    Args:
        file_object (file object): an opened file
        slcs (list):a list of slice specifications.
    Kwargs:
        delim (str): Defaults to '\t'.
        newline (str): Defaults to '\n'
        add_newline_q (bool): Defaults to True

    Returns:
        (str): Reads the line, splits it by slice, and then joins it back together.
    '''
    return delim.join(slices(readsplit(file_object, delim, newline), slcs)) \
    + (newline if add_newline_q else '')

def lineslice(line, slcs, delim='\t', newline='\n', add_newline_q=True):
    '''
    Args:
        line (str): a line in a file
        slcs (list):a list of slice specifications.
    Kwargs:
        delim (str): Defaults to '\t'.
        newline (str): Defaults to '\n'
        add_newline_q (bool): Defaults to True

    Returns:
        (str): Reads the line, splits it by slice, and then joins it back together.
    '''
    return delim.join(slices(linesplit(line, delim, newline), slcs)) \
    + (newline if add_newline_q else '')


def superdirs(location, up_to):
    '''
    Args:
        location (str): a full path.
        up_to (str): a sub-specficiation of full path. e.g. root --> dir_x
            where full path is: /dir_1/.../dir_x/.../dir_n
    Returns:
        (list): list of directories that come after up_to in location
    '''
    relative_path = location.replace(up_to, '').lstrip('/')
    sups = list(filter(lambda s: s is not '', relative_path.split('/')))
    return sups


def sharddir(location, up_to):
    '''
    Args:
        location (str): a full path.
        up_to (str): a sub-specficiation of full path. e.g. root --> dir_x
            where full path is: /dir_1/.../dir_x/.../dir_n
    Returns:
        (str): all the directories prior to up_to (dir_x)
    '''
    supdirs = os.path.join(*superdirs(location, up_to))
    return location.replace(supdirs, '').rstrip('/')

def shardname(location, up_to):
    '''
    Args:
        location (str): a full path.
        up_to (str): a sub-specficiation of full path. e.g. root --> dir_x
            where full path is: /dir_1/.../dir_x/.../dir_n
    Returns:
        (str): the name of the directory prior to up_to (dir_x).
    '''
    return os.path.basename(sharddir(location, up_to))

def shardloc(location, up_to):
    '''
    Args:
        location (str): a full path.
        up_to (str): a sub-specficiation of full path. e.g. root --> dir_x
            where full path is: /dir_1/.../dir_x/.../dir_n
    Returns:
        (str): the name of the directory prior to up_to (dir_x).
    '''
    return os.path.dirname(sharddir(location, up_to))




def readlines_split(file, delim='\t', newline='\n'):

    lines = []
    with open(file, 'r') as f:
        for line in f:
            lines.append(linesplit(line, delim, newline))
    return lines


def dir_from_cols(fields, cols):
    '''
    Args:
        fields (list): elements in a record.
        cols (list): list of column indicies to take.

    Returns:
        (str): the directory produced by taking fields[cols[i]] for each column.
            e.g. fields[cols[0]]/.../fields[cols[n]]
    '''
    return os.path.join(*[fields[i] for i in cols])

Answer 1

Mutable default arguments

sil_opts:dict = default_sil_options

This can create some nasty bugs, and some Python inspectors specifically recommend against doing this. The alternative is to set the default to None, and then write an if in your function to assign the default dictionary if the argument is None.

Redundant return

It seems like you intend for this to return None if something bad happens. That's not how it actually works - instead, when an exception is thrown, there is no return statement executed at all.

    with open(full_filename, 'r') as file:
        return sum(1 for line in file)
    return None

For this to work as you intend, you need to add a try/except.

Quote standard

You use a mix of double-quotes and single-quotes in your docstrings. PyCharm suggests double quotes, but whatever you do, get consistent.

Meaningful documentation

It's great that you've added a lot of documentation. The next step is to make it actually mean something. For example, this:

add_newline_q (bool): Defaults to True

is deeply unhelpful. The type and default are already obvious, due to the function signature. So what does this actually do?

Consider using more comprehensions

sups = list(filter(lambda s: s is not '', relative_path.split('/')))

may be more easily expressed as

sups = [s for s in relative_path.split('/') if s != '']

Generators

In the same sharddir function, rather than returning a list (which forces the use of memory), you can leave it as a generator:

return (s for s in relative_path.split('/') if s != '')

This allows the caller to decide whether they want to materialize the generator into memory, or only iterate over it.

Quickly isn't quickly

"""
Quickly returns the number of lines in a file.
"""

Unfortunately, this isn't a "quick" operation. Given that your focus here is on parallel performance, you need to rethink how to divide up your file into lines. You should not be doing a pre-parse step. Instead, the most reasonable approach is to probably get the file size, divide it by the number of workers, seek to the boundaries in the file between each chunk, find the actual file boundary, and then give each worker its start and end offset. Otherwise, counting each line in a file is both expensive and non-parallel.

Performance of readlines_split

def readlines_split(file, delim='\t', newline='\n'):

    lines = []
    with open(file, 'r') as f:
        for line in f:
            lines.append(linesplit(line, delim, newline))
    return lines

This function is performance-problematic for a few reasons. You're building up a list in memory, when you shouldn't. If you're assembling a "shard file" whose contents are totally unmodified from the source file other than being a subset, then this is not a good way to go. If the line count distribution between shards needs to be exact, then each worker should be counting the number of newlines in its chunk. If not, don't even count (and certainly don't create a list). Once the actual boundaries are determined, write out the chunk from the source to the destination, potentially in sub-chunks of a few MB.

Usage of kwargs

This chunk:

def foo(file, a, b, **kwargs):
    print(file, kwargs["dest"])

ppf.shard_file_apply(
    './short-lines', # sharded directory
    foo,
    args=[1,2],
    kwargs={
        "output_shard_name": 'short-lines-alt',
        "output_shard_loc": os.path.expanduser('~/Desktop')

shows you using kwargs in a somewhat awkward way. You should probably change the first instance to

def foo(file, a, b, dest):
    print(file, dest)

Any argument can still be a named argument from the caller, including dest.

As to the second instance, probably change to

ppf.shard_file_apply(
    './short-lines', # sharded directory
    foo,
    args=[1,2],
    output_shard_name='short-lines-alt',
    output_shard_loc=os.path.expanduser('~/Desktop')

Those last two argument can still be captured in kwargs without you passing a dictionary explicitly.