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?
