[bugfix] tolerate duplicate SID item IDs

The global uniqueness check and duplicate overflow rejection blocked large jobs on rare upstream duplicate IDs. Reuse the matched candidate list across duplicate overflow rows while preserving the input rows and collision accounting.

Author: WhiteSwan1

tzrec/tools/sid/resolve_sid_collisions.py

@@ -11,6 +11,13 @@
 
 r"""Offline best-effort SID collision resolution with TorchEasyRec-native I/O.
 
+Item IDs are expected to be unique in the upstream prediction output. To avoid
+blocking a large job on rare upstream violations, this tool does not perform a
+full-input uniqueness check or remove duplicate rows. Duplicate rows remain
+independent items in collision accounting and outputs. When candidate-strategy
+overflow rows share an item ID, they reuse one candidate list matched for that
+ID. Duplicate data should still be fixed upstream.
+
 The runner retains the first capacity items in each SID bucket, attempts to
 relocate overflow items using fixed-width last-layer candidates, delegates
 placement to the pure NumPy core, and writes item-level and grouped SID results
@@ -131,10 +138,20 @@ class _ItemIdLookup:
     def __init__(self, item_ids: np.ndarray) -> None:
         self._sorted_to_requested = np.argsort(item_ids, kind="stable")
         self._sorted_ids = item_ids[self._sorted_to_requested]
-        if self._sorted_ids.size > 1 and np.any(
-            self._sorted_ids[1:] == self._sorted_ids[:-1]
-        ):
-            raise ValueError("overflow item IDs must be unique.")
+        duplicate_sorted_rows = (
+            np.flatnonzero(self._sorted_ids[1:] == self._sorted_ids[:-1]) + 1
+        )
+        self._duplicate_requested_rows = self._sorted_to_requested[
+            duplicate_sorted_rows
+        ]
+        representative_sorted_rows = np.searchsorted(
+            self._sorted_ids,
+            self._sorted_ids[duplicate_sorted_rows],
+            side="left",
+        )
+        self._representative_requested_rows = self._sorted_to_requested[
+            representative_sorted_rows
+        ]
 
     def match(self, item_ids: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
         """Return matching source rows and requested-ID positions."""
@@ -146,6 +163,12 @@ def match(self, item_ids: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
             ]
         return source_rows, self._sorted_to_requested[positions[source_rows]]
 
+    def broadcast_duplicate_targets(self, values: np.ndarray) -> None:
+        """Copy representative values to duplicate requested-ID positions."""
+        values[self._duplicate_requested_rows] = values[
+            self._representative_requested_rows
+        ]
+
 
 @dataclass(frozen=True)
 class ResolveSidCollisionsConfig:
@@ -406,8 +429,6 @@ def _load_codes(self) -> Tuple[np.ndarray, np.ndarray]:
         if not id_chunks:
             raise ValueError("SID input is empty.")
         item_id_array = np.concatenate(id_chunks)
-        if np.unique(item_id_array).size != item_id_array.size:
-            raise ValueError("input item IDs must be unique.")
         code_matrix = np.concatenate(code_chunks, axis=0)
         if code_matrix.shape[1] < 1:
             raise ValueError("SID codes must have at least one layer.")
@@ -469,8 +490,6 @@ def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarra
             if source_rows.size == 0:
                 continue
 
-            if np.any(seen[target_rows]):
-                raise ValueError("candidate input contains duplicate item IDs.")
             selected = pc.take(batch[field], pa.array(source_rows, type=pa.int64()))
             batch_candidates = self._candidate_last_matrix(selected)
             if candidates is None:
@@ -489,6 +508,8 @@ def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarra
             raise ValueError(
                 "map has overflow items but candidate_codes yielded no candidates."
             )
+        item_id_lookup.broadcast_duplicate_targets(candidates)
+        item_id_lookup.broadcast_duplicate_targets(seen)
         if not np.all(seen):
             missing = np.flatnonzero(~seen)
             preview = ",".join(str(value) for value in missing[:10])

tzrec/tools/sid/resolve_sid_collisions_test.py

@@ -16,6 +16,7 @@
 from collections import Counter
 from unittest import mock
 
+import numpy as np
 import pyarrow as pa
 from parameterized import parameterized
 from pyarrow import csv, parquet
@@ -25,12 +26,21 @@
     CollisionResolutionRunner,
     ResolveSidCollisionsConfig,
 )
+from tzrec.utils.sid.collision import stable_order_hash
 from tzrec.utils.test_util import make_test_dir, parameterized_name_func
 
 
-def _parquet(path, item_ids, codes, candidate_codes=None):
+def _parquet(
+    path,
+    item_ids,
+    codes,
+    candidate_codes=None,
+    item_id_type=None,
+):
+    if item_id_type is None:
+        item_id_type = pa.int64()
     cols = {
-        "item_id": pa.array(item_ids, type=pa.int64()),
+        "item_id": pa.array(item_ids, type=item_id_type),
         "codes": pa.array(codes, type=pa.list_(pa.int64())),
     }
     if candidate_codes is not None:
@@ -824,12 +834,91 @@ def test_empty_input_raises(self) -> None:
         with self.assertRaisesRegex(ValueError, "SID input is empty"):
             self._run(inp, out, max_items_per_codebook=2)
 
-    def test_duplicate_item_id_raises(self) -> None:
+    @parameterized.expand(
+        [("same_batch", 100000), ("across_batches", 1)],
+        name_func=parameterized_name_func,
+    )
+    def test_candidate_tolerates_duplicate_item_ids(
+        self, _case_name, batch_size
+    ) -> None:
         inp = os.path.join(self.test_dir, "in.parquet")
         out = os.path.join(self.test_dir, "out")
-        _parquet(inp, [0, 1, 1], [[0, 0], [0, 1], [0, 2]])
-        with self.assertRaisesRegex(ValueError, "item IDs must be unique"):
-            self._run(inp, out, max_items_per_codebook=2)
+        distinct_ids = np.asarray(["a", "b"], dtype=object)
+        hash_order = np.argsort(stable_order_hash(distinct_ids))
+        keeper = distinct_ids[hash_order[0]]
+        duplicate = distinct_ids[hash_order[1]]
+        _parquet(
+            inp,
+            [keeper, duplicate, duplicate],
+            [[0, 0]] * 3,
+            [
+                [[0, 1], [0, 2]],
+                [[0, 1], [0, 2]],
+                [[0, 2], [0, 3]],
+            ],
+            item_id_type=pa.string(),
+        )
+
+        stats = self._run(
+            inp,
+            out,
+            batch_size=batch_size,
+            include_original=True,
+            max_items_per_codebook=1,
+        )
+
+        self.assertEqual(stats.total_items, 3)
+        self.assertEqual(stats.relocated_count, 2)
+        result = self._read_parquet(out)
+        self.assertEqual(result["item_id"], [keeper, duplicate, duplicate])
+        self.assertEqual(result["item_id"].count(duplicate), 2)
+        self._assert_map_matches_resolved_groups(out)
+        original_path, _ = self._group_paths(out)
+        original_groups = self._read_parquet(original_path)
+        self.assertCountEqual(
+            [item_id for group in original_groups["itemids"] for item_id in group],
+            [keeper, duplicate, duplicate],
+        )
+
+    def test_item_id_lookup_broadcasts_all_duplicate_targets(self) -> None:
+        lookup = resolve_sid_collisions._ItemIdLookup(
+            np.asarray(["b", "a", "b", "b"], dtype=object)
+        )
+        source_rows, target_rows = lookup.match(
+            np.asarray(["b", "missing", "a"], dtype=object)
+        )
+        np.testing.assert_array_equal(source_rows, [0, 2])
+        np.testing.assert_array_equal(target_rows, [0, 1])
+
+        values = np.asarray([[10, 11], [20, 21], [-1, -1], [-1, -1]])
+        lookup.broadcast_duplicate_targets(values)
+        np.testing.assert_array_equal(
+            values,
+            [[10, 11], [20, 21], [10, 11], [10, 11]],
+        )
+
+    def test_random_tolerates_duplicate_item_ids(self) -> None:
+        inp = os.path.join(self.test_dir, "in.parquet")
+        out = os.path.join(self.test_dir, "out")
+        _parquet(
+            inp,
+            ["b", "a", "a"],
+            [[0, 0]] * 3,
+            item_id_type=pa.string(),
+        )
+
+        stats = self._run(
+            inp,
+            out,
+            strategy="random",
+            random_num_candidates=8,
+            max_items_per_codebook=1,
+        )
+
+        self.assertEqual(stats.total_items, 3)
+        result = self._read_parquet(out)
+        self.assertEqual(result["item_id"], ["b", "a", "a"])
+        self._assert_map_matches_resolved_groups(out)
 
     def test_empty_codebook_token_raises(self) -> None:
         args = resolve_sid_collisions.build_parser().parse_args(