Miguel Brito
Posted on September 9, 2021
Introduction
In this post, I’m going to show what you can do to map a function that expects multiple arguments.
There are multiple ways of doing that and they work not only for the regular map
function, you can also use the trick to pass multiple parameters to Executor.map
from the concurrent.futures
and multiprocessing.Pool
.
Table of Contents
-
The Problem With the
map()
Function1.1. Solution 1 - Mapping Multiple Arguments with
itertools.starmap()
1.2. Solution 2 - Using
functools.partial
to “Freeze” the Arguments1.3. Solution 3 - Mapping Multiple Arguments by "Repeating" Them
-
Problem 2: Passing Multiple Parameters to multiprocessing
Pool.map
2.1. Using
pool.starmap
2.2. Using
partial()
-
3.1. Using
partial()
3.2. Using
repeat()
The Problem With the map()
Function
Let’s imagine that you have a function called sum_four
that takes 4 arguments and returns their sum.
>>> def sum_four(a, b, c, d):
return a + b + c + d
Let’s also suppose that you are solving a very specific problem that requires the first 3 arguments to be fixed. In this problem, you want to compare how the function behaves when you vary only the last parameter.
>>> a, b, c = 1, 2, 3
>>> sum_four(a=a, b=b, c=c, d=1)
7
>>> sum_four(a=a, b=b, c=c, d=2)
8
>>> sum_four(a=a, b=b, c=c, d=3)
9
>>> sum_four(a=a, b=b, c=c, d=4)
10
Now, say that you want to use map
, because you like functional programming, or maybe because you come from a language that encourages this paradigm.
Since only d
varies, we could store all potential values for d
we want to test in a list like this all_d_values = [1, 2, 3, 4]
.
The issue is, given a function and a list of single elements, if you want to pass that list to a map
function and it takes only one element, what can you do?
Solution 1 - Mapping Multiple Arguments with itertools.starmap()
The first solution is to not adopt the map
function but use itertools.starmap
instead. This function will take a function as arguments and an iterable of tuples. Then, starmap
will iterate over each tuple t
and call the function by unpacking the arguments, like this for t in tuples: function(*t)
.
To make things more clear, consider the following example.
>>> import itertools
>>> all_d_values = [1, 2, 3, 4]
>>> items = ((a, b, c, d) for d in all_d_values)
>>> list(items)
[(1, 2, 3, 1), (1, 2, 3, 2), (1, 2, 3, 3), (1, 2, 3, 4)]
>>> list(itertools.starmap(sum_four, items))
[7, 8, 9, 10]
As you can see, there’s a lot of repetition, which may inevitably consume a lot of memory if the list is big. To improve that I made items
as a generator, this way we only hold in memory the element we’ll be processing.
Solution 2 - Using functools.partial
to “Freeze” the Arguments
The second solution is to use currying and create a new partial function. According to the docs, partial()
will "freeze" some portion of a function’s arguments and/or keywords resulting in a new function with a simplified signature.
>>> import functools
>>> partial_sum_four = functools.partial(sum_four, a, b, c)
>>> partial_sum_four(3)
9
>>> list(map(partial_sum_four, all_d_values))
[7, 8, 9, 10]
Solution 3 - Mapping Multiple Arguments by "Repeating" Them
The third alternative is to use the itertools.repeat()
.
This function produces an iterator that returns object over and over again. It will run indefinitely if you don’t specify the times argument.
If we take a closer look at map()
's signature, it accepts a function and multiple iterables, map(function, iterable, ...)
.
According to its description,
If additional iterable arguments are passed, function must take that many arguments and is applied to the items from all iterables in parallel. With multiple iterables, the iterator stops when the shortest iterable is exhausted.
Bingo! We can make a
, b
and c
infitnite iterables by using itertools.repeat()
. As soon as all_d_values
is exhausted, which is the shortest iterable, map()
will stop.
>>> import itertools
>>> list(map(sum_four, itertools.repeat(a), itertools.repeat(b), itertools.repeat(c), all_d_values))
[7, 8, 9, 10]
To put it another way, using repeat()
is roughly equivalent to:
>>> list(map(sum_four, [1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], all_d_values))
[7, 8, 9, 10]
You don't need to worry too much about memory as repeat
produces the elements on the go. In fact, it returns a repeatobject
, not list
[ref] .
Problem 2: Passing Multiple Parameters to multiprocessing Pool.map
This problem is very similar to using the regular map()
. The only difference is that we need to pass multiple arguments to the multiprocessing's pool map.
Suppose that we want to speed up our code and run sum_four
in parallel using processes.
The good news is, you can use the solutions above, with one exception: Pool.map
only accepts one iterable. This means we cannot use repeat()
here. Let's see the alternatives.
Using pool.starmap
The Pool
class from multiprocessing
module implements a starmap
function that works the same way as its counterpart from the itertools
module.
>>> from multiprocessing import Pool
>>> import itertools
>>> def sum_four(a, b, c, d):
return a + b + c + d
>>> a, b, c = 1, 2, 3
>>> all_d_values = [1, 2, 3, 4]
>>> items = [(a, b, c, d) for d in all_d_values]
>>> items
[(1, 2, 3, 1), (1, 2, 3, 2), (1, 2, 3, 3), (1, 2, 3, 4)]
>>> with Pool(processes=4) as pool:
res = pool.starmap(sum_four, items)
>>> res
[7, 8, 9, 10]
Using partial()
As alternative, we can also rely on the good partial
function.
>>> import functools
>>> partial_sum_four = functools.partial(sum_four, a, b, c)
>>> with Pool(processes=4) as pool:
res = pool.map(partial_sum_four, all_d_values)
>>> res
[7, 8, 9, 10]
Problem 3: How to Pass Multiple Arguments to concurrent futures Executor.map
?
The concurrent.futures
module provides a high-level interface called Executor
to run callables asynchronously.
There are two different implementations available, a ThreadPoolExecutor
and a ProcessPoolExecutor
.
Contrary to multiprocessing.Pool
, a Executor
does not have a startmap()
function. However, its map()
implementation supports multiple iterables, which allow us to use repeat()
. Another difference is that Executor.map
returns a generator, not a list.
Using partial()
With a ProcessPoolExecutor (or ThreadPoolExecutor)
By "freezing" the arguments using partial
we use the map
method from ProcessPoolExecutor
like a regular map function. Since they both share the same interface, you can do the same interchangeably with a ThreadPoolExecutor
>>> from concurrent.futures import ProcessPoolExecutor
>>> import functools
>>> def sum_four(a, b, c, d):
return a + b + c + d
>>> a, b, c = 1, 2, 3
>>> all_d_values = [1, 2, 3, 4]
>>> partial_sum_four = functools.partial(sum_four, a, b, c)
>>> with ProcessPoolExecutor(max_workers=4) as pool:
res = list(pool.map(partial_sum_four, all_d_values))
>>> res
[7, 8, 9, 10]
Using repeat()
Again, we can just use itertools.repeat
to get the job done like the previous solutions.
>>> from concurrent.futures import ProcessPoolExecutor
>>> from itertools import repeat
>>> def sum_four(a, b, c, d):
return a + b + c + d
>>> a, b, c = 1, 2, 3
>>> all_d_values = [1, 2, 3, 4]
>>> with ProcessPoolExecutor(max_workers=4) as pool:
res = list(pool.map(sum_four, repeat(a), repeat(b), repeat(c), all_d_values))
>>> res
[7, 8, 9, 10]
Conclusion
That’s it for today, folks! I hope you’ve learned something different and useful. The map()
function makes Python feel like a functional programming language. map()
is available not only as a built-in function but also as methods in the multiprocessing
and concurrent.futures
module. In this article, I showed what I do to map functions that take several arguments.
Other posts you may like:
- How to Use datetime.timedelta in Python With Examples
- 73 Examples to Help You Master Python's f-strings
- How to Check if an Exception Is Raised (or Not) With pytest
- 3 Ways to Test API Client Applications in Python
- Everything You Need to Know About Python's Namedtuples
- The Best Way to Compare Two Dictionaries in Python
- 5 Hidden Python Features You Probably Never Heard Of
This post was originally published at https://miguendes.me/how-to-pass-multiple-arguments-to-a-map-function-in-python
See you next time!
Posted on September 9, 2021
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.