Prologue

I didn't give much thought to the computer I used for benchmarking the first time around - it was my daily-driver laptop with an i7-7660 and 32 gigs of DDR4-3200 - but lately I've gotten a bug up my butt about how modern software engineers don't care about their software's performance because they write/test their code on cutting-edge machines that drown their incompetency under an ocean of transistors. [1]

Therefore, this time around I decided to run all my tests on the oldest hardware I had laying around: my trusty Lenovo X201 with a Core i5-560M [2] and 8 gigs of DDR3-1600 RAM - I figured if the code runs fast on this machine it should run fast anywhere. To keep myself honest I decided to benchmark on Python 3.10 instead of 3.11 which has significant performance improvements.

My benchmarking methodology hadn't changed: I reboot before collecting data to mitigate performance impacts from other applications, empty any caches and reset any CPU-internal statistics that may impact timing; I run the benchmark five times, dropping the best/worst runs to mitigate the worst of random variance; I use Scalene while benchmarking to find potential areas for improvement:

for i in {1..5}; do scalene --outfile "run_${i}".json --json --cli --use-virtual-time --profile-all --profile-exclude threading bench.py; done

The benchmarking script is the same as last time:

from pathlib import Path
from time import perf_counter
from zipfile import ZipFile

from py_fast_trie import XFastTrie

def main():
    t = XFastTrie()
    start = perf_counter()

    for i in range(1, 2):
        name = f"primes{i}"
        p = Path(f"~/bench/{name}.zip").expanduser()
        ctr = 0

        with ZipFile(p) as z:
            with z.open(f"{name}.txt") as f:
                print(f"opened {name}.txt")
                for l in f.readlines()[2:]:
                    for n in l.split():
                        t += int(n)
                        ctr += 1

                        if ctr % 100000 == 0:
                            print(f"{ctr}: {perf_counter() - start}")

if __name__ == "__main__": main()

Baseline

Before I got started, I wanted some new baseline numbers to measure my progress:

I knew running the benchmark on a CPU almost 7 years older than before with an older memory architecture at half the speed would impact the numbers, but after so much work to break the 300 second barrier last time seeing an average runtime of 514.251 seconds with no runs faster than 500 seconds did hurt a bit. I turned that pain into fuel for improving my code and dove in.

__setitem__() Helper Function Inlining

I'd planned to focus solely on py-fast-trie this time around when I had a trademark Intrusive Thought - my X-fast trie implementation queries for values in dicts much more than it inserts them but __setitem()__ is still in my benchmark's critical path, shouldn't I take a look at it since I didn't touch it last time? Sure enough, it wasn't hard to find some dead weight:

def __setitem__(self, key: Hashable, value: Any) -> None:
    [...]

    # The index of the key in _keys and its related value in _values
    _, data_idx = self._lookup(key)

    [...]

    nearest_nbr = self._get_open_neighbor(expected_lookup_idx)

Two function calls in the critical path that could easily be inlined - and as I've said before, calling Python functions is costly. Inlining those functions netted a slight performance increase:

As the meme goes, it ain't much - the new average runtime of 510.131 seconds was only a 0.8% decrease - but it's honest work. I wasn't expecting much since dict insertion is such a small part of the critical path, hopefully the other changes I had in mind were more impactful.

Removing TrieNode Boilerplate

At this point I actually turned my attention to py-fast-trie and realized trying to be as Pythonic as possible was hurting more than helping - I made the variables in TrieNode "private" with an underscore, but they still needed to be manipulated from outside so every interaction with data in a TrieNode was mediated by some Java-esque property function:

def _get_left(self) -> TrieNode | None:
    return self._left

[...]

def _set_left(self, new_left: TrieNode | None) -> None:
    self._left = new_left

[...]

left = property(_get_left, _set_left)

I looked over my shoulder to make sure Guido couldn't watch me sin before I made my variables "public" and deleted the property functions; I also deleted some vanity properties that were differently-named wrappers around existing data. Let nobody tell you crime doesn't pay, my Python crimes netted me a tidy speedup:

My benchmark now ran in an average runtime of 475.970 seconds, a 6.7% decrease in runtime from inlining __setitem__() helper functions alone and a 7.4% decrease in runtime from baseline.

insert() Helper Function Inlining

Since my benchmark is focused on insertion, I figured I'd get the biggest return on investment by looking at insert() for the X-fast trie - and I found some more low-hanging fruit at the start of the function. As each leaf should have a reference to its predecessor and successor leaves, insert(), reasonably, calls predecessor() and successor() to find them:

leaf_pred = self.predecessor(value) if self._count > 0 else None
leaf_succ = self.successor(value) if self._count > 0 else None

However, both those functions just call _get_closest_leaf() and go one step in opposite directions - and insert() needs to go both directions anyway. Instead of making two function calls that both end up calling the same function anyway, I make the function call directly and get everything I need at once:

if self._count > 0:
    neighbor = self._get_closest_leaf(value)

    # The closest leaf is the predecessor of the value to be inserted
    if neighbor.value < value:
        leaf_pred = neighbor
        leaf_succ = neighbor.right

    # The closest leaf is the successor of the value to be inserted
    elif neighbor.value > value:
        leaf_succ = neighbor
        leaf_pred = neighbor.left

This resulted in the first big impact of the session:

With an average runtime of 425.106 seconds this last change resulted in a 10.7% decrease in runtime from removing TrieNode boilerplate and a 17.3% decrease in runtime from baseline.

Removing Redundant Input Sanitization

To make sure the library never attempted to operate on invalid input, all the functions where a user would theoretically interact with the library called a sanitization function on input:

def __iadd__(self, value: int | bytes) -> XFastTrie:
    value = self._to_int(value, self._maxlen)
    self.insert(value)
    return self

However, because users could theoretically call insert() directly it also called the sanitization function:

def insert(self, value: int | bytes) -> None:
    value = self._to_int(value, self._maxlen)
    [...]

It didn't make sense to call that function twice so I removed the redundant calls to _to_int() for all the helper functions for a slight speedup:

This change dropped the average runtime to 421.605 seconds, a 0.8% runtime decrease from inlining predecessor/successor leaf discovery and an 18.0% decrease in runtime from baseline - once again yielding a small win for little work.

get_closest_leaf() Value Conversion

I began this optimization session with two primary strategies I gained from the last go-around: inline as many function calls as is reasonable; and eliminate as much bitbanging as possible. Thankfully there wasn't much bitbanging in py-fast-trie to begin with, but there was some in _get_closest_leaf():

direction = value >> (self._maxlen - level - 2) & 1
descendant = ancestor.left if direction == 0 else ancestor.right

Since _get_closest_leaf() was in the critical path for insertion I converted the integer manipulation into string manipulation:

leg = bin(value)[2:].zfill(self._maxlen)[level + 1]
descendant = ancestor.left if leg == "0" else ancestor.right

Given what this strategy did for _get_displaced_neighbors(), I expected a tidy performance boost:

I was rewarded with another quick win of an average runtime...increase...to 435.346 seconds, a 3.3% increase over removing redundant input sanitization but still a 15.3% reduction in runtime from baseline. To be honest these results threw me for a loop and I didn't notice how well this "optimization" performed until I was analyzing the data for this post - the final runtimes were "close enough" to the one before it that I assumed it was another small win; had I noticed at the time I would have undone it, but since the change is in all my subsequent runs I kept it in.

Rewriting insert()

At this point I felt good about decreasing average runtimes by 80 seconds, but I was still a ways from my goal of a sub-300 second average. To make matters worse, since py-fast-trie had less code than py-hopscotch-dict the low-hanging fruit was pretty much gone by now - if I wanted any more gains I was going to have to start working for them. Thankfully I knew there was still one big bundle of suck waiting for me: the trie insertion function. [3]

I'd skimmed the function when I removed predecessor() and successor() calls earlier, but as I read it more thoroughly I realized there was no way to escape a complete rewrite. I didn't want to do it since I knew finding all the little bugs a rewrite would introduce would take longer than rewriting it to begin with; but then I remembered initially writing the function, and I remembered making it add intermediate nodes up from the leaf instead of down from the closest existing ancestor, and I knew what I had to do.

A day or so later Hypothesis said my new version of insert() was good enough, so I declared another small victory against computers. And what did I get for my efforts - was the agita over remembering how to handle descendant pointers worth it?

Yeah, I'd say so. Every run was almost twice as fast at every point as when I stopped trying to optimize last time - on slower hardware! Runtimes decreased to an average of 147.251 seconds, a 66.2% decrease over my bitbanging replacement and a 71.4% decrease from baseline. I said this the last time I was optimizing and I'll say it again here: it felt great to make numbers like this, but it meant I wrote some real dogshit code in the past - at least I'm atoning for my computer crimes now.

Adding __slots__

After the last set of results I felt like I could declare overall victory, but despite all this work to make everything run faster I was still using too damn much memory:

To be honest, I had no idea where it was going - 1 million 64-bit numbers should only take up 8MB of space so even with 20x overhead the trie should use far less than 1GB, plus prior attempts to use Scalene's memory leak detection feature resulted in alerts that put the false in false positive.

I don't remember why the idea came back to me, but I had a thought: "__slots__ is supposed to help with memory and you used it in py-hopscotch-dict without issue, why not try it in py-fast-trie too?"

class TrieNode:
    __slots__ = (
        "leaf",
        "left",
        "parent",
        "right",
        "value",
    )
[...]
class XFastTrie:
    __slots__ = (
        "_count",
        "_level_tables",
        "_max",
        "_maxlen",
        "_min",
        "_root",
    )

I didn't have any better ideas and it was quick to implement, so I added __slots__ on all the classes in py-fast-trie to surprising results:

My new average runtime was 134.280 seconds, an 8.8% decrease from rewriting insert() and a 73.9% decrease from baseline; but more importantly I managed to reduce memory usage by over 300MB from baseline!

get_closest_leaf() Value Deconversion

Everybody makes mistakes, and I definitely made one when I thought getting rid of bitbanging would speed up _get_closest_leaf() - but to be fair that has been the truth basically every other time I've tried it. Anyway, reverting the change did what I expected:

Average runtimes decreased slightly to 133.763 seconds, 0.4% from adding __slots__ and 74.0% overall. I thought the speedup would be greater since moving away from bitbanging led to a 3.3% increase, but I assume rewriting insert() may have decreased the benefit of _get_closest_leaf() being faster.

__contains__() Helper Function Inlining

As I was rewriting insert(), it finally clicked the first line calls __contains__() in py-hopscotch-dict but that function is still essentially boilerplate:

def __contains__(self, key: Hashable) -> bool:
    _, idx = self._lookup(key)
    return idx is not None

I decided to return to py-hopscotch-dict and inline the _lookup() call into __contains__() like I'd done for most of the other dict access functions:

def __contains__(self, key: Hashable) -> bool:
    result = False
    home_idx = abs(hash(key)) % self._size

    nbhd = self._nbhds[home_idx]
    for nbhd_idx, nbr in enumerate(bin(nbhd)[:1:-1]):
        if nbr == "1":
            neighbor = (home_idx + nbhd_idx) % self._size
            data_idx = self._indices[neighbor]

            if data_idx < 0:
                raise RuntimeError(
                    (
                        "Index {} has supposed displaced neighbor that points to "
                        "free index"
                    ).format(home_idx)
                )

            if self._keys[data_idx] == key:
                result = True
                break

    return result

The benefit was less than expected, but still welcome:

Runtimes now averaged 131.559 seconds, a 1.6% decrease over returning to bitbanging in _get_closest_leaf() and a 74.4% decrease in runtime overall.

__resize__() Refactoring

The Scalene reports suggested I was still spending a surprising amount of time in __resize__(), so I decided to see what performance improvements were possible there.

for data_idx, key in enumerate(self._keys):
    expected_lookup_idx = abs(hash(key)) % new_size

    open_nbr_idx = self._get_open_neighbor(expected_lookup_idx)
    if open_nbr_idx is None:
        if new_size - expected_lookup_idx < self._nbhd_size:
            free_up_idx = 0
        else:
            free_up_idx = expected_lookup_idx

        self._free_up(free_up_idx)
        open_nbr_idx = cast(int, self._get_open_neighbor(expected_lookup_idx))
    nbhd_idx = (open_nbr_idx - expected_lookup_idx) % new_size
    self._set_neighbor(expected_lookup_idx, nbhd_idx)
    self._set_lookup_index_info(open_nbr_idx, data=data_idx)

I couldn't help thinking eating two calls to _get_open_neighbor() every time a neighborhood was too full was a bit ridiculous, so I inlined the functionality as best I could:

for data_idx, key in enumerate(self._keys):
    found_nbr = False
    home_idx = abs(hash(key)) % new_size

    for offset in range(self._nbhd_size):
        nbr = (home_idx + offset) % self._size
        if self._indices[nbr] == self.FREE_ENTRY:
            self._set_neighbor(home_idx, offset)
            self._set_lookup_index_info(nbr, data=data_idx)
            found_nbr = True
            break

    if not found_nbr:
        free_up_idx = 0 if new_size - home_idx < self._nbhd_size else home_idx
        self._free_up(free_up_idx)

        for offset in range(self._nbhd_size):
            nbr = (home_idx + offset) % self._size
            if self._indices[nbr] == self.FREE_ENTRY:
                self._set_neighbor(home_idx, offset)
                self._set_lookup_index_info(nbr, data=data_idx)
                break

I still don't like writing essentially the same loop twice, but I can't argue with the results:

I now had an average runtime of 128.445 seconds, a 2.4% decrease from inlining helper functions in __contains__() and a 75.0% decrease from baseline.

Neighborhood Management Inlining

While I was refactoring _resize() I took a look at some of the other helper functions to remind myself how much "help" they actually provide:

def _set_neighbor(self, lookup_idx: int, nbhd_idx: int) -> None:
    if lookup_idx < 0 or nbhd_idx < 0:
        raise ValueError("Indexes cannot be negative")
    elif lookup_idx >= self._size:
        raise ValueError(f"Index {lookup_idx} outside array")
    elif nbhd_idx >= self._nbhd_size:
        raise ValueError("Trying to set neighbor outside neighborhood")

    self._nbhds[lookup_idx] |= 1 << nbhd_idx

def _clear_neighbor(self, lookup_idx: int, nbhd_idx: int) -> None:
    if lookup_idx < 0 or nbhd_idx < 0:
        raise ValueError("Indexes cannot be negative")
    elif lookup_idx >= self._size:
        raise ValueError(f"Index {lookup_idx} outside array")
    elif nbhd_idx >= self._nbhd_size:
        raise ValueError("Trying to clear neighbor outside neighborhood")

    self._nbhds[lookup_idx] &= ~(1 << nbhd_idx)

I know I originally wrote my code like this to "be more professional" and "avoid DRY", but eating a function call in a hot loop for one line of bitbanging is ridiculous and I would stand for it no longer. Even though I'm still not a fan of bitbanging I couldn't think of a better way that was still idempotent so I inlined the code as it was:

for data_idx, key in enumerate(self._keys):
    found_nbr = False
    home_idx = abs(hash(key)) % new_size

    for offset in range(self._nbhd_size):
        nbr = (home_idx + offset) % self._size
        if self._indices[nbr] == self.FREE_ENTRY:
            # Mark the neighbor storing the entry index as occupied
            self._nbhds[home_idx] |= 1 << offset
            self._set_lookup_index_info(nbr, data=data_idx)
            found_nbr = True
            break

    if not found_nbr:
        free_up_idx = 0 if new_size - home_idx < self._nbhd_size else home_idx
        self._free_up(free_up_idx)

        for offset in range(self._nbhd_size):
            nbr = (home_idx + offset) % self._size
            if self._indices[nbr] == self.FREE_ENTRY:
                # Mark the neighbor storing the entry index as occupied
                self._nbhds[home_idx] |= 1 << offset
                self._set_lookup_index_info(nbr, data=data_idx)
                break

At least one of my heuristics didn't let me down:

The average runtime decreased to 123.803 seconds, a 3.6% runtime decrease from refactoring _resize() and a 75.9% decrease from baseline.

Index Information Inlining

The keen-eyed reader may have noticed a second helper function in the snippet of _resize():

def _set_lookup_index_info(
    self, lookup_idx: int, data: Optional[int] = None, nbhd: Optional[int] = None
) -> None:
    if lookup_idx < 0:
        raise ValueError("Indexes cannot be negative")
    elif lookup_idx >= self._size:
        raise ValueError(f"Index {lookup_idx} outside array")

    if data is not None:
        self._indices[lookup_idx] = data

    if nbhd is not None:
        self._nbhds[lookup_idx] = nbhd

As it turned out I built a "two-line" helper function, but I only used one of the legs so it ended up being a one-line helper function as well. Obviously I got rid of it:

for data_idx, key in enumerate(self._keys):
    found_nbr = False
    home_idx = abs(hash(key)) % new_size

    for offset in range(self._nbhd_size):
        nbr = (home_idx + offset) % self._size
        if self._indices[nbr] == self.FREE_ENTRY:
            # Mark the neighbor storing the entry index as occupied
            self._nbhds[home_idx] |= 1 << offset
            self._indices[nbr] = data_idx
            found_nbr = True
            break

    if not found_nbr:
        free_up_idx = 0 if new_size - home_idx < self._nbhd_size else home_idx
        self._free_up(free_up_idx)

        for offset in range(self._nbhd_size):
            nbr = (home_idx + offset) % self._size
            if self._indices[nbr] == self.FREE_ENTRY:
                # Mark the neighbor storing the entry index as occupied
                self._nbhds[home_idx] |= 1 << offset
                self._indices[nbr] = data_idx
                break

The benefit here was also obvious:

At 117.979 seconds, the average benchmark runtime decreased 4.7% from inlining neighborhood management and 77.1% from baseline.

Wrapping Up

To be honest, I thought about stopping after I put bitbanging back in _get_closest_leaf() - I wanted to see a sub-120 second run, but had no idea how to shave off an additional 15 seconds - but I'm glad I stuck with it. Now my last goal is to insert the first 100k primes in under 10 seconds on this hardware - but I've put too much time into this already. I'm never happy realizing I write such unperformant garbage code but I am satisfied with the results I get when I sit down and try to improve it. I decided to play a data scientist and put all my data into a spreadsheet so I could do some math on it:

I'm glad I managed to ultimately reduce the standard deviation of my benchmark runs by roughly a factor of three, but I don't what I changed from when I cut it by a factor of 70.

Memory usage declined from 1109 MiB to 761 MiB, a 31% decrease pretty much entirely from using __slots__ on TrieNode. I didn't realize how many TrieNode objects I make in py-fast-trie; since there's so many of them a small savings goes a long way, but since there's basically nothing in them I don't have many options to make them smaller.

Just for funsies, I ran my benchmark script without Scalene on my homelab server with a Ryzen 5 5650 and DDR4-3200 memory: