Sample and partition the pooled corpus at passage level

`title` is the empty string for every cloud and fiqa passage, so two of the
four domains have exactly one title covering 72,442 and 61,022 passages, and
govt's titles are web-scrape artifacts with one 10,192-passage bucket. Only
clapnq has titles that identify a document.

Keeping whole titles therefore put the pooled gold floor at 146,543
passages: a budget of 120,000 yielded zero distractors, and 58 gold titles
alone accounted for 135,479 passages.

Passage level costs nothing the heterogeneous comparison needs. At alpha=0
the domain places a collection, so a query's gold is concentrated by
construction rather than by the atom, and a titleless domain now spreads
across its own collections instead of collapsing into one.

Claude-Session: https://claude.ai/code/session_01WhudUtZm6qqiuv8Y1sbwSc

evaluations/evaluations/datasets/mtrag_federated.py

@@ -301,7 +301,43 @@ def pooled_gold_ids(variant: str = "lastturn") -> set[str]:
 def load_pooled(
     budget: int = DEFAULT_BUDGET, seed: int = DEFAULT_SEED
 ) -> list[Mapping[str, Any]]:
-    return sample_records(load_pooled_records(), pooled_gold_ids(), budget, seed)
+    return sample_pooled_records(load_pooled_records(), pooled_gold_ids(), budget, seed)
+
+
+def sample_pooled_records(
+    records: Sequence[Mapping[str, Any]],
+    gold_ids: Iterable[str],
+    budget: int = DEFAULT_BUDGET,
+    seed: int = DEFAULT_SEED,
+) -> list[Mapping[str, Any]]:
+    """A fixed sub-corpus at passage level, keeping every gold passage.
+
+    The single-domain dataset keeps whole titles, which cannot work here: `title`
+    is the empty string for every cloud and fiqa passage, so two of the four
+    domains have exactly one title covering 72,442 and 61,022 passages. Whole
+    titles put the gold floor at 146,543 passages, leaving no distractors at any
+    budget below the entire corpus.
+
+    Passage level costs nothing this comparison needs: at alpha=0 the domain
+    places a collection, so a query's gold is concentrated by construction rather
+    than by the atom.
+    """
+    wanted = set(gold_ids)
+    by_id = {row["_id"]: row for row in records}
+    missing = sorted(wanted - set(by_id))
+    if missing:
+        raise ValueError(
+            f"{len(missing)} gold passages do not resolve to the pooled corpus, "
+            f"first few: {missing[:3]}"
+        )
+    kept = set(wanted)
+    others = [row["_id"] for row in records if row["_id"] not in wanted]
+    random.Random(seed).shuffle(others)
+    for passage_id in others:
+        if len(kept) >= budget:
+            break
+        kept.add(passage_id)
+    return [row for row in records if row["_id"] in kept]
 
 
 def partition_pooled(
@@ -310,11 +346,15 @@ def partition_pooled(
     alpha: float,
     seed: int = DEFAULT_SEED,
 ) -> dict[str, list[Mapping[str, Any]]]:
-    """Route pooled records to collections, honouring each record's domain."""
+    """Route pooled records to collections, honouring each record's domain.
+
+    Keyed on the passage id rather than the title, because two of the four
+    domains have no titles. See `sample_pooled_records`.
+    """
     names = pooled_collection_names(n)
     grouped: dict[str, list[Mapping[str, Any]]] = {name: [] for name in names}
     for row in records:
-        index = collection_of(row["title"], n, seed, alpha=alpha, domain=row["domain"])
+        index = collection_of(row["_id"], n, seed, alpha=alpha, domain=row["domain"])
         grouped[names[index]].append(row)
     return grouped
 

evaluations/tests/test_mtrag_federated.py

@@ -18,6 +18,7 @@ from evaluations.datasets.mtrag_federated import (
     partition_pooled,
     pooled_collection_names,
     pooled_database_paths,
+    sample_pooled_records,
     partition_records,
     pool_composition,
     sample_records,
@@ -421,3 +422,59 @@ class TestPooledPartition:
         for alpha in (0.0, 0.5, 1.0):
             grouped = partition_pooled(records, 8, alpha=alpha)
             assert sum(len(v) for v in grouped.values()) == len(records)
+
+
+class TestPooledSampler:
+    """Two of the four domains have no titles at all, so the pooled corpus is
+    sampled and partitioned at passage level rather than by title."""
+
+    @staticmethod
+    def _pool(per_domain: int = 40) -> list[dict[str, str]]:
+        return [
+            {
+                "_id": f"{domain}-{i}",
+                # cloud and fiqa carry an empty title upstream.
+                "title": "" if domain in ("cloud", "fiqa") else f"{domain} t{i}",
+                "text": "x",
+                "domain": domain,
+            }
+            for domain in DOMAINS
+            for i in range(per_domain)
+        ]
+
+    def test_keeps_every_gold_passage(self) -> None:
+        pool = self._pool()
+        gold = {"cloud-3", "fiqa-7", "clapnq-1", "govt-39"}
+        kept = sample_pooled_records(pool, gold, budget=20)
+        assert gold <= {row["_id"] for row in kept}
+
+    def test_respects_the_budget_above_the_gold_floor(self) -> None:
+        pool = self._pool()
+        kept = sample_pooled_records(pool, {"cloud-3"}, budget=25)
+        assert len(kept) == 25
+
+    def test_a_titleless_domain_does_not_drag_in_its_whole_corpus(self) -> None:
+        """The failure this replaces: whole-title keeping pulled all 72,442 cloud
+        passages in because they share one empty title."""
+        pool = self._pool()
+        kept = sample_pooled_records(pool, {"cloud-3"}, budget=10)
+        cloud = [row for row in kept if row["domain"] == "cloud"]
+        assert len(cloud) < 40, f"kept {len(cloud)} of 40 cloud passages"
+
+    def test_rejects_gold_the_pool_does_not_hold(self) -> None:
+        with pytest.raises(ValueError, match="do not resolve"):
+            sample_pooled_records(self._pool(), {"nope-1"}, budget=10)
+
+    def test_partition_is_passage_level_not_title_level(self) -> None:
+        """A titleless domain must still spread across its own collections."""
+        pool = self._pool(per_domain=200)
+        grouped = partition_pooled(pool, 8, alpha=0.0)
+        cloud_collections = {
+            name
+            for name, rows in grouped.items()
+            if any(row["domain"] == "cloud" for row in rows)
+        }
+        assert len(cloud_collections) == 2, (
+            f"cloud landed in {len(cloud_collections)} collections; with one empty "
+            "title a title-keyed partition would give 1"
+        )