While true

A colleague who was new to Python recently ran into some performance problems with a data processing script he'd written, which we had good reason to think was being held up by slow-running database queries. The script was something like this:

 for id in entities:
      res = run_expensive_query(id) #very slow
      process_result(res) #relatively fast

Other than tweaking the query, the obvious thing to try was to parallelize the for loop, but, since my colleague was new to Python (and programming in general), pointing him to the concurrent package was only going to confuse him. So I thought I'd write a decorator to do the job.

The 'Static' version

This decorator works statically (you'll see what I mean in a moment). It's designed to work with scripts like the one above, where entities is computed once before the for loop runs rather than dynamically changing. The upside is the beginner-friendly syntactic sugar @distribute_over(entities). The way it works is to offload the processing to a ThreadPoolExecutor using executor.map and append the items of the resulting generator to a list result:

from concurrent.futures import ThreadPoolExecutor

 def distribute_over(shards, threads=5):
     """Parallel for loop over an iterable 'shards', which is passed as the first argument of fn"""
     def wrapper(fn):
         def func_wrapper():
             result = []
             with ThreadPoolExecutor(max_workers=threads) as executor:
                 gen = executor.map(fn, shards)
                 for res in gen:
                     result.append(res)
             return result
         return func_wrapper
     return wrapper

Here's an example of usage:

 import time

 @distribute_over(range(10))
 def get_squares(i):
      """Return the ith square and sleeps for a second"""
      time.sleep(1) #Do something time consuming
      return i**2

 print(get_squares())

As you can see, the shards must be determined before the function is even defined, which might be a problem in some cases. But it was perfect for my colleague's use case, since he could just replace his for loop by a function definition and decorate it with @distribute_over(entities), which is just about as readable as it gets.

The 'Dynamic' version

This version is not quite as pretty in terms of syntax but it's more useful because you don't have to know shards ahead of time - you can just pass it in dynamically. Here is what it looks like:

 def distribute(fn, threads=5):
      """Parallel for loop over an iterable shards, which is passed as the first argument of fn"""
      def func_wrapper(shards):
           result = []
           with ThreadPoolExecutor(max_workers=threads) as executor:
                gen = executor.map(fn, shards)
                for res in gen:
                     result.append(res)
           return result
      return func_wrapper

And the example (with cubes this time):

 @distribute
 def get_cubes(i):
      """Return the ith cube and sleeps for a second"""
      time.sleep(1)
      return i**3

 print(get_cubes(range(10)))

Although it's more useful, it's not quite as obvious what's going on in the last line - and this is especially true if get_cubes was called something like analyze_entity, which suggests the argument should be a single entity rather than a list! So aptly naming the method to reflect that it's going to be called on an iterable is important to avoid confusion.

A 'Generator' version

As an afterthought, you could replace the result.append(res) line by yield res to do away with the result variable and end up with a generator instead, which you could call like

 for cube in get_cubes(range(10)):
      print(cube)

Not sure if that's any better (or if it even works, I haven't tried it), but it might satisfy some different use cases.

The code

The code is available from Github in this gist.