I have several files that have identifying strings that need to be replaced. There are >300,000 of these strings, and each string may appear more than once in any given file, or not at all. The files that need to be fixed can be several GB, comprising hundreds of millions of lines.

I first wrote a function that is rough but works, and then I parallelized it with Parallel Python by splitting each file into partitions that can be individually searched. On my test file, the function write a corrected file at about 30KB/min/processor. This is not going to be tenable for the larger files that I need to replace. Here's what I have so far:

#!/usr/bin/python
import sys
from sys import argv 
import collections
import os
import re
import subprocess
import pp

script, nproc, filename1, filename2, filename3 = argv 


def count_lines(file):
    # print "Counting lines..."
    out = subprocess.Popen(['wc', '-l', file], stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate()[0]
    line_count = int(out.split(' ')[1])
    #print(line_count)
    print line_count, "lines found."
    return line_count

def splitFile(file, line_count):
    print "Partitioning files..."
    with open(file) as f:
        partition = line_count/4
        f = f.readlines()
        p1 = f[0:partition]
        # print(p1)
        p2 = f[partition + 1: partition * 2]
        # print(p2)
        p3 = f[partition * 2 + 1: partition * 3]
        # print(p3)
        p4 = f[partition * 3 + 1: line_count]
        # print(p4)
        partitions = (p1, p2, p3, p4)
        print len(partitions), "partitions created."
        return partitions

def headerreplace(partition, file2, file3):
    # initialize dictionary
    d = {}
    with open(file2, 'w') as corrected, open(file3) as f:
        #create dictionary
        for line in f:
            line = line.rstrip()
            (key, val) = line.split(" ", 1)
            d[key] = val
        # parse original and print replacement to corrected
        count = 1
        for line in partition:
            #print line
            for key in d:
                if count < len(d):
                #print key
                    if key in line:
                        new_line = line.replace(key, d[key])
                        #print d[key]
                        corrected.write(new_line)
                        # print new_line
                        count = len(d)
                        break
                    else:
                        count += 1
                elif count == len(d):
                    corrected.write(line)
                    count = 0
                    # print line
        return "Finished job"



# tuple of all parallel python servers to connect with
ppservers = ()

if len(sys.argv) > 1:
    ncpus = int(sys.argv[1])
    # Creates jobserver with ncpus workers
    job_server = pp.Server(ncpus, ppservers=ppservers)
else:
    # Creates jobserver with automatically detected number of workers
    job_server = pp.Server(ppservers=ppservers)

print "Starting pp with", job_server.get_ncpus(), "workers"


def master():
    jobs = []
    line_count = count_lines(filename1)
    partitions = splitFile(filename1, line_count)
    corrected = ("corrected1.txt", "corrected2.txt", "corrected3.txt", "corrected4.txt")
    for partition, outfile in zip(partitions, corrected):
        jobs.append(job_server.submit(headerreplace, (partition, outfile, filename3)))
    for job in jobs:
        result = job()
        if result:
            print "Completed job"

master()

This code is rough, I know. I also haven't done Python optimization because the jobs I need done are usually short enough that it doesn't really matter (I'm a geneticist, not a programmer). I'm guessing the nested loops have a heavy cost, but I'm not sure if there is a way to improve upon that. I would greatly appreciate any pointers to speed this up.

====EDIT====

The Aho–Corasick algorithm has worked beautifully, and I've confirmed that the function I've written does what I want it to do. However, I'd still like to maintain parallelization with pp, but I've seem to have broken it.

#!/usr/bin/python
import sys
from sys import argv 
import collections
import os
import re
import subprocess
import pp
import ahocorasick

script, nproc, original, corrected, dict = argv 

def count_lines(filename):
    out = subprocess.Popen(['wc', '-l', filename], stdout=subprocess.PIPE, stderr=subprocess.STDOUT).communicate()[0]
    line_count = int(out.split(' ')[1])
    #print(line_count)
    print line_count, "lines found."
    return line_count

def splitFile(filename, line_count, nproc):
    print "Partitioning files..."
    with open(filename) as f:
        partition_size = line_count/nproc
        f = f.readlines()
        partitions = []
        count = 0
        for i in range(nproc):
            i = f[(partition_size + 1) * count:partition_size * (count + 1)]
            partitions.append(i)
            count += 1
        print len(partitions), "partitions created."
        return partitions

def makeAutomaton(filename):
    """Build an Aho-Corasick automaton from a dictionary file and return
    it. The lines in the dictionary file must consist of a key and a
    value separated by a space.
    """
    print "Making automaton..."
    automaton = ahocorasick.Automaton()
    with open(filename) as f:
        for line in f:
            key, value = line.rstrip().split(" ", 1)
            automaton.add_word(key, (key, value))
    automaton.make_automaton()
    return automaton

def headerreplace(automaton, partition, corrected):
    """Apply an Aho-Corasick automaton to an input file, replacing the
    first occurrence of a key in each line with the corresponding
    value, and writing the result to the output file."""
    with open(partition, 'r') as infile, open(corrected, 'w') as outfile:
        for line in infile:
            for end, (key, value) in automaton.iter(line):
                line = line[:end - len(key) + 1] + value + line[end + 1:]
                break # At most one replacement per line
            outfile.write(line)



# tuple of all parallel python servers to connect with
ppservers = ()

if len(sys.argv) > 1:
    ncpus = int(sys.argv[1])
    # Creates jobserver with ncpus workers
    job_server = pp.Server(ncpus, ppservers=ppservers)
else:
    # Creates jobserver with automatically detected number of workers
    job_server = pp.Server(ppservers=ppservers)

print "Starting pp with", job_server.get_ncpus(), "workers"


def master():
    nproc = int(sys.argv[1])
    jobs = []
    line_count = count_lines(original)
    partitions = splitFile(original, line_count, nproc)
    automaton = makeAutomaton(dict)
    # headerreplace(automaton, original, corrected)
    corrected_list = []
    for i in range(nproc):
        filename = str("%s%i.txt" % (corrected, i))
        corrected_list.append(filename)
    for infile, outfile in zip(partitions, corrected_list):
        jobs.append(job_server.submit(headerreplace, (automaton, infile, outfile)))
    for job in jobs:
        result = job()
        if result:
            print "Completed job"

master()

This results in the error TypeError: coercing to Unicode: need string or buffer, list found. Apparently it doesn't like that I'm iterating through a lists to submit separate jobs. If I edit jobs.append(job_server.submit(headerreplace, (automaton, infile, outfile))) to be jobs.append(job_server.submit(headerreplace, (automaton, original, corrected))) (i.e. the arguments), then the program runs splendidly (albeit, 4 times). What am I missing?

up vote 2 down vote accepted

The problem is in these loops:

for line in partition:
    for key in d:
        if key in line:

Every key is separately searched for in every line, and so if there are \$m\$ keys and \$n\$ lines, the overall runtime will be \$Ω(mn)\$.

Instead, you should search for all keys simultaneously, using the Aho–Corasick algorithm. Python doesn't come with an implementation of this algorithm, but you can use the pyahocorasick package. This will reduce the runtime to \$O(m + n)\$.

import ahocorasick

def make_automaton(filename):
    """Build an Aho-Corasick automaton from a dictionary file and return
    it. The lines in the dictionary file must consist of a key and a
    value separated by a space.

    """
    automaton = ahocorasick.Automaton()
    with open(filename) as f:
        for line in f:
            key, value = line.rstrip().split(" ", 1)
            automaton.add_word(key, (key, value))
    automaton.make_automaton()
    return automaton

def apply_automaton(automaton, input_filename, output_filename):
    """Apply an Aho-Corasick automaton to an input file, replacing the
    first occurrence of a key in each line with the corresponding
    value, and writing the result to the output file.

    """
    with open(input_filename) as infile, open(output_filename, 'w') as outfile:
        for line in infile:
            for end, (key, value) in automaton.iter(line):
                line = line[:end - len(key) + 1] + value + line[end + 1:]
                break # At most one replacement per line.
            outfile.write(line)

Update: With this change to the algorithm, it no longer makes sense to attempt to parallelize the processing of the file. The reason is that with an efficient search algorithm, the program is I/O bound, not CPU bound. Carrying out the replacement thus takes no longer than partitioning the file.

  • Thanks for this, it works well. I've managed to implement the Aho–Corasick algorithm, but it has broken my parallelization - see the edit above. – Nic Jul 24 '17 at 20:50
  • @Nic: see updated answer. – Gareth Rees Jul 25 '17 at 11:01
  • Ah, I see. Thank you. – Nic Jul 25 '17 at 12:46

Your Answer

By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.

Not the answer you're looking for? Browse other questions tagged or ask your own question.