14 December 2020

# A Faster Partition Function in Python

by Gianni Tedesco

Some pythonistas wonder what is the fastest way to write a function to partition a sequence of items in to two lists based on some predicate. One could, to be sure, filter the list and then filterfalse the list. But in this case the predicate will evaluated twice for each item - an inefficiency that some would balk at. That approach also doesn’t seem very parsimonious.

This problem has inspired several broad classes of solutions based on different needs and priorities and some very minor controversy about which is the most efficient. I hope to clear up the question of efficiency by introducing a new variant here today.

# An Itch You Just Can’t Scratch

Have you ever found yourself doing something like this:

foos = [x for x in input if x.type in known_foos]
bars = [x for x in input if x.type not in known_foos]


It’s a DRY itch. Quite irritating. It calls for a salve. Surely Python would have a solution for this? It has an expansive standard library…

# Could I Even Give Two Forks?

The suggestion in the itertools documentation is to use tee in an elegant way, along with filter and filterfalse, which are also both in the standard libraries. This approach also has the benefit of working with arbitrary generators. It can even, in principle, handle an infinitely long input sequence.

It’s quite clever, so it’s worth mentioning.

from itertools import tee, filterfalse

def partition(pred, it): # stdlib
a, b = tee(it)
return filterfalse(pred, a), filter(pred, b)


But the problem is that it’s also quite slow when you’re just going to be materialising the results immediately. To see why, imagine calling it:

hay, needles = partition(lambda item: item in needle_set, haystack)
return Results(list(hay), list(needles))


As soon as we materialise one of the returned generators, tee will internally materialise haystack in to a temporary list. And the predicates are still being evaluated twice, once for each fork of the tee.

But I can see this sort of thing being useful in a graph of asyncio tasks or something. Maybe one day I will find myself in a situation where I would need to use a cool technique like this.

# The Clever Trevors

But let’s forget about possibly-infinite generators and lazy evaluation and just consider the case of a list that we want to partition in to two new lists.

For that case case, I’ve seen a couple of clever approaches.

The first one is quite quick to type and nice and readable. It involves using the result of the predicate function to index in to the results tuple. It takes advantage of the fact that True and False cast to 1 and 0 when used to index a tuple.

def partition(pred, it): # trevor
ret = ([], [])
for item in it:
ret[pred(item)].append(item)
return ret


Honestly, this probably would have been my go-to approach, but it actually turns out to be quite slow.

The second is a very functional-style solution, straight out of some sort of Haskellers PhD thesis, proposed by stackoverflow member Mariy. This one uses functools.reduce, which is pythons version of foldl:

def partition(pred, it): # mariy
return reduce(lambda x, y: x[pred(y)].append(y) or x, it, ([], []))


In princple, this is the same as the above but it takes advantage of the fact that reduce performs some of the boilerplate for you.

I expect it to be a bit less efficient since it creates as many result tuples as there are elements in the input list. Repetetive object creation and destruction is a common source of overhead in python.

# The Plain Jane

According to gboffi, it turns out the fastest solution that stackoverflow came up with is just the straight forward version, credited to Mark Byers:

def partition(pred, it): # byers
ts = []
fs = []
for item in it:
if pred(item):
ts.append(item)
else:
fs.append(item)
return fs, ts


But, knowing something about optimizing python code, I can see a clear opportunity to speed this up.

# Don’t dis My Code, Man

We can disassemble this to python bytecode with dis. Generally speaking, the main determinant of performance in these kind of loops is the number of bytecode instructions which need to be dispatched per iteration.

So here is what the Byers version compiles to in cpython 3.9:

  2           0 BUILD_LIST               0
2 STORE_FAST               2 (ts)

3           4 BUILD_LIST               0
6 STORE_FAST               3 (fs)

10 GET_ITER
>>   12 FOR_ITER                34 (to 48)
14 STORE_FAST               4 (item)

20 CALL_FUNCTION            1
22 POP_JUMP_IF_FALSE       36

30 CALL_METHOD              1
32 POP_TOP
34 JUMP_ABSOLUTE           12

8     >>   36 LOAD_FAST                1 (fs)
42 CALL_METHOD              1
44 POP_TOP
46 JUMP_ABSOLUTE           12

9     >>   48 LOAD_FAST                3 (fs)
52 BUILD_TUPLE              2
54 RETURN_VALUE


The inner loop is between locations 12 and 46. It’s 18 instructions long.

We can see that the append code has some duplicated instructions in each branch (24-34, and 36-46), but since only one branch is taken at a time, that shouldn’t really be counted.

So if we measure the actual instruction path-length per iteration, it comes out to 12 instructions.

More importantly, we can see that inside that inner loop we’re doing a lookup of the append method, which is actually looking up a string in a dictionary. For this reason, the LOAD_GLOBAL and LOAD_METHOD calls are among the slowest of the basic instruction types in the python machine (not including the instructions which call out to python subroutines, of course).

# Python is Dynamic, and There is no Escape(-Analysis)

We have probably all heard that python is a dynamic language. But we may not know what all of the consequences of that fact are. For one thing, it means that attribute lookups are often expensive dictionary lookups. But it also means that even some very obvious optimisations simply cannot be made. For example, the python compiler cannot, in general, optimise:

# Program A
x = []
for a in b:
x.append(a)


in to:

# Program B
x = []
f = x.append
for a in b:
f(a)


The reason for this is that the method list.append is absolutely free to assign a different value to self.append. This means that programs A and B are (or at least, may be) semantically different. Therefore B is not a valid optimization of A.

Now you might think that, well, the bytecode compiler can know about the implementation of list.append because it’s built-in. Which is true. But it doesn’t stop any other code, for example in b.__iter__.next() from obtaining a reference to x via any number of mechanisms and then altering its attribute dict.

You wouldn’t even have to resort to reflection - the local variable x could just be a reference to an object which is also reachable in a global scope. And what variables there are in the global scope can be modified on the fly by any code so it wouldn’t be enough to do an escape-analysis at compile time.

Frankly there are just too many avenues in python for this sort of jiggery-pokery to take place. And these avenues are kind of the point of python.

# Locals are Fast In Python

OK, that’s interesting, but what if we just explicitly apply the above optimisation?

def partition(pred, it): # scara
ts = []
fs = []
t = ts.append
f = fs.append
for item in it:
if pred(item):
t(item)
else:
f(item)
return fs, ts


You can see from the bytecode that this has reduced the instruction path-length as we expected:

 28           0 BUILD_LIST               0
2 STORE_FAST               2 (ts)

29           4 BUILD_LIST               0
6 STORE_FAST               3 (fs)

12 STORE_FAST               4 (t)

18 STORE_FAST               5 (f)

22 GET_ITER
>>   24 FOR_ITER                30 (to 56)
26 STORE_FAST               6 (item)

32 CALL_FUNCTION            1
34 POP_JUMP_IF_FALSE       46

40 CALL_FUNCTION            1
42 POP_TOP
44 JUMP_ABSOLUTE           24

36     >>   46 LOAD_FAST                5 (f)
50 CALL_FUNCTION            1
52 POP_TOP
54 JUMP_ABSOLUTE           24

37     >>   56 LOAD_FAST                3 (fs)
60 BUILD_TUPLE              2
62 RETURN_VALUE


The inner loop takes place between locations 24 and 48. It’s 16 instructions long. And if we look at the instruction path length for the loop, it’s now 11 instructions long. One shorter than before. Of course, it’s that missing LOAD_ATTR instruction which has been hoisted out of the loop in to the function preamble.

# A Bit More Squeezing

I couldn’t figure out a way to get the inner loop any tighter but I did find a way to make the code more compact by removing duplicated code on both sides of the branch. It led to a tiny improvement in performance. Here it is:

def partition(pred, it): # scara2
ts = []
fs = []
t = ts.append
f = fs.append
for item in it:
(t if pred(item) else f)(item)
return fs, ts

 40           0 BUILD_LIST               0
2 STORE_FAST               2 (ts)

41           4 BUILD_LIST               0
6 STORE_FAST               3 (fs)

12 STORE_FAST               4 (t)

18 STORE_FAST               5 (f)

22 GET_ITER
>>   24 FOR_ITER                24 (to 50)
26 STORE_FAST               6 (item)

32 CALL_FUNCTION            1
34 POP_JUMP_IF_FALSE       40
38 JUMP_FORWARD             2 (to 42)
44 CALL_FUNCTION            1
46 POP_TOP
48 JUMP_ABSOLUTE           24

46     >>   50 LOAD_FAST                3 (fs)
54 BUILD_TUPLE              2
56 RETURN_VALUE


Now we’re down to a 13 instruction inner loop with an 11 instruction path-length.

# Benchmarks

For benchmarking I loaded a dictionary of some half a million words, created a set of all the words that begin with “s”, and then partitioned the dictionary based on membership of the “s” set.

Tests were performed on an i7-6600U laptop with python 3.9.0 and PYTHONHASHSEED=0

dic = tuple(Path('/usr/share/dict/words').read_text().splitlines())
needles = frozenset((word for word in dic if word.startswith('s')))
for func_name in ('stdlib', 'mariy', 'trevor', 'byers', 'scara', 'scara2'):
print(func_name, timeit(
f'{func_name}(pred, dic)',
setup=f'from __main__ import dic, pred, {func_name}',
number=32,
))

Variant Time for 32 iterations (seconds)
stdlib 3.5664224450010806
mariy 3.532839963911101
trevor 2.5254638858605176
byers 2.356995061971247
scara 2.1279552809428424
scara2 2.087127824081108

# A Parting Shot

Here it is with all the mypy --strict typing goodness added:

from typing import List, Any, Callable, Iterable, TypeVar, Tuple

T = TypeVar('T')

def partition(pred: Callable[[T], bool], it: Iterable[T]) \
-> Tuple[List[T], List[T]]:
ts: List[T] = []
fs: List[T] = []
t = ts.append
f = fs.append
for item in it:
(t if pred(item) else f)(item)
return fs, ts

tags: performance - python