Three pieces, all read-only or additive — no shard mutation, no
schema-version bump:
1. shard_for_document(document_root, M) in arborist/document.py.
Pure function: int(document_root[:8], 16) % M. 22 tests cover
determinism, range-bounds, near-uniform distribution (±5pp at
N=20k), and seven lock-in fixtures so peers will disagree
loudly if anyone changes the formula.
2. corpus_shard_count meta field + get/set helpers in store.py.
Lives in the existing key/value meta table; SCHEMA_VERSION
stays at v9.8.0 (the DDL doesn't change and source_root is
layout-independent, so cache records survive a reshard).
Legacy shards (without the field) return None; reshard tool
populates it on every target shard at migration time.
3. Pre-migration snapshot captured to
bench/results/pre-migration-snapshot.json:
docs 3,468,392 (3,468,226 globally unique)
chunks 6,235,764
edges 90,593,537
audit 3,468,403
This is the reference set post-reshard row counts must match.
4. Audit-event extraction script writes all 3.47M events from
all 4 shards to /tmp/audit-events.ndjson (2.0 GB) for the
Option-A canonical-chain consolidation step. Verifies chain
integrity on extract — all 4 source chains report 0 breaks.
5. Fixed a wrong chunk count in docs/corpus-history.md
(had ~3.54M/shard; actual is ~1.56M/shard) and added the
edge-count column (~22.6M/shard, 90.6M total). 6.24M chunks
total, not 14.12M.
Tests: 29 new pass (22 routing + 7 meta). No existing tests
touched.
101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
"""Tests for ``shard_for_document`` (#000065).
|
||
|
||
The routing helper picks a target shard index from a content hash.
|
||
Identical answer on every peer; uniform distribution across shards.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
|
||
import pytest
|
||
|
||
from arborist.document import shard_for_document
|
||
|
||
|
||
def _root(text: str) -> str:
|
||
"""Return a synthetic hex sha256 to use as a document_root."""
|
||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||
|
||
|
||
class TestDeterminism:
|
||
def test_same_input_same_output(self):
|
||
root = _root("the rain in spain")
|
||
assert shard_for_document(root, 4) == shard_for_document(root, 4)
|
||
|
||
def test_independent_of_prior_calls(self):
|
||
a = _root("a"), 4
|
||
b = _root("b"), 4
|
||
first = (shard_for_document(*a), shard_for_document(*b))
|
||
second = (shard_for_document(*a), shard_for_document(*b))
|
||
assert first == second
|
||
|
||
|
||
class TestBounds:
|
||
@pytest.mark.parametrize("M", [1, 2, 4, 8, 16, 64])
|
||
def test_within_range(self, M):
|
||
for i in range(200):
|
||
root = _root(f"doc-{i}")
|
||
idx = shard_for_document(root, M)
|
||
assert 0 <= idx < M
|
||
|
||
def test_M_equals_one_always_zero(self):
|
||
for i in range(50):
|
||
assert shard_for_document(_root(f"doc-{i}"), 1) == 0
|
||
|
||
def test_M_zero_rejected(self):
|
||
with pytest.raises(ValueError):
|
||
shard_for_document(_root("x"), 0)
|
||
|
||
def test_M_negative_rejected(self):
|
||
with pytest.raises(ValueError):
|
||
shard_for_document(_root("x"), -1)
|
||
|
||
|
||
class TestUniformity:
|
||
"""First-32-bit prefix of SHA-256 is uniform; check that distribution
|
||
is near-uniform across shards at meaningful sample sizes."""
|
||
|
||
@pytest.mark.parametrize("M", [2, 4, 8])
|
||
def test_distribution_within_5pct(self, M):
|
||
N = 20_000
|
||
counts = [0] * M
|
||
for i in range(N):
|
||
counts[shard_for_document(_root(f"doc-{i}"), M)] += 1
|
||
expected = N / M
|
||
for i, c in enumerate(counts):
|
||
ratio = c / expected
|
||
assert 0.95 <= ratio <= 1.05, (
|
||
f"shard {i}: {c} hits, expected ~{expected:.0f} "
|
||
f"({ratio:.3f}× — outside 5%)"
|
||
)
|
||
|
||
def test_M_changes_partition(self):
|
||
"""When M doubles, docs can move shard. Confirm at least one
|
||
sampled doc lands in different shards under different M values."""
|
||
roots = [_root(s) for s in ("YouTube", "SQLite", "arborist", "Wikipedia")]
|
||
any_varied = False
|
||
for r in roots:
|
||
seen = {shard_for_document(r, m) for m in (2, 4, 8, 16)}
|
||
if len(seen) >= 2:
|
||
any_varied = True
|
||
break
|
||
assert any_varied, "no sampled root produced different shards under varying M"
|
||
|
||
|
||
class TestKnownRoots:
|
||
"""Lock-in fixtures: if these change, peers will disagree."""
|
||
|
||
@pytest.mark.parametrize(
|
||
"root,M,expected",
|
||
[
|
||
("00000000" + "0" * 56, 4, 0),
|
||
("ffffffff" + "0" * 56, 4, 3),
|
||
("00000004" + "0" * 56, 4, 0),
|
||
("00000003" + "0" * 56, 4, 3),
|
||
("12345678" + "0" * 56, 4, 0),
|
||
("12345679" + "0" * 56, 4, 1),
|
||
("ffffffff" + "0" * 56, 1, 0),
|
||
],
|
||
)
|
||
def test_fixture(self, root, M, expected):
|
||
assert shard_for_document(root, M) == expected
|