From d97a70d90c20c4f15bb8ba58f5d8391f97793879 Mon Sep 17 00:00:00 2001 From: Thomas Waldmann Date: Tue, 11 Aug 2026 15:14:08 +0200 Subject: [PATCH] items(): support iterating over a key-prefix partition, fixes #49 Add optional prefix_bits / prefix arguments to HashTable.items() and HashTableNT.items() to iterate only over the items whose key starts with the given bit prefix (first 4 key bytes interpreted big-endian, like _get_index does, so a partition is a contiguous range of the sorted key space). Non-matching entries are skipped at C level without creating Python objects for them, so scanning the table once per partition is cheap. As the keys are random bytes, this partitions the items into 2 ** prefix_bits roughly equally sized, disjoint sets, e.g. to process a huge hash table in batches with a small memory footprint (see borgbackup/borg#9886). Co-Authored-By: Claude Fable 5 --- CHANGES.rst | 7 +++++- src/borghash/HashTable.pyx | 28 +++++++++++++++++++++++- src/borghash/HashTableNT.pyx | 9 ++++++-- tests/hashtable_test.py | 42 ++++++++++++++++++++++++++++++++++++ tests/hashtablent_test.py | 14 ++++++++++++ 5 files changed, 96 insertions(+), 4 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 2fdbf5f..8fb649a 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,9 +1,14 @@ Changelog ========= -Version 0.1.2 (unreleased) +Version 0.2.0 (unreleased) -------------------------- +- ``HashTable.items()`` / ``HashTableNT.items()``: add optional ``prefix_bits`` / + ``prefix`` arguments to iterate only over the items whose key starts with the + given bit prefix. As the keys are random bytes, this partitions the items into + ``2 ** prefix_bits`` roughly equally sized, disjoint sets, e.g. to process a + huge hash table in batches with a small memory footprint, #49. - Require ``key_size >= 4`` to avoid out-of-bounds reads in ``_get_index``, #42. Version 0.1.1 (2026-02-09) diff --git a/src/borghash/HashTable.pyx b/src/borghash/HashTable.pyx index 67284e9..1fcd1a8 100644 --- a/src/borghash/HashTable.pyx +++ b/src/borghash/HashTable.pyx @@ -226,13 +226,39 @@ cdef class HashTable: del self[key] return value - def items(self) -> Iterator[tuple[bytes, bytes]]: + def items(self, *, prefix_bits: int = 0, prefix: int = 0) -> Iterator[tuple[bytes, bytes]]: + """ + Iterate over items, optionally only over items whose key starts with the given bit prefix. + + prefix_bits: number of leading key bits to compare, 0..32 (0 means: no filtering). + prefix: expected value of these leading key bits, given in the integer's low bits, + thus: 0 <= prefix < 2 ** prefix_bits. + + As the keys are expected to be random bytes, this partitions the keys into + 2 ** prefix_bits roughly equally sized, disjoint sets, while only creating + bytes objects for the matching keys/values. As the prefix compares the keys' + leading bits, each set is a contiguous range of the sorted key space. + """ + if not 0 <= prefix_bits <= 32: + raise ValueError("prefix_bits must be in range 0..32.") + if not 0 <= prefix < (1 << prefix_bits): + raise ValueError("prefix must be in range 0..(2 ** prefix_bits - 1).") cdef size_t i cdef uint32_t kv_index + cdef uint8_t* key_ptr + cdef uint32_t key32 + cdef int shift = 32 - prefix_bits + cdef uint32_t wanted = prefix + cdef bint filtering = prefix_bits > 0 self.stats_iter += 1 for i in range(self.capacity): kv_index = self.table[i] if kv_index not in (FREE_BUCKET, TOMBSTONE_BUCKET): + if filtering: + key_ptr = self.keys + kv_index * self.ksize + key32 = (key_ptr[0] << 24) | (key_ptr[1] << 16) | (key_ptr[2] << 8) | key_ptr[3] + if (key32 >> shift) != wanted: + continue key = self.keys[kv_index * self.ksize:(kv_index + 1) * self.ksize] value = self.values[kv_index * self.vsize:(kv_index + 1) * self.vsize] yield key, value diff --git a/src/borghash/HashTableNT.pyx b/src/borghash/HashTableNT.pyx index 02b0e92..289416c 100644 --- a/src/borghash/HashTableNT.pyx +++ b/src/borghash/HashTableNT.pyx @@ -90,8 +90,13 @@ cdef class HashTableNT: self._check_key(key) return key in self.inner - def items(self) -> Iterator[tuple[bytes, Any]]: - for key, binary_value in self.inner.items(): + def items(self, *, prefix_bits: int = 0, prefix: int = 0) -> Iterator[tuple[bytes, Any]]: + """ + Iterate over items, optionally only over items whose key starts with the given bit prefix. + + See HashTable.items for the prefix_bits / prefix semantics. + """ + for key, binary_value in self.inner.items(prefix_bits=prefix_bits, prefix=prefix): yield (key, self._to_namedtuple_value(binary_value)) def __len__(self) -> int: diff --git a/tests/hashtable_test.py b/tests/hashtable_test.py index 4622a43..7cf6ce4 100644 --- a/tests/hashtable_test.py +++ b/tests/hashtable_test.py @@ -68,6 +68,48 @@ def test_items(ht12): assert (key2, value2) in items +@pytest.mark.parametrize("prefix_bits", [1, 2, 3, 8, 9, 32]) +def test_items_prefix(ht, prefix_bits): + # pseudo-random keys, so (for small prefix_bits) all partitions should get some keys + expected = {} + for x in range(1000): + key = H2(x) + value = key[:4] + ht[key] = value + prefix = int.from_bytes(key[:4], "big") >> (32 - prefix_bits) + expected.setdefault(prefix, set()).add((key, value)) + # only iterate over the actually occupied partitions (2 ** 32 would take a while)... + collected = [] + for prefix in expected: + items = set(ht.items(prefix_bits=prefix_bits, prefix=prefix)) + assert items == expected[prefix] + collected.extend(items) + # together, the occupied partitions have everything, exactly once: + assert len(collected) == len(set(collected)) == 1000 + # ... but an unoccupied partition (if any) must yield nothing: + unoccupied = next((p for p in range(2 ** prefix_bits) if p not in expected), None) + if unoccupied is not None: + assert list(ht.items(prefix_bits=prefix_bits, prefix=unoccupied)) == [] + + +def test_items_prefix_zero_bits(ht12): + # prefix_bits=0 means: no filtering + assert set(ht12.items(prefix_bits=0, prefix=0)) == set(ht12.items()) + + +def test_items_prefix_validation(ht12): + with pytest.raises(ValueError): + list(ht12.items(prefix_bits=-1)) + with pytest.raises(ValueError): + list(ht12.items(prefix_bits=33)) + with pytest.raises(ValueError): + list(ht12.items(prefix_bits=0, prefix=1)) + with pytest.raises(ValueError): + list(ht12.items(prefix_bits=2, prefix=4)) + with pytest.raises(ValueError): + list(ht12.items(prefix_bits=2, prefix=-1)) + + def test_len(ht12): assert len(ht12) == 2 diff --git a/tests/hashtablent_test.py b/tests/hashtablent_test.py index cad2498..38720ac 100644 --- a/tests/hashtablent_test.py +++ b/tests/hashtablent_test.py @@ -62,6 +62,20 @@ def test_items(ntht12): assert (key2, value2) in items +def test_items_prefix(ntht): + prefix_bits = 4 + expected = {} + for x in range(100): + key = H2(x) + value = value_type(x, x + 1, x + 2) + ntht[key] = value + prefix = key[0] >> (8 - prefix_bits) + expected.setdefault(prefix, set()).add((key, value)) + for prefix in range(2 ** prefix_bits): + items = set(ntht.items(prefix_bits=prefix_bits, prefix=prefix)) + assert items == expected.get(prefix, set()) + + def test_len(ntht12): assert len(ntht12) == 2