items(): support iterating over a key-prefix partition

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)

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 = <uint32_t> 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

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:

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
 

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