From f47a8957073289a73a48a03c8af168884156ad76 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Mon, 6 Jul 2026 02:43:04 +0000 Subject: [PATCH 01/29] [feat] SID: add offline codebook collision-prevention tool Deterministic offline post-process over predicted SID rows that caps items per SID bucket at --max_items_per_codebook and reassigns the overflow. Strategies: - candidate (default): reassign to explicit nearest-neighbor candidate rows. - random: draw a uniform-random free within-band last-layer code, no candidate input required; a baseline that ignores semantic nearest-neighbor proximity and is still reproducible given --seed. Local (CSV/Parquet) and ODPS-SQL backends. Capacity/strategy are CLI-only policy (deliberately kept out of all protos). Output columns: item_id, origin_codebook, codebook, index. Includes unit tests. Co-Authored-By: Claude Opus 4.8 --- tzrec/tools/sid/__init__.py | 12 + tzrec/tools/sid/collision_prevention.py | 1231 ++++++++++++++++++ tzrec/tools/sid/collision_prevention_test.py | 591 +++++++++ 3 files changed, 1834 insertions(+) create mode 100644 tzrec/tools/sid/__init__.py create mode 100644 tzrec/tools/sid/collision_prevention.py create mode 100644 tzrec/tools/sid/collision_prevention_test.py diff --git a/tzrec/tools/sid/__init__.py b/tzrec/tools/sid/__init__.py new file mode 100644 index 00000000..c4aae99f --- /dev/null +++ b/tzrec/tools/sid/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SID offline tools.""" diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py new file mode 100644 index 00000000..b4b702a7 --- /dev/null +++ b/tzrec/tools/sid/collision_prevention.py @@ -0,0 +1,1231 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline SID codebook collision-prevention tool. + +The default ('candidate') allocator is a deterministic post-process over predicted +SID rows plus explicit candidate SID rows. The opt-in 'random' strategy instead +generates random within-band last-layer candidates (no candidate input required) as +a baseline that ignores semantic nearest-neighbor proximity; it is still fully +reproducible given ``seed``. +""" + +import argparse +import hashlib +import random +from collections import OrderedDict, defaultdict +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import pyarrow as pa + +# Register the local reader/writer classes used through create_reader/writer. +import tzrec.datasets.csv_dataset # noqa: F401 +import tzrec.datasets.parquet_dataset # noqa: F401 +from tzrec.datasets.dataset import create_reader, create_writer +from tzrec.utils.logging_util import logger + + +@dataclass(frozen=True) +class RawSidRow: + """One raw item -> SID row.""" + + item_id: Any + item_key: str + origin_codebook: str + + +@dataclass(frozen=True) +class CandidateSidRow: + """One item -> candidate SID row.""" + + item_key: str + candidate_codebook: str + priority: int + score: float + + +@dataclass(frozen=True) +class AssignedSidRow: + """One final item -> SID assignment row.""" + + item_id: Any + item_key: str + origin_codebook: str + codebook: str + index: int + + +@dataclass(frozen=True) +class AssignmentStats: + """Summary statistics for a collision-prevention run.""" + + total_items: int + raw_collision_buckets: int + final_collision_buckets: int + reassigned_count: int + unassigned_count: int + iteration_count: int + max_final_bucket_size: int + + +@dataclass(frozen=True) +class OdpsTableRef: + """Parsed ODPS table URI.""" + + project: str + table: str + partitions: Tuple[str, ...] + schema: Optional[str] = None + + @classmethod + def parse(cls, path: str) -> "OdpsTableRef": + """Parse an ``odps://project/tables/table[/pt=value]`` path.""" + parts = path.split("/") + if len(parts) < 5 or parts[0] != "odps:" or parts[3] != "tables": + raise ValueError( + f"invalid ODPS path {path!r}; expected " + "odps://project/tables/table[/pt=value]" + ) + table = parts[4] + schema = None + if "." in table: + schema, table = table.split(".", 1) + return cls( + project=parts[2], + table=table, + partitions=tuple(p for p in parts[5:] if p), + schema=schema, + ) + + @property + def table_name(self) -> str: + """Fully qualified ODPS table name.""" + table = f"{self.schema}.{self.table}" if self.schema else self.table + return f"{self.project}.{table}" + + @property + def partition_predicate(self) -> str: + """SQL predicate suffix for this table's partition path.""" + if not self.partitions: + return "" + predicates = [] + for key, value in self._partition_pairs(): + predicates.append(f"{key}='{value}'") + return " AND " + " AND ".join(predicates) + + @property + def insert_target(self) -> str: + """SQL INSERT target including partition spec.""" + if not self.partitions: + return self.table_name + specs = [f"{key}='{value}'" for key, value in self._partition_pairs()] + return f"{self.table_name} PARTITION ({','.join(specs)})" + + @property + def partition_schema(self) -> str: + """CREATE TABLE partition schema suffix.""" + if not self.partitions: + return "" + fields = [f"{key} STRING" for key, _ in self._partition_pairs()] + return f" PARTITIONED BY ({','.join(fields)})" + + def _partition_pairs(self) -> List[Tuple[str, str]]: + pairs = [] + for part in self.partitions: + if "=" not in part: + raise ValueError(f"invalid ODPS partition segment: {part!r}") + key, value = part.split("=", 1) + pairs.append((key, value)) + return pairs + + +class SidCollisionAssigner: + """Deterministically assign overflow SID rows to explicit candidates.""" + + def __init__( + self, + capacity: int, + max_iters: int = 50, + seed: int = 2026, + score_order: str = "lower", + unassigned_policy: str = "error", + strategy: str = "candidate", + code_delimiter: str = ",", + random_last_layer_size: Optional[int] = None, + random_num_candidates: int = 64, + ) -> None: + if capacity < 1: + raise ValueError(f"capacity must be >= 1, got {capacity}") + if score_order not in ("lower", "higher"): + raise ValueError("score_order must be 'lower' or 'higher'.") + if unassigned_policy not in ("error", "drop", "keep_original"): + raise ValueError( + "unassigned_policy must be one of: error, drop, keep_original." + ) + if strategy not in ("candidate", "random"): + raise ValueError("strategy must be 'candidate' or 'random'.") + if strategy == "random": + if random_last_layer_size is None or random_last_layer_size < 2: + raise ValueError( + "strategy='random' requires random_last_layer_size >= 2." + ) + if random_num_candidates < 1: + raise ValueError("random_num_candidates must be >= 1.") + self.capacity = capacity + self.max_iters = max_iters + self.seed = seed + self.score_order = score_order + self.unassigned_policy = unassigned_policy + self.strategy = strategy + self.code_delimiter = code_delimiter + self.random_last_layer_size = random_last_layer_size + self.random_num_candidates = random_num_candidates + + def assign( + self, + raw_rows: Sequence[RawSidRow], + candidate_rows: Sequence[CandidateSidRow], + ) -> Tuple[List[AssignedSidRow], AssignmentStats]: + """Assign overflow SID rows to non-full candidate codebooks.""" + if len({row.item_key for row in raw_rows}) != len(raw_rows): + raise ValueError("raw_rows contains duplicate item_id values.") + + raw_by_item = {row.item_key: row for row in raw_rows} + by_origin: Dict[str, List[RawSidRow]] = defaultdict(list) + for row in raw_rows: + by_origin[row.origin_codebook].append(row) + + assigned: List[AssignedSidRow] = [] + assigned_items = set() + code_counts: Dict[str, int] = defaultdict(int) + + for codebook, rows in by_origin.items(): + for index, row in enumerate( + sorted(rows, key=self._assignment_sort_key)[: self.capacity], + start=1, + ): + assigned.append( + AssignedSidRow( + item_id=row.item_id, + item_key=row.item_key, + origin_codebook=row.origin_codebook, + codebook=codebook, + index=index, + ) + ) + assigned_items.add(row.item_key) + code_counts[codebook] += 1 + + overflow_items = set(raw_by_item) - assigned_items + if self.strategy == "random": + candidate_rows = self._generate_random_candidates( + overflow_items, raw_by_item + ) + elif overflow_items and not candidate_rows: + raise ValueError( + "raw SID input has overflow rows, but no explicit candidate input was " + "provided." + ) + candidates_by_item = self._dedup_candidates(candidate_rows, raw_by_item) + + iteration_count = self._assign_candidates( + raw_by_item, + candidates_by_item, + assigned, + assigned_items, + code_counts, + ) + unassigned = self._handle_unassigned( + raw_by_item, + assigned, + assigned_items, + code_counts, + ) + final_counts = self._final_counts(assigned) + stats = AssignmentStats( + total_items=len(raw_rows), + raw_collision_buckets=self._raw_collision_buckets(raw_rows), + final_collision_buckets=sum( + 1 for count in final_counts.values() if count > self.capacity + ), + reassigned_count=sum( + 1 for row in assigned if row.origin_codebook != row.codebook + ), + unassigned_count=( + len(unassigned) if self.unassigned_policy != "keep_original" else 0 + ), + iteration_count=iteration_count, + max_final_bucket_size=max(final_counts.values()) if final_counts else 0, + ) + return sorted(assigned, key=lambda r: (r.codebook, r.index, r.item_key)), stats + + @staticmethod + def _stable_hash(*parts: Any) -> int: + h = hashlib.blake2b(digest_size=8) + for part in parts: + h.update(str(part).encode("utf-8")) + h.update(b"\x1f") + return int.from_bytes(h.digest(), byteorder="big", signed=False) + + def _raw_collision_buckets(self, raw_rows: Sequence[RawSidRow]) -> int: + counts: Dict[str, int] = defaultdict(int) + for row in raw_rows: + counts[row.origin_codebook] += 1 + return sum(1 for count in counts.values() if count > self.capacity) + + def _generate_random_candidates( + self, + overflow_items: Iterable[str], + raw_by_item: Dict[str, RawSidRow], + ) -> List[CandidateSidRow]: + """Generate random within-band last-layer candidates for overflow items. + + Keeps every layer except the last and draws distinct random codes for the + last layer in ``[0, random_last_layer_size)``, excluding the origin code, so + each item stays in its own ``(prefix)`` band. Deterministic given ``seed``; + unlike the 'candidate' strategy it ignores nearest-neighbor proximity. + """ + assert self.random_last_layer_size is not None + size = self.random_last_layer_size + num_draws = min(self.random_num_candidates, size - 1) + rows: List[CandidateSidRow] = [] + for item_key in overflow_items: + raw = raw_by_item[item_key] + parts = raw.origin_codebook.split(self.code_delimiter) + prefix = parts[:-1] + try: + origin_last: Optional[int] = int(parts[-1]) + except ValueError: + origin_last = None + rng = random.Random(self._stable_hash(self.seed, item_key)) + drawn: set = set() + while len(drawn) < num_draws: + value = rng.randrange(size) + if value in drawn or value == origin_last: + continue + drawn.add(value) + rows.append( + CandidateSidRow( + item_key=item_key, + candidate_codebook=self.code_delimiter.join( + [*prefix, str(value)] + ), + priority=1, + score=0.0, + ) + ) + return rows + + def _assignment_sort_key(self, row: RawSidRow) -> Tuple[int, str]: + return ( + self._stable_hash(self.seed, row.origin_codebook, row.item_key), + row.item_key, + ) + + def _candidate_sort_key( + self, + row: CandidateSidRow, + ) -> Tuple[int, float, int, str, str]: + score = row.score if self.score_order == "lower" else -row.score + return ( + row.priority, + score, + self._stable_hash(self.seed, row.item_key, row.candidate_codebook), + row.item_key, + row.candidate_codebook, + ) + + def _dedup_candidates( + self, + candidate_rows: Sequence[CandidateSidRow], + raw_by_item: Dict[str, RawSidRow], + ) -> Dict[str, List[CandidateSidRow]]: + dedup_candidates: Dict[Tuple[str, str], CandidateSidRow] = {} + for row in candidate_rows: + if row.item_key not in raw_by_item: + continue + key = (row.item_key, row.candidate_codebook) + current = dedup_candidates.get(key) + if current is None or self._candidate_sort_key( + row + ) < self._candidate_sort_key(current): + dedup_candidates[key] = row + + candidates_by_item: Dict[str, List[CandidateSidRow]] = defaultdict(list) + for row in dedup_candidates.values(): + candidates_by_item[row.item_key].append(row) + return candidates_by_item + + def _assign_candidates( + self, + raw_by_item: Dict[str, RawSidRow], + candidates_by_item: Dict[str, List[CandidateSidRow]], + assigned: List[AssignedSidRow], + assigned_items: set, + code_counts: Dict[str, int], + ) -> int: + iteration_count = 0 + for iteration in range(self.max_iters): + unassigned = set(raw_by_item) - assigned_items + if not unassigned: + break + + available = self._available_candidates( + unassigned, + candidates_by_item, + code_counts, + ) + if not available: + break + + accepted = self._select_candidates(available, code_counts) + progress = 0 + for candidate in accepted: + if candidate.item_key in assigned_items: + continue + if code_counts[candidate.candidate_codebook] >= self.capacity: + continue + raw = raw_by_item[candidate.item_key] + code_counts[candidate.candidate_codebook] += 1 + assigned.append( + AssignedSidRow( + item_id=raw.item_id, + item_key=raw.item_key, + origin_codebook=raw.origin_codebook, + codebook=candidate.candidate_codebook, + index=code_counts[candidate.candidate_codebook], + ) + ) + assigned_items.add(candidate.item_key) + progress += 1 + + iteration_count = iteration + 1 + if progress == 0: + break + return iteration_count + + def _available_candidates( + self, + unassigned: Sequence[str], + candidates_by_item: Dict[str, List[CandidateSidRow]], + code_counts: Dict[str, int], + ) -> List[CandidateSidRow]: + available = [] + for item_key in unassigned: + for candidate in candidates_by_item.get(item_key, []): + if code_counts[candidate.candidate_codebook] < self.capacity: + available.append(candidate) + return available + + def _select_candidates( + self, + available: Sequence[CandidateSidRow], + code_counts: Dict[str, int], + ) -> List[CandidateSidRow]: + selected_by_codebook: List[CandidateSidRow] = [] + by_codebook: Dict[str, List[CandidateSidRow]] = defaultdict(list) + for candidate in available: + by_codebook[candidate.candidate_codebook].append(candidate) + + for codebook, rows in by_codebook.items(): + remaining = self.capacity - code_counts[codebook] + if remaining <= 0: + continue + selected_by_codebook.extend( + sorted(rows, key=self._candidate_sort_key)[:remaining] + ) + + best_by_item: Dict[str, CandidateSidRow] = {} + for candidate in selected_by_codebook: + current = best_by_item.get(candidate.item_key) + if current is None or self._candidate_sort_key( + candidate + ) < self._candidate_sort_key(current): + best_by_item[candidate.item_key] = candidate + + return sorted( + best_by_item.values(), + key=lambda r: (r.candidate_codebook, self._candidate_sort_key(r)), + ) + + def _handle_unassigned( + self, + raw_by_item: Dict[str, RawSidRow], + assigned: List[AssignedSidRow], + assigned_items: set, + code_counts: Dict[str, int], + ) -> List[str]: + unassigned = sorted(set(raw_by_item) - assigned_items) + if not unassigned: + return unassigned + + if self.unassigned_policy == "error": + preview = ",".join(unassigned[:10]) + raise RuntimeError( + f"{len(unassigned)} items could not be assigned within capacity; " + f"first unassigned item_ids: {preview}" + ) + if self.unassigned_policy == "keep_original": + for item_key in unassigned: + raw = raw_by_item[item_key] + code_counts[raw.origin_codebook] += 1 + assigned.append( + AssignedSidRow( + item_id=raw.item_id, + item_key=raw.item_key, + origin_codebook=raw.origin_codebook, + codebook=raw.origin_codebook, + index=code_counts[raw.origin_codebook], + ) + ) + return unassigned + + @staticmethod + def _final_counts(rows: Sequence[AssignedSidRow]) -> Dict[str, int]: + final_counts: Dict[str, int] = defaultdict(int) + for row in rows: + final_counts[row.codebook] += 1 + return final_counts + + +class LocalCollisionRunner: + """CSV/Parquet SID collision-prevention runner.""" + + def __init__(self, args: argparse.Namespace) -> None: + self.args = args + + def run(self) -> AssignmentStats: + """Run local CSV/Parquet collision prevention.""" + raw_rows = self._load_raw_sid_rows() + candidate_rows = self._load_candidate_rows() + rows, stats = SidCollisionAssigner( + capacity=self.args.max_items_per_codebook, + max_iters=self.args.max_iters, + seed=self.args.seed, + score_order=self.args.score_order, + unassigned_policy=self.args.unassigned_policy, + strategy=self.args.strategy, + code_delimiter=self.args.code_delimiter, + random_last_layer_size=self.args.random_last_layer_size, + random_num_candidates=self.args.random_num_candidates, + ).assign(raw_rows, candidate_rows) + self._write_assignments(rows) + if self.args.diagnostics_output_path: + self._write_diagnostics(stats) + logger.info("SID collision prevention finished: %s", stats) + return stats + + @staticmethod + def _cell_to_code(value: Any, delimiter: str) -> str: + if value is None: + raise ValueError("SID code value cannot be null.") + if isinstance(value, (list, tuple)): + return delimiter.join(str(v) for v in value) + return str(value) + + @classmethod + def _split_compact_candidates(cls, value: Any, delimiter: str) -> List[str]: + if value is None: + return [] + if isinstance(value, (list, tuple)): + return [cls._cell_to_code(v, ",") for v in value if v is not None] + text = str(value).strip() + if not text: + return [] + return [part.strip() for part in text.split(delimiter) if part.strip()] + + @staticmethod + def _array_to_pylist(batch: Dict[str, pa.Array], field: str) -> List[Any]: + if field not in batch: + raise ValueError( + f"required field {field!r} not found; available fields: " + f"{sorted(batch.keys())}" + ) + return batch[field].to_pylist() + + def _read_batches( + self, + input_path: str, + reader_type: Optional[str], + ) -> Iterable[Dict[str, pa.Array]]: + reader = create_reader( + input_path=input_path, + batch_size=self.args.batch_size, + reader_type=reader_type, + quota_name=self.args.odps_data_quota_name, + ) + yield from reader.to_batches() + + def _load_raw_sid_rows(self) -> List[RawSidRow]: + rows: List[RawSidRow] = [] + seen = set() + code_fields = self._code_fields() + code_field = None if code_fields else self.args.code_field + if bool(code_field) == bool(code_fields): + raise ValueError("Set exactly one of --code_field or --code_fields.") + + for batch in self._read_batches(self.args.input_path, self.args.reader_type): + item_ids = self._array_to_pylist(batch, self.args.item_id_field) + if code_field: + codes = self._array_to_pylist(batch, code_field) + origin_codes = [ + self._cell_to_code(v, self.args.code_delimiter) for v in codes + ] + else: + assert code_fields is not None + code_columns = [self._array_to_pylist(batch, f) for f in code_fields] + origin_codes = [ + self.args.code_delimiter.join(str(col[i]) for col in code_columns) + for i in range(len(item_ids)) + ] + + for item_id, origin_codebook in zip(item_ids, origin_codes): + item_key = str(item_id) + if item_key in seen: + raise ValueError(f"duplicate item_id in raw SID input: {item_key}") + seen.add(item_key) + rows.append( + RawSidRow( + item_id=item_id, + item_key=item_key, + origin_codebook=origin_codebook, + ) + ) + + if not rows: + raise ValueError("raw SID input is empty.") + return rows + + def _load_candidate_rows(self) -> List[CandidateSidRow]: + rows: List[CandidateSidRow] = [] + if not self.args.candidate_input_path: + return rows + + candidate_codebook_field = ( + None + if self.args.compact_candidate_field + else self.args.candidate_codebook_field + ) + if bool(candidate_codebook_field) == bool(self.args.compact_candidate_field): + raise ValueError( + "Set exactly one of --candidate_codebook_field or " + "--compact_candidate_field." + ) + + reader_type = self.args.candidate_reader_type or self.args.reader_type + item_id_field = self.args.candidate_item_id_field or self.args.item_id_field + for batch in self._read_batches(self.args.candidate_input_path, reader_type): + item_ids = self._array_to_pylist(batch, item_id_field) + priorities = ( + self._array_to_pylist(batch, self.args.priority_field) + if self.args.priority_field and self.args.priority_field in batch + else None + ) + scores = ( + self._array_to_pylist(batch, self.args.score_field) + if self.args.score_field and self.args.score_field in batch + else None + ) + if candidate_codebook_field: + candidates = self._array_to_pylist(batch, candidate_codebook_field) + for i, (item_id, candidate) in enumerate(zip(item_ids, candidates)): + rows.append( + CandidateSidRow( + item_key=str(item_id), + candidate_codebook=self._cell_to_code( + candidate, + self.args.code_delimiter, + ), + priority=int(priorities[i]) + if priorities is not None + else 1, + score=float(scores[i]) if scores is not None else 0.0, + ) + ) + else: + compact_values = self._array_to_pylist( + batch, + self.args.compact_candidate_field, + ) + for item_id, compact_value in zip(item_ids, compact_values): + for priority, candidate in enumerate( + self._split_compact_candidates( + compact_value, + self.args.candidate_delimiter, + ), + start=1, + ): + rows.append( + CandidateSidRow( + item_key=str(item_id), + candidate_codebook=candidate, + priority=priority, + score=0.0, + ) + ) + + return rows + + def _code_fields(self) -> Optional[List[str]]: + if not self.args.code_fields: + return None + return [ + field.strip() for field in self.args.code_fields.split(",") if field.strip() + ] + + @staticmethod + def _item_id_array(rows: Sequence[AssignedSidRow]) -> pa.Array: + values = [row.item_id for row in rows] + if all(isinstance(v, int) and not isinstance(v, bool) for v in values): + return pa.array(values, type=pa.int64()) + return pa.array([str(v) for v in values], type=pa.string()) + + def _write_assignments(self, rows: Sequence[AssignedSidRow]) -> None: + writer = create_writer( + self.args.output_path, + writer_type=self.args.writer_type, + quota_name=self.args.odps_data_quota_name, + world_size=1, + ) + writer.write( + OrderedDict( + [ + ("item_id", self._item_id_array(rows)), + ( + "origin_codebook", + pa.array( + [row.origin_codebook for row in rows], + type=pa.string(), + ), + ), + ( + "codebook", + pa.array([row.codebook for row in rows], type=pa.string()), + ), + ("index", pa.array([row.index for row in rows], type=pa.int64())), + ] + ) + ) + writer.close() + + def _write_diagnostics(self, stats: AssignmentStats) -> None: + writer = create_writer( + self.args.diagnostics_output_path, + writer_type=self.args.writer_type, + quota_name=self.args.odps_data_quota_name, + world_size=1, + ) + writer.write( + OrderedDict( + [ + ("total_items", pa.array([stats.total_items], type=pa.int64())), + ( + "raw_collision_buckets", + pa.array([stats.raw_collision_buckets], type=pa.int64()), + ), + ( + "final_collision_buckets", + pa.array([stats.final_collision_buckets], type=pa.int64()), + ), + ( + "reassigned_count", + pa.array([stats.reassigned_count], type=pa.int64()), + ), + ( + "unassigned_count", + pa.array([stats.unassigned_count], type=pa.int64()), + ), + ( + "iteration_count", + pa.array([stats.iteration_count], type=pa.int64()), + ), + ( + "max_final_bucket_size", + pa.array([stats.max_final_bucket_size], type=pa.int64()), + ), + ] + ) + ) + writer.close() + + +class OdpsSqlGenerator: + """Generate deterministic MaxCompute SQL for canonical candidate tables.""" + + def __init__(self, args: argparse.Namespace) -> None: + if not args.candidate_input_path: + raise ValueError("--candidate_input_path is required for --backend odps.") + if args.code_fields: + raise ValueError( + "--backend odps currently supports --code_field, not split code_fields." + ) + if args.compact_candidate_field: + raise ValueError( + "--backend odps expects canonical candidate rows; compact candidates " + "are supported in local CSV/Parquet mode." + ) + + self.args = args + self.raw_ref = OdpsTableRef.parse(args.input_path) + self.candidate_ref = OdpsTableRef.parse(args.candidate_input_path) + self.output_ref = OdpsTableRef.parse(args.output_path) + self.prefix = args.temp_prefix or "tmp_sid_collision" + self.assigned = f"{self.prefix}_assigned" + self.selected = f"{self.prefix}_selected" + self.counts = f"{self.prefix}_counts" + + def generate(self) -> List[str]: + """Generate the SQL statements for one ODPS collision-prevention run.""" + sqls = [ + "SET odps.sql.type.system.odps2=true", + self._create_assigned_sql(), + self._create_selected_sql(), + self._create_counts_sql(), + self._create_output_sql(), + self._initial_assignment_sql(), + ] + for i in range(self.args.max_iters): + sqls.extend( + [ + self._refresh_counts_sql(), + self._select_candidates_sql(), + self._insert_selected_sql(), + self._remaining_unassigned_sql(f"iteration {i + 1}"), + ] + ) + sqls.append(self._remaining_unassigned_sql("final")) + if self.args.unassigned_policy == "keep_original": + sqls.extend([self._refresh_counts_sql(), self._keep_original_sql()]) + sqls.append(self._final_insert_sql()) + return sqls + + @property + def _raw_table(self) -> str: + return self.raw_ref.table_name + + @property + def _candidate_table(self) -> str: + return self.candidate_ref.table_name + + @property + def _raw_predicate(self) -> str: + return self.raw_ref.partition_predicate + + @property + def _candidate_predicate(self) -> str: + return self.candidate_ref.partition_predicate + + @property + def _score_expr(self) -> str: + if self.args.score_field: + return f"CAST({self.args.score_field} AS DOUBLE)" + return "0.0" + + @property + def _priority_expr(self) -> str: + if self.args.priority_field: + return f"CAST({self.args.priority_field} AS BIGINT)" + return "1" + + @property + def _score_order(self) -> str: + return "score ASC" if self.args.score_order == "lower" else "score DESC" + + def _create_assigned_sql(self) -> str: + return ( + f"CREATE TABLE IF NOT EXISTS {self.assigned} (" + "item_id STRING, origin_codebook STRING, codebook STRING, `index` BIGINT" + f") LIFECYCLE {self.args.odps_lifecycle}" + ) + + def _create_selected_sql(self) -> str: + return ( + f"CREATE TABLE IF NOT EXISTS {self.selected} (" + "item_id STRING, origin_codebook STRING, codebook STRING, " + "priority BIGINT, score DOUBLE" + f") LIFECYCLE {self.args.odps_lifecycle}" + ) + + def _create_counts_sql(self) -> str: + return ( + f"CREATE TABLE IF NOT EXISTS {self.counts} (" + "codebook STRING, cnt BIGINT" + f") LIFECYCLE {self.args.odps_lifecycle}" + ) + + def _create_output_sql(self) -> str: + return ( + f"CREATE TABLE IF NOT EXISTS {self.output_ref.table_name} (" + "item_id STRING, origin_codebook STRING, codebook STRING, `index` BIGINT" + f"){self.output_ref.partition_schema}" + ) + + def _initial_assignment_sql(self) -> str: + return ( + f"INSERT OVERWRITE TABLE {self.assigned}\n" + "SELECT item_id, origin_codebook, origin_codebook AS codebook,\n" + " rn AS `index`\n" + "FROM (\n" + f" SELECT CAST({self.args.item_id_field} AS STRING) AS item_id,\n" + f" CAST({self.args.code_field} AS STRING) AS origin_codebook,\n" + " ROW_NUMBER() OVER (\n" + f" PARTITION BY CAST({self.args.code_field} AS STRING)\n" + " ORDER BY ABS(HASH(CONCAT(" + f"'{self.args.seed}', ':', CAST({self.args.item_id_field} AS STRING))))\n" + " ) AS rn\n" + f" FROM {self._raw_table}\n" + f" WHERE 1=1{self._raw_predicate}\n" + ") t\n" + f"WHERE rn <= {self.args.max_items_per_codebook}" + ) + + def _refresh_counts_sql(self) -> str: + return ( + f"INSERT OVERWRITE TABLE {self.counts}\n" + "SELECT codebook, COUNT(*) AS cnt\n" + f"FROM {self.assigned}\n" + "GROUP BY codebook" + ) + + def _select_candidates_sql(self) -> str: + return ( + f"INSERT OVERWRITE TABLE {self.selected}\n" + "SELECT item_id, origin_codebook, codebook, priority, score\n" + "FROM (\n" + " SELECT c.item_id, c.origin_codebook, c.codebook, " + "c.priority, c.score,\n" + " ROW_NUMBER() OVER (\n" + " PARTITION BY c.codebook\n" + f" ORDER BY c.priority ASC, c.{self._score_order}, " + "ABS(HASH(CONCAT(" + f"'{self.args.seed}', ':', c.item_id, ':', c.codebook)))\n" + " ) AS rn,\n" + " COALESCE(cnt.cnt, 0) AS current_cnt\n" + " FROM (\n" + " SELECT CAST(" + f"{self.args.candidate_item_id_field or self.args.item_id_field}" + " AS STRING) AS item_id,\n" + " CAST(" + f"{self.args.candidate_origin_codebook_field} AS STRING" + ") AS origin_codebook,\n" + " CAST(" + f"{self.args.candidate_codebook_field}" + " AS STRING) AS codebook,\n" + f" {self._priority_expr} AS priority,\n" + f" {self._score_expr} AS score\n" + f" FROM {self._candidate_table}\n" + f" WHERE 1=1{self._candidate_predicate}\n" + " ) c\n" + " INNER JOIN (\n" + f" SELECT CAST({self.args.item_id_field} AS STRING) AS item_id\n" + f" FROM {self._raw_table}\n" + f" WHERE 1=1{self._raw_predicate}\n" + " ) r ON c.item_id = r.item_id\n" + f" LEFT OUTER JOIN {self.assigned} a ON c.item_id = a.item_id\n" + f" LEFT OUTER JOIN {self.counts} cnt ON c.codebook = cnt.codebook\n" + " WHERE a.item_id IS NULL\n" + f" AND COALESCE(cnt.cnt, 0) < {self.args.max_items_per_codebook}\n" + ") ranked\n" + f"WHERE rn <= {self.args.max_items_per_codebook} - current_cnt" + ) + + def _insert_selected_sql(self) -> str: + # ``codebook_rn`` must rank ONLY the surviving (``item_rn = 1``) rows so + # the emitted ``index`` stays dense and contiguous with the rows already + # in ``assigned`` (counted by ``current_cnt``). Ranking before the + # ``item_rn = 1`` filter leaves a gap whenever a codebook's top-ranked + # candidate wins a different codebook (its row here is dropped): that gap + # both wastes capacity and is re-issued next iteration as a duplicate + # ``(codebook, index)`` pair, since ``current_cnt`` (a COUNT of inserted + # rows) then undercounts the real slot high-water mark. Hence the nested + # shape: pick each item's best codebook first, THEN densely number the + # survivors per codebook. + return ( + f"INSERT INTO TABLE {self.assigned}\n" + "SELECT item_id, origin_codebook, codebook,\n" + " current_cnt + codebook_rn AS `index`\n" + "FROM (\n" + " SELECT item_id, origin_codebook, codebook, current_cnt,\n" + " ROW_NUMBER() OVER (\n" + " PARTITION BY codebook\n" + f" ORDER BY priority ASC, {self._score_order}, " + "ABS(HASH(CONCAT(" + f"'{self.args.seed}', ':', item_id, ':', codebook)))\n" + " ) AS codebook_rn\n" + " FROM (\n" + " SELECT s.item_id, s.origin_codebook, s.codebook, " + "s.priority, s.score,\n" + " ROW_NUMBER() OVER (\n" + " PARTITION BY s.item_id\n" + f" ORDER BY s.priority ASC, s.{self._score_order}, " + "ABS(HASH(CONCAT(" + f"'{self.args.seed}', ':', s.item_id, ':', s.codebook)))\n" + " ) AS item_rn,\n" + " COALESCE(cnt.cnt, 0) AS current_cnt\n" + f" FROM {self.selected} s\n" + f" LEFT OUTER JOIN {self.counts} cnt " + "ON s.codebook = cnt.codebook\n" + " ) x\n" + " WHERE item_rn = 1\n" + ") y\n" + f"WHERE current_cnt + codebook_rn <= {self.args.max_items_per_codebook}" + ) + + def _remaining_unassigned_sql(self, label: str) -> str: + return ( + f"-- {label}: remaining_unassigned\n" + "SELECT COUNT(*) AS remaining_unassigned\n" + f"FROM {self._raw_table} r\n" + f"LEFT OUTER JOIN {self.assigned} a\n" + f"ON CAST(r.{self.args.item_id_field} AS STRING) = a.item_id\n" + f"WHERE a.item_id IS NULL{self._raw_predicate}" + ) + + def _keep_original_sql(self) -> str: + return ( + f"INSERT INTO TABLE {self.assigned}\n" + "SELECT u.item_id, u.origin_codebook, " + "u.origin_codebook AS codebook,\n" + " COALESCE(cnt.cnt, 0) + u.rn AS `index`\n" + "FROM (\n" + f" SELECT CAST(r.{self.args.item_id_field} AS STRING) AS item_id,\n" + f" CAST(r.{self.args.code_field} AS STRING) AS " + "origin_codebook,\n" + " ROW_NUMBER() OVER (\n" + f" PARTITION BY CAST(r.{self.args.code_field} AS STRING)\n" + " ORDER BY ABS(HASH(CONCAT(" + f"'{self.args.seed}', ':', CAST(r.{self.args.item_id_field} AS STRING))))\n" + " ) AS rn\n" + f" FROM {self._raw_table} r\n" + f" LEFT OUTER JOIN {self.assigned} a\n" + f" ON CAST(r.{self.args.item_id_field} AS STRING) = a.item_id\n" + f" WHERE a.item_id IS NULL{self._raw_predicate}\n" + ") u\n" + f"LEFT OUTER JOIN {self.counts} cnt\n" + "ON u.origin_codebook = cnt.codebook" + ) + + def _final_insert_sql(self) -> str: + return ( + f"INSERT OVERWRITE TABLE {self.output_ref.insert_target}\n" + "SELECT item_id, origin_codebook, codebook, `index`\n" + f"FROM {self.assigned}" + ) + + +class OdpsCollisionRunner: + """ODPS SID collision-prevention runner.""" + + def __init__(self, args: argparse.Namespace) -> None: + self.args = args + + def run(self) -> None: + """Run MaxCompute SQL collision prevention.""" + sqls = OdpsSqlGenerator(self.args).generate() + if self.args.dry_run_sql: + print(";\n\n".join(sqls) + ";") + return + + from odps import ODPS + + from tzrec.datasets.odps_dataset import _create_odps_account + + output_ref = OdpsTableRef.parse(self.args.output_path) + account, endpoint = _create_odps_account() + odps = ODPS(account=account, project=output_ref.project, endpoint=endpoint) + for sql in sqls: + logger.info("Executing ODPS SQL:\n%s", sql) + instance = odps.execute_sql(sql) + instance.wait_for_success() + if not self._is_remaining_unassigned_sql(sql): + continue + + remaining_unassigned = self._read_scalar(instance) + logger.info( + "ODPS remaining unassigned item count: %s", + remaining_unassigned, + ) + if not self._is_final_remaining_unassigned_sql(sql): + continue + self._handle_final_unassigned(remaining_unassigned) + + @staticmethod + def _is_remaining_unassigned_sql(sql: str) -> bool: + return " AS remaining_unassigned" in sql + + @staticmethod + def _is_final_remaining_unassigned_sql(sql: str) -> bool: + return sql.lstrip().startswith("-- final:") + + @staticmethod + def _read_scalar(instance: Any) -> int: + with instance.open_reader() as reader: + for record in reader: + return int(record[0]) + return 0 + + def _handle_final_unassigned(self, remaining_unassigned: int) -> None: + if remaining_unassigned == 0: + return + if self.args.unassigned_policy == "error": + raise RuntimeError( + f"{remaining_unassigned} items could not be assigned within " + "capacity in ODPS collision-prevention run." + ) + if self.args.unassigned_policy == "drop": + logger.warning( + "Dropping %s unassigned items because --unassigned_policy=drop.", + remaining_unassigned, + ) + + +def assign_sid_collisions( + raw_rows: Sequence[RawSidRow], + candidate_rows: Sequence[CandidateSidRow], + capacity: int, + max_iters: int = 50, + seed: int = 2026, + score_order: str = "lower", + unassigned_policy: str = "error", + strategy: str = "candidate", + code_delimiter: str = ",", + random_last_layer_size: Optional[int] = None, + random_num_candidates: int = 64, +) -> Tuple[List[AssignedSidRow], AssignmentStats]: + """Assign overflow SID rows to non-full candidate codebooks.""" + return SidCollisionAssigner( + capacity=capacity, + max_iters=max_iters, + seed=seed, + score_order=score_order, + unassigned_policy=unassigned_policy, + strategy=strategy, + code_delimiter=code_delimiter, + random_last_layer_size=random_last_layer_size, + random_num_candidates=random_num_candidates, + ).assign(raw_rows, candidate_rows) + + +def run_local(args: argparse.Namespace) -> AssignmentStats: + """Run local CSV/Parquet collision prevention.""" + return LocalCollisionRunner(args).run() + + +def generate_odps_sql(args: argparse.Namespace) -> List[str]: + """Generate deterministic MaxCompute SQL for canonical candidate tables.""" + return OdpsSqlGenerator(args).generate() + + +def run_odps(args: argparse.Namespace) -> None: + """Run MaxCompute SQL collision prevention.""" + OdpsCollisionRunner(args).run() + + +def build_parser() -> argparse.ArgumentParser: + """Build the command line argument parser.""" + parser = argparse.ArgumentParser( + description="Prevent SID codebook collisions with explicit candidates." + ) + parser.add_argument("--backend", choices=["local", "odps"], default="local") + parser.add_argument("--input_path", required=True) + parser.add_argument("--output_path", required=True) + parser.add_argument("--candidate_input_path", default=None) + parser.add_argument("--diagnostics_output_path", default=None) + parser.add_argument( + "--reader_type", choices=["CsvReader", "ParquetReader"], default=None + ) + parser.add_argument( + "--candidate_reader_type", + choices=["CsvReader", "ParquetReader"], + default=None, + ) + parser.add_argument( + "--writer_type", + choices=["CsvWriter", "ParquetWriter"], + default="ParquetWriter", + ) + parser.add_argument("--batch_size", type=int, default=4096) + parser.add_argument("--item_id_field", default="item_id") + parser.add_argument("--candidate_item_id_field", default=None) + parser.add_argument("--code_field", default="codes") + parser.add_argument( + "--code_fields", + default=None, + help="Comma-separated split code fields. Mutually exclusive with code_field.", + ) + parser.add_argument("--code_delimiter", default=",") + parser.add_argument("--candidate_codebook_field", default="candidate_codebook") + parser.add_argument("--candidate_origin_codebook_field", default="origin_codebook") + parser.add_argument( + "--compact_candidate_field", + default=None, + help="Compact string/list candidate field for local mode.", + ) + parser.add_argument( + "--candidate_delimiter", + default="|", + help="Delimiter for compact string candidate lists.", + ) + parser.add_argument("--priority_field", default="priority") + parser.add_argument("--score_field", default="score") + parser.add_argument("--score_order", choices=["lower", "higher"], default="lower") + parser.add_argument("--max_items_per_codebook", type=int, required=True) + parser.add_argument("--max_iters", type=int, default=50) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument( + "--unassigned_policy", + choices=["error", "drop", "keep_original"], + default="error", + ) + parser.add_argument( + "--strategy", + choices=["candidate", "random"], + default="candidate", + help="Reassignment strategy: 'candidate' uses explicit nearest-neighbor " + "candidate rows; 'random' draws random within-band last-layer codes and " + "needs no candidate input (local backend only).", + ) + parser.add_argument( + "--random_last_layer_size", + type=int, + default=None, + help="Code-space size of the last SID layer; required for --strategy random.", + ) + parser.add_argument( + "--random_num_candidates", + type=int, + default=64, + help="Random last-layer codes drawn per overflow item for --strategy random.", + ) + parser.add_argument("--odps_data_quota_name", default="pay-as-you-go") + parser.add_argument("--temp_prefix", default=None) + parser.add_argument("--odps_lifecycle", type=int, default=7) + parser.add_argument("--dry_run_sql", action="store_true", default=False) + return parser + + +def main() -> None: + """Command line entrypoint.""" + parser = build_parser() + args = parser.parse_args() + if args.backend == "local": + run_local(args) + else: + if args.strategy == "random": + raise NotImplementedError( + "strategy='random' is only supported for --backend local." + ) + run_odps(args) + + +if __name__ == "__main__": + main() diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py new file mode 100644 index 00000000..880fb34e --- /dev/null +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -0,0 +1,591 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shutil +import sys +import tempfile +import types +import unittest +from collections import Counter, defaultdict +from unittest import mock + +import pyarrow as pa +from pyarrow import csv, parquet + +from tzrec.tools.sid.collision_prevention import ( + CandidateSidRow, + OdpsSqlGenerator, + RawSidRow, + assign_sid_collisions, + build_parser, + generate_odps_sql, + run_local, + run_odps, +) + + +class SidCollisionPreventionTest(unittest.TestCase): + def setUp(self) -> None: + if not os.path.exists("./tmp"): + os.makedirs("./tmp") + self.test_dir = tempfile.mkdtemp(prefix="tzrec_", dir="./tmp") + + def tearDown(self) -> None: + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + + def test_assign_sid_collisions_respects_capacity(self) -> None: + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_2", "item_2", "A"), + RawSidRow("item_3", "item_3", "B"), + ] + candidate_rows = [ + CandidateSidRow("item_0", "C", 1, 0.1), + CandidateSidRow("item_1", "C", 1, 0.1), + CandidateSidRow("item_2", "C", 1, 0.1), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, + candidate_rows, + capacity=2, + seed=7, + ) + + self.assertEqual(len(assigned), 4) + self.assertEqual(len({row.item_key for row in assigned}), 4) + self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 2) + self.assertEqual(stats.raw_collision_buckets, 1) + self.assertEqual(stats.final_collision_buckets, 0) + self.assertEqual(stats.reassigned_count, 1) + self.assertEqual(stats.unassigned_count, 0) + + def test_random_strategy_reassigns_within_band_without_candidates(self) -> None: + # Three items collide on SID (lv1,lv2,lv3) = (1,2,3); capacity 1. + raw_rows = [ + RawSidRow("item_0", "item_0", "1,2,3"), + RawSidRow("item_1", "item_1", "1,2,3"), + RawSidRow("item_2", "item_2", "1,2,3"), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, + [], # random needs no candidate input + capacity=1, + strategy="random", + random_last_layer_size=16, + seed=7, + ) + + self.assertEqual(len(assigned), 3) + codebooks = {row.item_key: row.codebook for row in assigned} + # every item keeps its (lv1,lv2) band, only the last layer varies ... + for codebook in codebooks.values(): + self.assertTrue(codebook.startswith("1,2,")) + # ... capacity 1 keeps exactly one item at the origin SID and reassigns the + # other two to distinct random last-layer codes -> near-injective. + self.assertEqual(len(set(codebooks.values())), 3) + self.assertEqual(sum(1 for cb in codebooks.values() if cb == "1,2,3"), 1) + self.assertEqual(stats.final_collision_buckets, 0) + self.assertEqual(stats.reassigned_count, 2) + self.assertEqual(stats.unassigned_count, 0) + + def test_random_strategy_is_deterministic_given_seed(self) -> None: + raw_rows = [RawSidRow(f"item_{i}", f"item_{i}", "0,0,0") for i in range(4)] + kwargs = dict(capacity=1, strategy="random", random_last_layer_size=64, seed=11) + first, _ = assign_sid_collisions(raw_rows, [], **kwargs) + second, _ = assign_sid_collisions(raw_rows, [], **kwargs) + self.assertEqual( + {r.item_key: r.codebook for r in first}, + {r.item_key: r.codebook for r in second}, + ) + + def test_random_strategy_requires_last_layer_size(self) -> None: + raw_rows = [RawSidRow("item_0", "item_0", "1,2,3")] + with self.assertRaisesRegex(ValueError, "random_last_layer_size"): + assign_sid_collisions(raw_rows, [], capacity=1, strategy="random") + + def test_missing_candidates_errors_on_overflow(self) -> None: + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + ] + with self.assertRaisesRegex(ValueError, "no explicit candidate input"): + assign_sid_collisions(raw_rows, [], capacity=1) + + def test_duplicate_candidates_do_not_consume_capacity_twice(self) -> None: + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_2", "item_2", "A"), + ] + candidate_rows = [ + CandidateSidRow("item_0", "C", 1, 0.1), + CandidateSidRow("item_0", "C", 2, 0.2), + CandidateSidRow("item_1", "C", 1, 0.1), + CandidateSidRow("item_2", "C", 1, 0.1), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, + candidate_rows, + capacity=2, + seed=7, + ) + + self.assertEqual(len(assigned), 3) + self.assertEqual(stats.unassigned_count, 0) + self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 2) + + def test_local_csv_outputs_codebooks_as_strings(self) -> None: + raw_path = os.path.join(self.test_dir, "raw.csv") + cand_path = os.path.join(self.test_dir, "cand.csv") + out_dir = os.path.join(self.test_dir, "out") + csv.write_csv( + pa.table( + { + "item_id": ["1", "2", "3"], + "codes": ["A", "A", "A"], + } + ), + raw_path, + ) + csv.write_csv( + pa.table( + { + "item_id": ["1", "2", "3"], + "candidate_codebook": ["C", "C", "C"], + "priority": [1, 1, 1], + "score": [0.1, 0.1, 0.1], + } + ), + cand_path, + ) + + args = build_parser().parse_args( + [ + "--input_path", + raw_path, + "--candidate_input_path", + cand_path, + "--output_path", + out_dir, + "--reader_type", + "CsvReader", + "--writer_type", + "CsvWriter", + "--max_items_per_codebook", + "2", + ] + ) + stats = run_local(args) + + self.assertEqual(stats.reassigned_count, 1) + result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) + self.assertEqual(result.schema.field("origin_codebook").type, pa.string()) + self.assertEqual(result.schema.field("codebook").type, pa.string()) + self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) + + def test_local_parquet_accepts_list_codes(self) -> None: + raw_path = os.path.join(self.test_dir, "raw.parquet") + cand_path = os.path.join(self.test_dir, "cand.parquet") + out_dir = os.path.join(self.test_dir, "out_parquet") + parquet.write_table( + pa.table( + { + "item_id": pa.array([1, 2, 3], type=pa.int64()), + "codes": pa.array([[1, 2], [1, 2], [1, 2]]), + } + ), + raw_path, + ) + parquet.write_table( + pa.table( + { + "item_id": pa.array([1, 2, 3], type=pa.int64()), + "candidate_codebook": ["1,3", "1,3", "1,3"], + "priority": [1, 1, 1], + "score": [0.1, 0.1, 0.1], + } + ), + cand_path, + ) + + args = build_parser().parse_args( + [ + "--input_path", + raw_path, + "--candidate_input_path", + cand_path, + "--output_path", + out_dir, + "--writer_type", + "ParquetWriter", + "--max_items_per_codebook", + "2", + ] + ) + stats = run_local(args) + + self.assertEqual(stats.reassigned_count, 1) + result = parquet.read_table(os.path.join(out_dir, "part-0.parquet")) + self.assertIn("1,2", set(result["origin_codebook"].to_pylist())) + self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) + + def test_local_csv_accepts_split_codes_and_compact_candidates(self) -> None: + raw_path = os.path.join(self.test_dir, "raw_split.csv") + cand_path = os.path.join(self.test_dir, "cand_compact.csv") + out_dir = os.path.join(self.test_dir, "out_split") + csv.write_csv( + pa.table( + { + "item_id": ["1", "2", "3"], + "code_0": ["A", "A", "A"], + "code_1": ["B", "B", "B"], + } + ), + raw_path, + ) + csv.write_csv( + pa.table( + { + "item_id": ["1", "2", "3"], + "sorted_index": ["A|C", "A|C", "A|C"], + } + ), + cand_path, + ) + + args = build_parser().parse_args( + [ + "--input_path", + raw_path, + "--candidate_input_path", + cand_path, + "--output_path", + out_dir, + "--reader_type", + "CsvReader", + "--writer_type", + "CsvWriter", + "--code_field", + "", + "--code_fields", + "code_0,code_1", + "--compact_candidate_field", + "sorted_index", + "--max_items_per_codebook", + "2", + ] + ) + stats = run_local(args) + + self.assertEqual(stats.reassigned_count, 1) + result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) + self.assertIn("A,B", set(result["origin_codebook"].to_pylist())) + self.assertIn("A", set(result["codebook"].to_pylist())) + self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) + + def test_generate_odps_sql_uses_tool_capacity_and_no_random(self) -> None: + args = build_parser().parse_args( + [ + "--backend", + "odps", + "--input_path", + "odps://proj/tables/raw_sid/ds=20260630", + "--candidate_input_path", + "odps://proj/tables/cand_sid/ds=20260630", + "--output_path", + "odps://proj/tables/final_sid/ds=20260630", + "--max_items_per_codebook", + "5", + "--max_iters", + "1", + "--dry_run_sql", + ] + ) + sql = "\n".join(generate_odps_sql(args)) + self.assertIn("<= 5", sql) + self.assertIn("PARTITION (ds='20260630')", sql) + self.assertIn("INNER JOIN", sql) + self.assertIn("FROM proj.raw_sid", sql) + self.assertIn("-- final: remaining_unassigned", sql) + self.assertNotIn("rand()", sql.lower()) + self.assertNotIn("last_layer_random", sql) + + def test_generate_odps_sql_keep_original_policy(self) -> None: + args = build_parser().parse_args( + [ + "--backend", + "odps", + "--input_path", + "odps://proj/tables/raw_sid/ds=20260630", + "--candidate_input_path", + "odps://proj/tables/cand_sid/ds=20260630", + "--output_path", + "odps://proj/tables/final_sid/ds=20260630", + "--max_items_per_codebook", + "5", + "--max_iters", + "1", + "--unassigned_policy", + "keep_original", + ] + ) + sql = "\n".join(generate_odps_sql(args)) + + self.assertIn("-- final: remaining_unassigned", sql) + self.assertIn("u.origin_codebook AS codebook", sql) + self.assertIn("COALESCE(cnt.cnt, 0) + u.rn AS `index`", sql) + self.assertIn("WHERE a.item_id IS NULL AND ds='20260630'", sql) + + def test_run_odps_errors_on_remaining_unassigned(self) -> None: + executed_sqls = [] + + class FakeReader: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def __iter__(self): + return iter([(1,)]) + + class FakeInstance: + def __init__(self, sql): + self.sql = sql + + def wait_for_success(self): + return None + + def open_reader(self): + return FakeReader() + + class FakeODPS: + def __init__(self, **kwargs): + pass + + def execute_sql(self, sql): + executed_sqls.append(sql) + return FakeInstance(sql) + + args = build_parser().parse_args( + [ + "--backend", + "odps", + "--input_path", + "odps://proj/tables/raw_sid/ds=20260630", + "--candidate_input_path", + "odps://proj/tables/cand_sid/ds=20260630", + "--output_path", + "odps://proj/tables/final_sid/ds=20260630", + "--max_items_per_codebook", + "5", + "--max_iters", + "0", + ] + ) + + fake_odps_module = types.SimpleNamespace(ODPS=FakeODPS) + with mock.patch.dict(sys.modules, {"odps": fake_odps_module}): + with mock.patch( + "tzrec.datasets.odps_dataset._create_odps_account", + return_value=("account", "endpoint"), + ): + with self.assertRaisesRegex(RuntimeError, "could not be assigned"): + run_odps(args) + + executed_sql = "\n".join(executed_sqls) + self.assertIn("-- final: remaining_unassigned", executed_sql) + self.assertNotIn("INSERT OVERWRITE TABLE proj.final_sid", executed_sql) + + def test_odps_insert_selected_ranks_survivors_densely(self) -> None: + # Regression guard for the duplicate-(codebook, index) bug: the dense + # `codebook_rn` that feeds `index` must be computed AFTER the + # `WHERE item_rn = 1` survivor filter, otherwise dropped rows leave index + # gaps that reappear next iteration as duplicate slots. In the fixed + # nesting the `codebook_rn` window is defined in the outer query (over the + # already-filtered survivors) and therefore textually precedes the inner + # `item_rn` window; the buggy version co-located them with `item_rn` first. + args = build_parser().parse_args( + [ + "--backend", + "odps", + "--input_path", + "odps://proj/tables/raw_sid/ds=20260630", + "--candidate_input_path", + "odps://proj/tables/cand_sid/ds=20260630", + "--output_path", + "odps://proj/tables/final_sid/ds=20260630", + "--max_items_per_codebook", + "5", + "--max_iters", + "3", + ] + ) + sql = OdpsSqlGenerator(args)._insert_selected_sql() + + self.assertIn("current_cnt + codebook_rn AS `index`", sql) + self.assertIn("WHERE item_rn = 1", sql) + self.assertLess(sql.index("AS codebook_rn"), sql.index("AS item_rn")) + + def test_local_reassignment_indices_are_unique_and_dense(self) -> None: + # After reassignment every (codebook, index) pair must be unique and each + # bucket's indices must form a contiguous 1..N run -- the invariant the + # ODPS backend must also uphold. (Which two items overflow A is decided by + # a seeded hash, so every item carries the same B fallback.) + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_2", "item_2", "A"), + RawSidRow("item_3", "item_3", "A"), + ] + candidate_rows = [ + CandidateSidRow("item_0", "B", 1, 0.1), + CandidateSidRow("item_1", "B", 1, 0.2), + CandidateSidRow("item_2", "B", 1, 0.3), + CandidateSidRow("item_3", "B", 1, 0.4), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, + candidate_rows, + capacity=2, + seed=7, + ) + + self.assertEqual(len(assigned), 4) + self.assertEqual(stats.unassigned_count, 0) + pairs = [(row.codebook, row.index) for row in assigned] + self.assertEqual(len(pairs), len(set(pairs))) + by_codebook = defaultdict(list) + for codebook, index in pairs: + by_codebook[codebook].append(index) + for indices in by_codebook.values(): + self.assertEqual(sorted(indices), list(range(1, len(indices) + 1))) + + def test_local_drop_policy_omits_unplaceable_items(self) -> None: + # A (capacity 1) keeps one item; the other two overflow and both want B, + # which fits only one -> one item stays unplaceable and reaches + # _handle_unassigned's drop branch. + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_2", "item_2", "A"), + ] + candidate_rows = [ + CandidateSidRow("item_0", "B", 1, 0.1), + CandidateSidRow("item_1", "B", 1, 0.2), + CandidateSidRow("item_2", "B", 1, 0.3), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, + candidate_rows, + capacity=1, + seed=7, + unassigned_policy="drop", + ) + + self.assertEqual(stats.unassigned_count, 1) + # The unplaceable item is dropped: A keeps 1, B keeps 1, nothing else. + self.assertEqual(len(assigned), 2) + self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 1) + + def test_local_keep_original_readds_over_capacity(self) -> None: + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_2", "item_2", "A"), + ] + candidate_rows = [ + CandidateSidRow("item_0", "B", 1, 0.1), + CandidateSidRow("item_1", "B", 1, 0.2), + CandidateSidRow("item_2", "B", 1, 0.3), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, + candidate_rows, + capacity=1, + seed=7, + unassigned_policy="keep_original", + ) + + self.assertEqual(stats.unassigned_count, 0) + self.assertEqual(len(assigned), 3) + # The unplaceable item is re-added at its origin: keep_original is the + # only policy allowed to exceed capacity (A now holds 2 > capacity 1). + self.assertEqual(Counter(row.codebook for row in assigned)["A"], 2) + self.assertEqual(stats.max_final_bucket_size, 2) + + def test_local_error_policy_raises_on_unplaceable(self) -> None: + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_2", "item_2", "A"), + ] + candidate_rows = [ + CandidateSidRow("item_0", "B", 1, 0.1), + CandidateSidRow("item_1", "B", 1, 0.2), + CandidateSidRow("item_2", "B", 1, 0.3), + ] + + # Distinct from the "no explicit candidate input" ValueError: candidates + # are provided but cannot place every overflow item. + with self.assertRaisesRegex(RuntimeError, "could not be assigned"): + assign_sid_collisions( + raw_rows, + candidate_rows, + capacity=1, + seed=7, + unassigned_policy="error", + ) + + def test_local_score_order_higher_prefers_high_score(self) -> None: + # One item overflows A (capacity 1); it can go to B (score 0.1) or + # C (score 0.9). score_order flips which score wins. Both items carry both + # candidates so the choice is independent of which item the seeded hash + # picks to overflow. + raw_rows = [ + RawSidRow("item_0", "item_0", "A"), + RawSidRow("item_1", "item_1", "A"), + ] + candidate_rows = [ + CandidateSidRow("item_0", "B", 1, 0.1), + CandidateSidRow("item_0", "C", 1, 0.9), + CandidateSidRow("item_1", "B", 1, 0.1), + CandidateSidRow("item_1", "C", 1, 0.9), + ] + + lower, _ = assign_sid_collisions( + raw_rows, candidate_rows, capacity=1, seed=7, score_order="lower" + ) + higher, _ = assign_sid_collisions( + raw_rows, candidate_rows, capacity=1, seed=7, score_order="higher" + ) + + reassigned_lower = next( + row for row in lower if row.origin_codebook != row.codebook + ) + reassigned_higher = next( + row for row in higher if row.origin_codebook != row.codebook + ) + self.assertEqual(reassigned_lower.codebook, "B") + self.assertEqual(reassigned_higher.codebook, "C") + + +if __name__ == "__main__": + unittest.main() From 2a2e4dedc61e89e5f6c6e2087e44c6ac8f7f0fb9 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 8 Jul 2026 08:26:30 +0000 Subject: [PATCH 02/29] [refactor] SID collision tool: use off-the-shelf readers/writers, drop ODPS SQL Follow hitrate.py -- route every backend through create_reader/create_writer (which already speak CSV/Parquet/ODPS) instead of a parallel hand-written ODPS-SQL path. Removes OdpsTableRef, OdpsSqlGenerator (~264 lines of SQL generation), OdpsCollisionRunner, and the generate_odps_sql/run_odps entry points. LocalCollisionRunner -> CollisionRunner (now serves ODPS too via --reader_type/--writer_type OdpsReader/OdpsWriter); run_local -> run; drop the --backend / --temp_prefix / --odps_lifecycle / --dry_run_sql CLI surface. SidCollisionAssigner (the assignment algorithm) is unchanged. Tool 1231 -> 814 lines; tests 588 -> 441 (dropped the 4 ODPS-SQL tests). 14 tests pass; ruff and format clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 455 +------------------ tzrec/tools/sid/collision_prevention_test.py | 164 +------ 2 files changed, 26 insertions(+), 593 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index b4b702a7..64fc73a1 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -77,77 +77,6 @@ class AssignmentStats: max_final_bucket_size: int -@dataclass(frozen=True) -class OdpsTableRef: - """Parsed ODPS table URI.""" - - project: str - table: str - partitions: Tuple[str, ...] - schema: Optional[str] = None - - @classmethod - def parse(cls, path: str) -> "OdpsTableRef": - """Parse an ``odps://project/tables/table[/pt=value]`` path.""" - parts = path.split("/") - if len(parts) < 5 or parts[0] != "odps:" or parts[3] != "tables": - raise ValueError( - f"invalid ODPS path {path!r}; expected " - "odps://project/tables/table[/pt=value]" - ) - table = parts[4] - schema = None - if "." in table: - schema, table = table.split(".", 1) - return cls( - project=parts[2], - table=table, - partitions=tuple(p for p in parts[5:] if p), - schema=schema, - ) - - @property - def table_name(self) -> str: - """Fully qualified ODPS table name.""" - table = f"{self.schema}.{self.table}" if self.schema else self.table - return f"{self.project}.{table}" - - @property - def partition_predicate(self) -> str: - """SQL predicate suffix for this table's partition path.""" - if not self.partitions: - return "" - predicates = [] - for key, value in self._partition_pairs(): - predicates.append(f"{key}='{value}'") - return " AND " + " AND ".join(predicates) - - @property - def insert_target(self) -> str: - """SQL INSERT target including partition spec.""" - if not self.partitions: - return self.table_name - specs = [f"{key}='{value}'" for key, value in self._partition_pairs()] - return f"{self.table_name} PARTITION ({','.join(specs)})" - - @property - def partition_schema(self) -> str: - """CREATE TABLE partition schema suffix.""" - if not self.partitions: - return "" - fields = [f"{key} STRING" for key, _ in self._partition_pairs()] - return f" PARTITIONED BY ({','.join(fields)})" - - def _partition_pairs(self) -> List[Tuple[str, str]]: - pairs = [] - for part in self.partitions: - if "=" not in part: - raise ValueError(f"invalid ODPS partition segment: {part!r}") - key, value = part.split("=", 1) - pairs.append((key, value)) - return pairs - - class SidCollisionAssigner: """Deterministically assign overflow SID rows to explicit candidates.""" @@ -497,14 +426,18 @@ def _final_counts(rows: Sequence[AssignedSidRow]) -> Dict[str, int]: return final_counts -class LocalCollisionRunner: - """CSV/Parquet SID collision-prevention runner.""" +class CollisionRunner: + """SID collision-prevention runner over the standard dataset reader/writer. + + The backend (CSV / Parquet / ODPS) is chosen by the reader/writer type, so + the same path serves local files and MaxCompute tables. + """ def __init__(self, args: argparse.Namespace) -> None: self.args = args def run(self) -> AssignmentStats: - """Run local CSV/Parquet collision prevention.""" + """Read inputs, assign, and write outputs via create_reader/create_writer.""" raw_rows = self._load_raw_sid_rows() candidate_rows = self._load_candidate_rows() rows, stats = SidCollisionAssigner( @@ -758,336 +691,6 @@ def _write_diagnostics(self, stats: AssignmentStats) -> None: writer.close() -class OdpsSqlGenerator: - """Generate deterministic MaxCompute SQL for canonical candidate tables.""" - - def __init__(self, args: argparse.Namespace) -> None: - if not args.candidate_input_path: - raise ValueError("--candidate_input_path is required for --backend odps.") - if args.code_fields: - raise ValueError( - "--backend odps currently supports --code_field, not split code_fields." - ) - if args.compact_candidate_field: - raise ValueError( - "--backend odps expects canonical candidate rows; compact candidates " - "are supported in local CSV/Parquet mode." - ) - - self.args = args - self.raw_ref = OdpsTableRef.parse(args.input_path) - self.candidate_ref = OdpsTableRef.parse(args.candidate_input_path) - self.output_ref = OdpsTableRef.parse(args.output_path) - self.prefix = args.temp_prefix or "tmp_sid_collision" - self.assigned = f"{self.prefix}_assigned" - self.selected = f"{self.prefix}_selected" - self.counts = f"{self.prefix}_counts" - - def generate(self) -> List[str]: - """Generate the SQL statements for one ODPS collision-prevention run.""" - sqls = [ - "SET odps.sql.type.system.odps2=true", - self._create_assigned_sql(), - self._create_selected_sql(), - self._create_counts_sql(), - self._create_output_sql(), - self._initial_assignment_sql(), - ] - for i in range(self.args.max_iters): - sqls.extend( - [ - self._refresh_counts_sql(), - self._select_candidates_sql(), - self._insert_selected_sql(), - self._remaining_unassigned_sql(f"iteration {i + 1}"), - ] - ) - sqls.append(self._remaining_unassigned_sql("final")) - if self.args.unassigned_policy == "keep_original": - sqls.extend([self._refresh_counts_sql(), self._keep_original_sql()]) - sqls.append(self._final_insert_sql()) - return sqls - - @property - def _raw_table(self) -> str: - return self.raw_ref.table_name - - @property - def _candidate_table(self) -> str: - return self.candidate_ref.table_name - - @property - def _raw_predicate(self) -> str: - return self.raw_ref.partition_predicate - - @property - def _candidate_predicate(self) -> str: - return self.candidate_ref.partition_predicate - - @property - def _score_expr(self) -> str: - if self.args.score_field: - return f"CAST({self.args.score_field} AS DOUBLE)" - return "0.0" - - @property - def _priority_expr(self) -> str: - if self.args.priority_field: - return f"CAST({self.args.priority_field} AS BIGINT)" - return "1" - - @property - def _score_order(self) -> str: - return "score ASC" if self.args.score_order == "lower" else "score DESC" - - def _create_assigned_sql(self) -> str: - return ( - f"CREATE TABLE IF NOT EXISTS {self.assigned} (" - "item_id STRING, origin_codebook STRING, codebook STRING, `index` BIGINT" - f") LIFECYCLE {self.args.odps_lifecycle}" - ) - - def _create_selected_sql(self) -> str: - return ( - f"CREATE TABLE IF NOT EXISTS {self.selected} (" - "item_id STRING, origin_codebook STRING, codebook STRING, " - "priority BIGINT, score DOUBLE" - f") LIFECYCLE {self.args.odps_lifecycle}" - ) - - def _create_counts_sql(self) -> str: - return ( - f"CREATE TABLE IF NOT EXISTS {self.counts} (" - "codebook STRING, cnt BIGINT" - f") LIFECYCLE {self.args.odps_lifecycle}" - ) - - def _create_output_sql(self) -> str: - return ( - f"CREATE TABLE IF NOT EXISTS {self.output_ref.table_name} (" - "item_id STRING, origin_codebook STRING, codebook STRING, `index` BIGINT" - f"){self.output_ref.partition_schema}" - ) - - def _initial_assignment_sql(self) -> str: - return ( - f"INSERT OVERWRITE TABLE {self.assigned}\n" - "SELECT item_id, origin_codebook, origin_codebook AS codebook,\n" - " rn AS `index`\n" - "FROM (\n" - f" SELECT CAST({self.args.item_id_field} AS STRING) AS item_id,\n" - f" CAST({self.args.code_field} AS STRING) AS origin_codebook,\n" - " ROW_NUMBER() OVER (\n" - f" PARTITION BY CAST({self.args.code_field} AS STRING)\n" - " ORDER BY ABS(HASH(CONCAT(" - f"'{self.args.seed}', ':', CAST({self.args.item_id_field} AS STRING))))\n" - " ) AS rn\n" - f" FROM {self._raw_table}\n" - f" WHERE 1=1{self._raw_predicate}\n" - ") t\n" - f"WHERE rn <= {self.args.max_items_per_codebook}" - ) - - def _refresh_counts_sql(self) -> str: - return ( - f"INSERT OVERWRITE TABLE {self.counts}\n" - "SELECT codebook, COUNT(*) AS cnt\n" - f"FROM {self.assigned}\n" - "GROUP BY codebook" - ) - - def _select_candidates_sql(self) -> str: - return ( - f"INSERT OVERWRITE TABLE {self.selected}\n" - "SELECT item_id, origin_codebook, codebook, priority, score\n" - "FROM (\n" - " SELECT c.item_id, c.origin_codebook, c.codebook, " - "c.priority, c.score,\n" - " ROW_NUMBER() OVER (\n" - " PARTITION BY c.codebook\n" - f" ORDER BY c.priority ASC, c.{self._score_order}, " - "ABS(HASH(CONCAT(" - f"'{self.args.seed}', ':', c.item_id, ':', c.codebook)))\n" - " ) AS rn,\n" - " COALESCE(cnt.cnt, 0) AS current_cnt\n" - " FROM (\n" - " SELECT CAST(" - f"{self.args.candidate_item_id_field or self.args.item_id_field}" - " AS STRING) AS item_id,\n" - " CAST(" - f"{self.args.candidate_origin_codebook_field} AS STRING" - ") AS origin_codebook,\n" - " CAST(" - f"{self.args.candidate_codebook_field}" - " AS STRING) AS codebook,\n" - f" {self._priority_expr} AS priority,\n" - f" {self._score_expr} AS score\n" - f" FROM {self._candidate_table}\n" - f" WHERE 1=1{self._candidate_predicate}\n" - " ) c\n" - " INNER JOIN (\n" - f" SELECT CAST({self.args.item_id_field} AS STRING) AS item_id\n" - f" FROM {self._raw_table}\n" - f" WHERE 1=1{self._raw_predicate}\n" - " ) r ON c.item_id = r.item_id\n" - f" LEFT OUTER JOIN {self.assigned} a ON c.item_id = a.item_id\n" - f" LEFT OUTER JOIN {self.counts} cnt ON c.codebook = cnt.codebook\n" - " WHERE a.item_id IS NULL\n" - f" AND COALESCE(cnt.cnt, 0) < {self.args.max_items_per_codebook}\n" - ") ranked\n" - f"WHERE rn <= {self.args.max_items_per_codebook} - current_cnt" - ) - - def _insert_selected_sql(self) -> str: - # ``codebook_rn`` must rank ONLY the surviving (``item_rn = 1``) rows so - # the emitted ``index`` stays dense and contiguous with the rows already - # in ``assigned`` (counted by ``current_cnt``). Ranking before the - # ``item_rn = 1`` filter leaves a gap whenever a codebook's top-ranked - # candidate wins a different codebook (its row here is dropped): that gap - # both wastes capacity and is re-issued next iteration as a duplicate - # ``(codebook, index)`` pair, since ``current_cnt`` (a COUNT of inserted - # rows) then undercounts the real slot high-water mark. Hence the nested - # shape: pick each item's best codebook first, THEN densely number the - # survivors per codebook. - return ( - f"INSERT INTO TABLE {self.assigned}\n" - "SELECT item_id, origin_codebook, codebook,\n" - " current_cnt + codebook_rn AS `index`\n" - "FROM (\n" - " SELECT item_id, origin_codebook, codebook, current_cnt,\n" - " ROW_NUMBER() OVER (\n" - " PARTITION BY codebook\n" - f" ORDER BY priority ASC, {self._score_order}, " - "ABS(HASH(CONCAT(" - f"'{self.args.seed}', ':', item_id, ':', codebook)))\n" - " ) AS codebook_rn\n" - " FROM (\n" - " SELECT s.item_id, s.origin_codebook, s.codebook, " - "s.priority, s.score,\n" - " ROW_NUMBER() OVER (\n" - " PARTITION BY s.item_id\n" - f" ORDER BY s.priority ASC, s.{self._score_order}, " - "ABS(HASH(CONCAT(" - f"'{self.args.seed}', ':', s.item_id, ':', s.codebook)))\n" - " ) AS item_rn,\n" - " COALESCE(cnt.cnt, 0) AS current_cnt\n" - f" FROM {self.selected} s\n" - f" LEFT OUTER JOIN {self.counts} cnt " - "ON s.codebook = cnt.codebook\n" - " ) x\n" - " WHERE item_rn = 1\n" - ") y\n" - f"WHERE current_cnt + codebook_rn <= {self.args.max_items_per_codebook}" - ) - - def _remaining_unassigned_sql(self, label: str) -> str: - return ( - f"-- {label}: remaining_unassigned\n" - "SELECT COUNT(*) AS remaining_unassigned\n" - f"FROM {self._raw_table} r\n" - f"LEFT OUTER JOIN {self.assigned} a\n" - f"ON CAST(r.{self.args.item_id_field} AS STRING) = a.item_id\n" - f"WHERE a.item_id IS NULL{self._raw_predicate}" - ) - - def _keep_original_sql(self) -> str: - return ( - f"INSERT INTO TABLE {self.assigned}\n" - "SELECT u.item_id, u.origin_codebook, " - "u.origin_codebook AS codebook,\n" - " COALESCE(cnt.cnt, 0) + u.rn AS `index`\n" - "FROM (\n" - f" SELECT CAST(r.{self.args.item_id_field} AS STRING) AS item_id,\n" - f" CAST(r.{self.args.code_field} AS STRING) AS " - "origin_codebook,\n" - " ROW_NUMBER() OVER (\n" - f" PARTITION BY CAST(r.{self.args.code_field} AS STRING)\n" - " ORDER BY ABS(HASH(CONCAT(" - f"'{self.args.seed}', ':', CAST(r.{self.args.item_id_field} AS STRING))))\n" - " ) AS rn\n" - f" FROM {self._raw_table} r\n" - f" LEFT OUTER JOIN {self.assigned} a\n" - f" ON CAST(r.{self.args.item_id_field} AS STRING) = a.item_id\n" - f" WHERE a.item_id IS NULL{self._raw_predicate}\n" - ") u\n" - f"LEFT OUTER JOIN {self.counts} cnt\n" - "ON u.origin_codebook = cnt.codebook" - ) - - def _final_insert_sql(self) -> str: - return ( - f"INSERT OVERWRITE TABLE {self.output_ref.insert_target}\n" - "SELECT item_id, origin_codebook, codebook, `index`\n" - f"FROM {self.assigned}" - ) - - -class OdpsCollisionRunner: - """ODPS SID collision-prevention runner.""" - - def __init__(self, args: argparse.Namespace) -> None: - self.args = args - - def run(self) -> None: - """Run MaxCompute SQL collision prevention.""" - sqls = OdpsSqlGenerator(self.args).generate() - if self.args.dry_run_sql: - print(";\n\n".join(sqls) + ";") - return - - from odps import ODPS - - from tzrec.datasets.odps_dataset import _create_odps_account - - output_ref = OdpsTableRef.parse(self.args.output_path) - account, endpoint = _create_odps_account() - odps = ODPS(account=account, project=output_ref.project, endpoint=endpoint) - for sql in sqls: - logger.info("Executing ODPS SQL:\n%s", sql) - instance = odps.execute_sql(sql) - instance.wait_for_success() - if not self._is_remaining_unassigned_sql(sql): - continue - - remaining_unassigned = self._read_scalar(instance) - logger.info( - "ODPS remaining unassigned item count: %s", - remaining_unassigned, - ) - if not self._is_final_remaining_unassigned_sql(sql): - continue - self._handle_final_unassigned(remaining_unassigned) - - @staticmethod - def _is_remaining_unassigned_sql(sql: str) -> bool: - return " AS remaining_unassigned" in sql - - @staticmethod - def _is_final_remaining_unassigned_sql(sql: str) -> bool: - return sql.lstrip().startswith("-- final:") - - @staticmethod - def _read_scalar(instance: Any) -> int: - with instance.open_reader() as reader: - for record in reader: - return int(record[0]) - return 0 - - def _handle_final_unassigned(self, remaining_unassigned: int) -> None: - if remaining_unassigned == 0: - return - if self.args.unassigned_policy == "error": - raise RuntimeError( - f"{remaining_unassigned} items could not be assigned within " - "capacity in ODPS collision-prevention run." - ) - if self.args.unassigned_policy == "drop": - logger.warning( - "Dropping %s unassigned items because --unassigned_policy=drop.", - remaining_unassigned, - ) - - def assign_sid_collisions( raw_rows: Sequence[RawSidRow], candidate_rows: Sequence[CandidateSidRow], @@ -1115,19 +718,9 @@ def assign_sid_collisions( ).assign(raw_rows, candidate_rows) -def run_local(args: argparse.Namespace) -> AssignmentStats: - """Run local CSV/Parquet collision prevention.""" - return LocalCollisionRunner(args).run() - - -def generate_odps_sql(args: argparse.Namespace) -> List[str]: - """Generate deterministic MaxCompute SQL for canonical candidate tables.""" - return OdpsSqlGenerator(args).generate() - - -def run_odps(args: argparse.Namespace) -> None: - """Run MaxCompute SQL collision prevention.""" - OdpsCollisionRunner(args).run() +def run(args: argparse.Namespace) -> AssignmentStats: + """Run collision prevention over the configured reader/writer backend.""" + return CollisionRunner(args).run() def build_parser() -> argparse.ArgumentParser: @@ -1135,22 +728,23 @@ def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Prevent SID codebook collisions with explicit candidates." ) - parser.add_argument("--backend", choices=["local", "odps"], default="local") parser.add_argument("--input_path", required=True) parser.add_argument("--output_path", required=True) parser.add_argument("--candidate_input_path", default=None) parser.add_argument("--diagnostics_output_path", default=None) parser.add_argument( - "--reader_type", choices=["CsvReader", "ParquetReader"], default=None + "--reader_type", + choices=["CsvReader", "ParquetReader", "OdpsReader"], + default=None, ) parser.add_argument( "--candidate_reader_type", - choices=["CsvReader", "ParquetReader"], + choices=["CsvReader", "ParquetReader", "OdpsReader"], default=None, ) parser.add_argument( "--writer_type", - choices=["CsvWriter", "ParquetWriter"], + choices=["CsvWriter", "ParquetWriter", "OdpsWriter"], default="ParquetWriter", ) parser.add_argument("--batch_size", type=int, default=4096) @@ -1168,7 +762,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--compact_candidate_field", default=None, - help="Compact string/list candidate field for local mode.", + help="Compact string/list candidate field.", ) parser.add_argument( "--candidate_delimiter", @@ -1192,7 +786,7 @@ def build_parser() -> argparse.ArgumentParser: default="candidate", help="Reassignment strategy: 'candidate' uses explicit nearest-neighbor " "candidate rows; 'random' draws random within-band last-layer codes and " - "needs no candidate input (local backend only).", + "needs no candidate input.", ) parser.add_argument( "--random_last_layer_size", @@ -1207,24 +801,13 @@ def build_parser() -> argparse.ArgumentParser: help="Random last-layer codes drawn per overflow item for --strategy random.", ) parser.add_argument("--odps_data_quota_name", default="pay-as-you-go") - parser.add_argument("--temp_prefix", default=None) - parser.add_argument("--odps_lifecycle", type=int, default=7) - parser.add_argument("--dry_run_sql", action="store_true", default=False) return parser def main() -> None: """Command line entrypoint.""" - parser = build_parser() - args = parser.parse_args() - if args.backend == "local": - run_local(args) - else: - if args.strategy == "random": - raise NotImplementedError( - "strategy='random' is only supported for --backend local." - ) - run_odps(args) + args = build_parser().parse_args() + run(args) if __name__ == "__main__": diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 880fb34e..f42db90d 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -11,25 +11,19 @@ import os import shutil -import sys import tempfile -import types import unittest from collections import Counter, defaultdict -from unittest import mock import pyarrow as pa from pyarrow import csv, parquet from tzrec.tools.sid.collision_prevention import ( CandidateSidRow, - OdpsSqlGenerator, RawSidRow, assign_sid_collisions, build_parser, - generate_odps_sql, - run_local, - run_odps, + run, ) @@ -189,7 +183,7 @@ def test_local_csv_outputs_codebooks_as_strings(self) -> None: "2", ] ) - stats = run_local(args) + stats = run(args) self.assertEqual(stats.reassigned_count, 1) result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) @@ -236,7 +230,7 @@ def test_local_parquet_accepts_list_codes(self) -> None: "2", ] ) - stats = run_local(args) + stats = run(args) self.assertEqual(stats.reassigned_count, 1) result = parquet.read_table(os.path.join(out_dir, "part-0.parquet")) @@ -289,7 +283,7 @@ def test_local_csv_accepts_split_codes_and_compact_candidates(self) -> None: "2", ] ) - stats = run_local(args) + stats = run(args) self.assertEqual(stats.reassigned_count, 1) result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) @@ -297,155 +291,11 @@ def test_local_csv_accepts_split_codes_and_compact_candidates(self) -> None: self.assertIn("A", set(result["codebook"].to_pylist())) self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) - def test_generate_odps_sql_uses_tool_capacity_and_no_random(self) -> None: - args = build_parser().parse_args( - [ - "--backend", - "odps", - "--input_path", - "odps://proj/tables/raw_sid/ds=20260630", - "--candidate_input_path", - "odps://proj/tables/cand_sid/ds=20260630", - "--output_path", - "odps://proj/tables/final_sid/ds=20260630", - "--max_items_per_codebook", - "5", - "--max_iters", - "1", - "--dry_run_sql", - ] - ) - sql = "\n".join(generate_odps_sql(args)) - self.assertIn("<= 5", sql) - self.assertIn("PARTITION (ds='20260630')", sql) - self.assertIn("INNER JOIN", sql) - self.assertIn("FROM proj.raw_sid", sql) - self.assertIn("-- final: remaining_unassigned", sql) - self.assertNotIn("rand()", sql.lower()) - self.assertNotIn("last_layer_random", sql) - - def test_generate_odps_sql_keep_original_policy(self) -> None: - args = build_parser().parse_args( - [ - "--backend", - "odps", - "--input_path", - "odps://proj/tables/raw_sid/ds=20260630", - "--candidate_input_path", - "odps://proj/tables/cand_sid/ds=20260630", - "--output_path", - "odps://proj/tables/final_sid/ds=20260630", - "--max_items_per_codebook", - "5", - "--max_iters", - "1", - "--unassigned_policy", - "keep_original", - ] - ) - sql = "\n".join(generate_odps_sql(args)) - - self.assertIn("-- final: remaining_unassigned", sql) - self.assertIn("u.origin_codebook AS codebook", sql) - self.assertIn("COALESCE(cnt.cnt, 0) + u.rn AS `index`", sql) - self.assertIn("WHERE a.item_id IS NULL AND ds='20260630'", sql) - - def test_run_odps_errors_on_remaining_unassigned(self) -> None: - executed_sqls = [] - - class FakeReader: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - return False - - def __iter__(self): - return iter([(1,)]) - - class FakeInstance: - def __init__(self, sql): - self.sql = sql - - def wait_for_success(self): - return None - - def open_reader(self): - return FakeReader() - - class FakeODPS: - def __init__(self, **kwargs): - pass - - def execute_sql(self, sql): - executed_sqls.append(sql) - return FakeInstance(sql) - - args = build_parser().parse_args( - [ - "--backend", - "odps", - "--input_path", - "odps://proj/tables/raw_sid/ds=20260630", - "--candidate_input_path", - "odps://proj/tables/cand_sid/ds=20260630", - "--output_path", - "odps://proj/tables/final_sid/ds=20260630", - "--max_items_per_codebook", - "5", - "--max_iters", - "0", - ] - ) - - fake_odps_module = types.SimpleNamespace(ODPS=FakeODPS) - with mock.patch.dict(sys.modules, {"odps": fake_odps_module}): - with mock.patch( - "tzrec.datasets.odps_dataset._create_odps_account", - return_value=("account", "endpoint"), - ): - with self.assertRaisesRegex(RuntimeError, "could not be assigned"): - run_odps(args) - - executed_sql = "\n".join(executed_sqls) - self.assertIn("-- final: remaining_unassigned", executed_sql) - self.assertNotIn("INSERT OVERWRITE TABLE proj.final_sid", executed_sql) - - def test_odps_insert_selected_ranks_survivors_densely(self) -> None: - # Regression guard for the duplicate-(codebook, index) bug: the dense - # `codebook_rn` that feeds `index` must be computed AFTER the - # `WHERE item_rn = 1` survivor filter, otherwise dropped rows leave index - # gaps that reappear next iteration as duplicate slots. In the fixed - # nesting the `codebook_rn` window is defined in the outer query (over the - # already-filtered survivors) and therefore textually precedes the inner - # `item_rn` window; the buggy version co-located them with `item_rn` first. - args = build_parser().parse_args( - [ - "--backend", - "odps", - "--input_path", - "odps://proj/tables/raw_sid/ds=20260630", - "--candidate_input_path", - "odps://proj/tables/cand_sid/ds=20260630", - "--output_path", - "odps://proj/tables/final_sid/ds=20260630", - "--max_items_per_codebook", - "5", - "--max_iters", - "3", - ] - ) - sql = OdpsSqlGenerator(args)._insert_selected_sql() - - self.assertIn("current_cnt + codebook_rn AS `index`", sql) - self.assertIn("WHERE item_rn = 1", sql) - self.assertLess(sql.index("AS codebook_rn"), sql.index("AS item_rn")) - def test_local_reassignment_indices_are_unique_and_dense(self) -> None: # After reassignment every (codebook, index) pair must be unique and each - # bucket's indices must form a contiguous 1..N run -- the invariant the - # ODPS backend must also uphold. (Which two items overflow A is decided by - # a seeded hash, so every item carries the same B fallback.) + # bucket's indices must form a contiguous 1..N run. (Which two items + # overflow A is decided by a seeded hash, so every item carries the same + # B fallback.) raw_rows = [ RawSidRow("item_0", "item_0", "A"), RawSidRow("item_1", "item_1", "A"), From 96abd9df7032d055cd7fa15a911c5f0b434b140f Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 8 Jul 2026 08:36:54 +0000 Subject: [PATCH 03/29] [chore] SID collision tool: dedup writer boilerplate, drop dead CLI arg - extract CollisionRunner._write_table (create_writer/write/close) so the assignments and diagnostics writers share one path; plain dict instead of OrderedDict (dicts are ordered on the 3.10 target). - drop the unused --candidate_origin_codebook_field CLI arg -- an orphan from the removed ODPS-SQL path, never read by the runner. 14 tests pass; ruff + format clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 92 +++++++++---------------- 1 file changed, 33 insertions(+), 59 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 64fc73a1..b0c7e940 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -21,7 +21,7 @@ import argparse import hashlib import random -from collections import OrderedDict, defaultdict +from collections import defaultdict from dataclasses import dataclass from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple @@ -622,73 +622,48 @@ def _item_id_array(rows: Sequence[AssignedSidRow]) -> pa.Array: return pa.array(values, type=pa.int64()) return pa.array([str(v) for v in values], type=pa.string()) - def _write_assignments(self, rows: Sequence[AssignedSidRow]) -> None: + def _write_table(self, output_path: str, columns: Dict[str, pa.Array]) -> None: writer = create_writer( - self.args.output_path, + output_path, writer_type=self.args.writer_type, quota_name=self.args.odps_data_quota_name, world_size=1, ) - writer.write( - OrderedDict( - [ - ("item_id", self._item_id_array(rows)), - ( - "origin_codebook", - pa.array( - [row.origin_codebook for row in rows], - type=pa.string(), - ), - ), - ( - "codebook", - pa.array([row.codebook for row in rows], type=pa.string()), - ), - ("index", pa.array([row.index for row in rows], type=pa.int64())), - ] - ) - ) + writer.write(columns) writer.close() + def _write_assignments(self, rows: Sequence[AssignedSidRow]) -> None: + self._write_table( + self.args.output_path, + { + "item_id": self._item_id_array(rows), + "origin_codebook": pa.array( + [row.origin_codebook for row in rows], type=pa.string() + ), + "codebook": pa.array([row.codebook for row in rows], type=pa.string()), + "index": pa.array([row.index for row in rows], type=pa.int64()), + }, + ) + def _write_diagnostics(self, stats: AssignmentStats) -> None: - writer = create_writer( + self._write_table( self.args.diagnostics_output_path, - writer_type=self.args.writer_type, - quota_name=self.args.odps_data_quota_name, - world_size=1, + { + "total_items": pa.array([stats.total_items], type=pa.int64()), + "raw_collision_buckets": pa.array( + [stats.raw_collision_buckets], type=pa.int64() + ), + "final_collision_buckets": pa.array( + [stats.final_collision_buckets], type=pa.int64() + ), + "reassigned_count": pa.array([stats.reassigned_count], type=pa.int64()), + "unassigned_count": pa.array([stats.unassigned_count], type=pa.int64()), + "iteration_count": pa.array([stats.iteration_count], type=pa.int64()), + "max_final_bucket_size": pa.array( + [stats.max_final_bucket_size], type=pa.int64() + ), + }, ) - writer.write( - OrderedDict( - [ - ("total_items", pa.array([stats.total_items], type=pa.int64())), - ( - "raw_collision_buckets", - pa.array([stats.raw_collision_buckets], type=pa.int64()), - ), - ( - "final_collision_buckets", - pa.array([stats.final_collision_buckets], type=pa.int64()), - ), - ( - "reassigned_count", - pa.array([stats.reassigned_count], type=pa.int64()), - ), - ( - "unassigned_count", - pa.array([stats.unassigned_count], type=pa.int64()), - ), - ( - "iteration_count", - pa.array([stats.iteration_count], type=pa.int64()), - ), - ( - "max_final_bucket_size", - pa.array([stats.max_final_bucket_size], type=pa.int64()), - ), - ] - ) - ) - writer.close() def assign_sid_collisions( @@ -758,7 +733,6 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--code_delimiter", default=",") parser.add_argument("--candidate_codebook_field", default="candidate_codebook") - parser.add_argument("--candidate_origin_codebook_field", default="origin_codebook") parser.add_argument( "--compact_candidate_field", default=None, From 33778280232205ee2161637d6e0a5f8b4ccaf8e9 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 8 Jul 2026 08:45:48 +0000 Subject: [PATCH 04/29] [perf] SID collision tool: memoize candidate sort key _candidate_sort_key's blake2b tie-breaker is pure but was recomputed on every pass of the up-to-max_iters assignment loop -- in the per-codebook sorts and the two-sided best-candidate comparisons. Cache it on the frozen CandidateSidRow so each unique key is hashed exactly once, cutting the assignment phase from O(max_iters x N) blake2b digests to O(N). Output is identical (same keys, same order). 14 tests pass (incl. the determinism tests); ruff + format clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index b0c7e940..4b23b0a4 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -118,6 +118,9 @@ def __init__( self.code_delimiter = code_delimiter self.random_last_layer_size = random_last_layer_size self.random_num_candidates = random_num_candidates + self._candidate_sort_keys: Dict[ + CandidateSidRow, Tuple[int, float, int, str, str] + ] = {} def assign( self, @@ -264,14 +267,22 @@ def _candidate_sort_key( self, row: CandidateSidRow, ) -> Tuple[int, float, int, str, str]: + # Memoized: this key's blake2b tie-breaker is pure but is compared inside + # the up-to-max_iters assignment loop, so recomputing it would dominate + # the phase. CandidateSidRow is frozen/hashable, so cache on the row. + cached = self._candidate_sort_keys.get(row) + if cached is not None: + return cached score = row.score if self.score_order == "lower" else -row.score - return ( + key = ( row.priority, score, self._stable_hash(self.seed, row.item_key, row.candidate_codebook), row.item_key, row.candidate_codebook, ) + self._candidate_sort_keys[row] = key + return key def _dedup_candidates( self, From a6ced50a14bbfe3081f19e939275836438329139 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 8 Jul 2026 08:55:15 +0000 Subject: [PATCH 05/29] [perf] SID collision tool: project raw-read columns via selected_cols Pass selected_cols=[item_id, *code_fields] into create_reader on the raw SID read, so wide / ODPS source tables no longer decode unused columns (create_reader already forwards selected_cols to the reader; hitrate.py does the same). The candidate read still reads all columns because its priority/score fields are optional and projecting them would fail when absent. 14 tests pass; ruff + format clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 4b23b0a4..8ad8cd9b 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -500,10 +500,12 @@ def _read_batches( self, input_path: str, reader_type: Optional[str], + selected_cols: Optional[List[str]] = None, ) -> Iterable[Dict[str, pa.Array]]: reader = create_reader( input_path=input_path, batch_size=self.args.batch_size, + selected_cols=selected_cols, reader_type=reader_type, quota_name=self.args.odps_data_quota_name, ) @@ -517,7 +519,12 @@ def _load_raw_sid_rows(self) -> List[RawSidRow]: if bool(code_field) == bool(code_fields): raise ValueError("Set exactly one of --code_field or --code_fields.") - for batch in self._read_batches(self.args.input_path, self.args.reader_type): + # Project only the required columns so wide / ODPS source tables don't + # decode unused columns (create_reader forwards selected_cols to the reader). + selected = [self.args.item_id_field, *(code_fields or [code_field])] + for batch in self._read_batches( + self.args.input_path, self.args.reader_type, selected_cols=selected + ): item_ids = self._array_to_pylist(batch, self.args.item_id_field) if code_field: codes = self._array_to_pylist(batch, code_field) From d69d2af14566ded21bb5f82291e8c0c025fb73cc Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 8 Jul 2026 09:01:31 +0000 Subject: [PATCH 06/29] [refactor] SID collision tool: derive writer from reader when unset Mirror hitrate.py -- capture the main-input reader class during the raw read and default --writer_type to it (CsvReader -> CsvWriter, etc.) instead of hard-defaulting to ParquetWriter, so the output backend matches the input's when unspecified. Explicit --writer_type still wins, and ODPS output still resolves to OdpsWriter via create_writer's path detection. Adds a test for the derivation path. 15 tests pass; ruff + format clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 27 +++++++++++++++++--- tzrec/tools/sid/collision_prevention_test.py | 24 +++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 8ad8cd9b..f6680578 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -446,6 +446,9 @@ class CollisionRunner: def __init__(self, args: argparse.Namespace) -> None: self.args = args + # Class name of the main-input reader, captured during the raw read; used + # to derive the writer type when --writer_type is unset. + self._input_reader_cls_name: Optional[str] = None def run(self) -> AssignmentStats: """Read inputs, assign, and write outputs via create_reader/create_writer.""" @@ -501,6 +504,7 @@ def _read_batches( input_path: str, reader_type: Optional[str], selected_cols: Optional[List[str]] = None, + capture_reader_cls: bool = False, ) -> Iterable[Dict[str, pa.Array]]: reader = create_reader( input_path=input_path, @@ -509,6 +513,8 @@ def _read_batches( reader_type=reader_type, quota_name=self.args.odps_data_quota_name, ) + if capture_reader_cls: + self._input_reader_cls_name = reader.__class__.__name__ yield from reader.to_batches() def _load_raw_sid_rows(self) -> List[RawSidRow]: @@ -523,7 +529,10 @@ def _load_raw_sid_rows(self) -> List[RawSidRow]: # decode unused columns (create_reader forwards selected_cols to the reader). selected = [self.args.item_id_field, *(code_fields or [code_field])] for batch in self._read_batches( - self.args.input_path, self.args.reader_type, selected_cols=selected + self.args.input_path, + self.args.reader_type, + selected_cols=selected, + capture_reader_cls=True, ): item_ids = self._array_to_pylist(batch, self.args.item_id_field) if code_field: @@ -640,10 +649,20 @@ def _item_id_array(rows: Sequence[AssignedSidRow]) -> pa.Array: return pa.array(values, type=pa.int64()) return pa.array([str(v) for v in values], type=pa.string()) + def _writer_type(self) -> Optional[str]: + # Derive the writer from the input reader (hitrate.py idiom) when unset, + # so the output backend matches the input's. ODPS output still resolves + # to OdpsWriter via create_writer's path detection regardless. + if self.args.writer_type: + return self.args.writer_type + if self._input_reader_cls_name: + return self._input_reader_cls_name.replace("Reader", "Writer") + return None + def _write_table(self, output_path: str, columns: Dict[str, pa.Array]) -> None: writer = create_writer( output_path, - writer_type=self.args.writer_type, + writer_type=self._writer_type(), quota_name=self.args.odps_data_quota_name, world_size=1, ) @@ -738,7 +757,9 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--writer_type", choices=["CsvWriter", "ParquetWriter", "OdpsWriter"], - default="ParquetWriter", + default=None, + help="Output writer; defaults to matching the input reader " + "(CsvReader -> CsvWriter, etc.).", ) parser.add_argument("--batch_size", type=int, default=4096) parser.add_argument("--item_id_field", default="item_id") diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index f42db90d..f5df7051 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -191,6 +191,30 @@ def test_local_csv_outputs_codebooks_as_strings(self) -> None: self.assertEqual(result.schema.field("codebook").type, pa.string()) self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) + def test_writer_type_defaults_to_matching_the_reader(self) -> None: + # --writer_type unset: the writer is derived from the input reader + # (CsvReader -> CsvWriter), so CSV in yields CSV out. + raw_path = os.path.join(self.test_dir, "raw_derive.csv") + out_dir = os.path.join(self.test_dir, "out_derive") + csv.write_csv( + pa.table({"item_id": ["1", "2", "3"], "codes": ["A", "A", "A"]}), + raw_path, + ) + args = build_parser().parse_args( + [ + "--input_path", + raw_path, + "--output_path", + out_dir, + "--max_items_per_codebook", + "3", + ] + ) + run(args) + # A CSV part file the CSV reader can parse => CsvWriter was derived. + result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) + self.assertEqual(result.num_rows, 3) + def test_local_parquet_accepts_list_codes(self) -> None: raw_path = os.path.join(self.test_dir, "raw.parquet") cand_path = os.path.join(self.test_dir, "cand.parquet") From f9de751ea0bc7a3b4e871f5725e5beaa30624ebd Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 8 Jul 2026 09:24:48 +0000 Subject: [PATCH 07/29] [refactor] SID collision tool: read candidates from the single input table The SID model emits `codes` and the candidate SIDs in the same rows, so drop --candidate_input_path (and the now-redundant --candidate_reader_type / --candidate_item_id_field): raw SIDs and candidates both come from --input_path, read in one pass. Candidates are loaded only for --strategy candidate (random synthesizes its own) and only when the candidate column is present, so a plain SID table still runs (overflow then follows --unassigned_policy). Merges the two loaders into _load_rows with _origin_codes / _candidate_rows helpers; tests use single-table fixtures. 15 tests pass; ruff + format clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 200 ++++++++++--------- tzrec/tools/sid/collision_prevention_test.py | 39 +--- 2 files changed, 109 insertions(+), 130 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index f6680578..d89e3330 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -452,8 +452,7 @@ def __init__(self, args: argparse.Namespace) -> None: def run(self) -> AssignmentStats: """Read inputs, assign, and write outputs via create_reader/create_writer.""" - raw_rows = self._load_raw_sid_rows() - candidate_rows = self._load_candidate_rows() + raw_rows, candidate_rows = self._load_rows() rows, stats = SidCollisionAssigner( capacity=self.args.max_items_per_codebook, max_iters=self.args.max_iters, @@ -517,122 +516,142 @@ def _read_batches( self._input_reader_cls_name = reader.__class__.__name__ yield from reader.to_batches() - def _load_raw_sid_rows(self) -> List[RawSidRow]: - rows: List[RawSidRow] = [] - seen = set() + def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: + """Read raw SID rows and candidate rows from the single input table. + + The model emits ``codes`` and the candidate SIDs in the same rows, so + candidates come from the same table. They are read only for the + ``candidate`` strategy (``random`` synthesizes its own) and only when the + candidate column is present, so a plain SID table still runs (overflow + then follows ``--unassigned_policy``). + """ code_fields = self._code_fields() code_field = None if code_fields else self.args.code_field if bool(code_field) == bool(code_fields): raise ValueError("Set exactly one of --code_field or --code_fields.") - # Project only the required columns so wide / ODPS source tables don't - # decode unused columns (create_reader forwards selected_cols to the reader). - selected = [self.args.item_id_field, *(code_fields or [code_field])] + candidate_field, is_compact = None, False + if self.args.strategy == "candidate": + codebook_field = ( + None + if self.args.compact_candidate_field + else self.args.candidate_codebook_field + ) + if bool(codebook_field) == bool(self.args.compact_candidate_field): + raise ValueError( + "Set exactly one of --candidate_codebook_field or " + "--compact_candidate_field." + ) + is_compact = bool(self.args.compact_candidate_field) + candidate_field = self.args.compact_candidate_field or codebook_field + + # Project to the SID columns only when candidate columns aren't also read + # from the same table (their optional priority/score can't be projected). + selected_cols = ( + [self.args.item_id_field, *(code_fields or [code_field])] + if candidate_field is None + else None + ) + + raw_rows: List[RawSidRow] = [] + candidate_rows: List[CandidateSidRow] = [] + seen: set = set() for batch in self._read_batches( self.args.input_path, self.args.reader_type, - selected_cols=selected, + selected_cols=selected_cols, capture_reader_cls=True, ): item_ids = self._array_to_pylist(batch, self.args.item_id_field) - if code_field: - codes = self._array_to_pylist(batch, code_field) - origin_codes = [ - self._cell_to_code(v, self.args.code_delimiter) for v in codes - ] - else: - assert code_fields is not None - code_columns = [self._array_to_pylist(batch, f) for f in code_fields] - origin_codes = [ - self.args.code_delimiter.join(str(col[i]) for col in code_columns) - for i in range(len(item_ids)) - ] - + origin_codes = self._origin_codes( + batch, code_field, code_fields, len(item_ids) + ) for item_id, origin_codebook in zip(item_ids, origin_codes): item_key = str(item_id) if item_key in seen: - raise ValueError(f"duplicate item_id in raw SID input: {item_key}") + raise ValueError(f"duplicate item_id in SID input: {item_key}") seen.add(item_key) - rows.append( + raw_rows.append( RawSidRow( item_id=item_id, item_key=item_key, origin_codebook=origin_codebook, ) ) + if candidate_field is not None and candidate_field in batch: + candidate_rows.extend(self._candidate_rows(batch, item_ids, is_compact)) - if not rows: - raise ValueError("raw SID input is empty.") - return rows - - def _load_candidate_rows(self) -> List[CandidateSidRow]: - rows: List[CandidateSidRow] = [] - if not self.args.candidate_input_path: - return rows + if not raw_rows: + raise ValueError("SID input is empty.") + return raw_rows, candidate_rows - candidate_codebook_field = ( - None - if self.args.compact_candidate_field - else self.args.candidate_codebook_field - ) - if bool(candidate_codebook_field) == bool(self.args.compact_candidate_field): - raise ValueError( - "Set exactly one of --candidate_codebook_field or " - "--compact_candidate_field." - ) + def _origin_codes( + self, + batch: Dict[str, pa.Array], + code_field: Optional[str], + code_fields: Optional[List[str]], + n: int, + ) -> List[str]: + if code_field: + codes = self._array_to_pylist(batch, code_field) + return [self._cell_to_code(v, self.args.code_delimiter) for v in codes] + assert code_fields is not None + columns = [self._array_to_pylist(batch, f) for f in code_fields] + return [ + self.args.code_delimiter.join(str(col[i]) for col in columns) + for i in range(n) + ] - reader_type = self.args.candidate_reader_type or self.args.reader_type - item_id_field = self.args.candidate_item_id_field or self.args.item_id_field - for batch in self._read_batches(self.args.candidate_input_path, reader_type): - item_ids = self._array_to_pylist(batch, item_id_field) - priorities = ( - self._array_to_pylist(batch, self.args.priority_field) - if self.args.priority_field and self.args.priority_field in batch - else None - ) - scores = ( - self._array_to_pylist(batch, self.args.score_field) - if self.args.score_field and self.args.score_field in batch - else None + def _candidate_rows( + self, + batch: Dict[str, pa.Array], + item_ids: List[Any], + is_compact: bool, + ) -> List[CandidateSidRow]: + rows: List[CandidateSidRow] = [] + if is_compact: + compact_values = self._array_to_pylist( + batch, self.args.compact_candidate_field ) - if candidate_codebook_field: - candidates = self._array_to_pylist(batch, candidate_codebook_field) - for i, (item_id, candidate) in enumerate(zip(item_ids, candidates)): + for item_id, compact_value in zip(item_ids, compact_values): + for priority, candidate in enumerate( + self._split_compact_candidates( + compact_value, self.args.candidate_delimiter + ), + start=1, + ): rows.append( CandidateSidRow( item_key=str(item_id), - candidate_codebook=self._cell_to_code( - candidate, - self.args.code_delimiter, - ), - priority=int(priorities[i]) - if priorities is not None - else 1, - score=float(scores[i]) if scores is not None else 0.0, + candidate_codebook=candidate, + priority=priority, + score=0.0, ) ) - else: - compact_values = self._array_to_pylist( - batch, - self.args.compact_candidate_field, - ) - for item_id, compact_value in zip(item_ids, compact_values): - for priority, candidate in enumerate( - self._split_compact_candidates( - compact_value, - self.args.candidate_delimiter, - ), - start=1, - ): - rows.append( - CandidateSidRow( - item_key=str(item_id), - candidate_codebook=candidate, - priority=priority, - score=0.0, - ) - ) + return rows + priorities = ( + self._array_to_pylist(batch, self.args.priority_field) + if self.args.priority_field and self.args.priority_field in batch + else None + ) + scores = ( + self._array_to_pylist(batch, self.args.score_field) + if self.args.score_field and self.args.score_field in batch + else None + ) + candidates = self._array_to_pylist(batch, self.args.candidate_codebook_field) + for i, (item_id, candidate) in enumerate(zip(item_ids, candidates)): + rows.append( + CandidateSidRow( + item_key=str(item_id), + candidate_codebook=self._cell_to_code( + candidate, self.args.code_delimiter + ), + priority=int(priorities[i]) if priorities is not None else 1, + score=float(scores[i]) if scores is not None else 0.0, + ) + ) return rows def _code_fields(self) -> Optional[List[str]]: @@ -742,18 +761,12 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--input_path", required=True) parser.add_argument("--output_path", required=True) - parser.add_argument("--candidate_input_path", default=None) parser.add_argument("--diagnostics_output_path", default=None) parser.add_argument( "--reader_type", choices=["CsvReader", "ParquetReader", "OdpsReader"], default=None, ) - parser.add_argument( - "--candidate_reader_type", - choices=["CsvReader", "ParquetReader", "OdpsReader"], - default=None, - ) parser.add_argument( "--writer_type", choices=["CsvWriter", "ParquetWriter", "OdpsWriter"], @@ -763,7 +776,6 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--batch_size", type=int, default=4096) parser.add_argument("--item_id_field", default="item_id") - parser.add_argument("--candidate_item_id_field", default=None) parser.add_argument("--code_field", default="codes") parser.add_argument( "--code_fields", diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index f5df7051..db5a2d80 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -144,35 +144,24 @@ def test_duplicate_candidates_do_not_consume_capacity_twice(self) -> None: def test_local_csv_outputs_codebooks_as_strings(self) -> None: raw_path = os.path.join(self.test_dir, "raw.csv") - cand_path = os.path.join(self.test_dir, "cand.csv") out_dir = os.path.join(self.test_dir, "out") csv.write_csv( pa.table( { "item_id": ["1", "2", "3"], "codes": ["A", "A", "A"], - } - ), - raw_path, - ) - csv.write_csv( - pa.table( - { - "item_id": ["1", "2", "3"], "candidate_codebook": ["C", "C", "C"], "priority": [1, 1, 1], "score": [0.1, 0.1, 0.1], } ), - cand_path, + raw_path, ) args = build_parser().parse_args( [ "--input_path", raw_path, - "--candidate_input_path", - cand_path, "--output_path", out_dir, "--reader_type", @@ -217,35 +206,24 @@ def test_writer_type_defaults_to_matching_the_reader(self) -> None: def test_local_parquet_accepts_list_codes(self) -> None: raw_path = os.path.join(self.test_dir, "raw.parquet") - cand_path = os.path.join(self.test_dir, "cand.parquet") out_dir = os.path.join(self.test_dir, "out_parquet") parquet.write_table( pa.table( { "item_id": pa.array([1, 2, 3], type=pa.int64()), "codes": pa.array([[1, 2], [1, 2], [1, 2]]), - } - ), - raw_path, - ) - parquet.write_table( - pa.table( - { - "item_id": pa.array([1, 2, 3], type=pa.int64()), "candidate_codebook": ["1,3", "1,3", "1,3"], "priority": [1, 1, 1], "score": [0.1, 0.1, 0.1], } ), - cand_path, + raw_path, ) args = build_parser().parse_args( [ "--input_path", raw_path, - "--candidate_input_path", - cand_path, "--output_path", out_dir, "--writer_type", @@ -263,7 +241,6 @@ def test_local_parquet_accepts_list_codes(self) -> None: def test_local_csv_accepts_split_codes_and_compact_candidates(self) -> None: raw_path = os.path.join(self.test_dir, "raw_split.csv") - cand_path = os.path.join(self.test_dir, "cand_compact.csv") out_dir = os.path.join(self.test_dir, "out_split") csv.write_csv( pa.table( @@ -271,26 +248,16 @@ def test_local_csv_accepts_split_codes_and_compact_candidates(self) -> None: "item_id": ["1", "2", "3"], "code_0": ["A", "A", "A"], "code_1": ["B", "B", "B"], - } - ), - raw_path, - ) - csv.write_csv( - pa.table( - { - "item_id": ["1", "2", "3"], "sorted_index": ["A|C", "A|C", "A|C"], } ), - cand_path, + raw_path, ) args = build_parser().parse_args( [ "--input_path", raw_path, - "--candidate_input_path", - cand_path, "--output_path", out_dir, "--reader_type", From 5e024e59dcc7315e3bd52f202a7f02839476deda Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 8 Jul 2026 09:29:45 +0000 Subject: [PATCH 08/29] [refactor] SID collision tool: drop --code_fields (single code column only) The SID lands in one `codes` column (a scalar string or a list cell), so remove the split-across-columns --code_fields path along with its _code_fields helper and the now-trivial _origin_codes helper (folded into the read loop). --code_field (default "codes") is the sole origin-SID source. 15 tests pass; ruff + format clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 46 +++----------------- tzrec/tools/sid/collision_prevention_test.py | 13 ++---- 2 files changed, 10 insertions(+), 49 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index d89e3330..d6642a68 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -525,11 +525,6 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: candidate column is present, so a plain SID table still runs (overflow then follows ``--unassigned_policy``). """ - code_fields = self._code_fields() - code_field = None if code_fields else self.args.code_field - if bool(code_field) == bool(code_fields): - raise ValueError("Set exactly one of --code_field or --code_fields.") - candidate_field, is_compact = None, False if self.args.strategy == "candidate": codebook_field = ( @@ -548,7 +543,7 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: # Project to the SID columns only when candidate columns aren't also read # from the same table (their optional priority/score can't be projected). selected_cols = ( - [self.args.item_id_field, *(code_fields or [code_field])] + [self.args.item_id_field, self.args.code_field] if candidate_field is None else None ) @@ -563,10 +558,8 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: capture_reader_cls=True, ): item_ids = self._array_to_pylist(batch, self.args.item_id_field) - origin_codes = self._origin_codes( - batch, code_field, code_fields, len(item_ids) - ) - for item_id, origin_codebook in zip(item_ids, origin_codes): + code_cells = self._array_to_pylist(batch, self.args.code_field) + for item_id, code_cell in zip(item_ids, code_cells): item_key = str(item_id) if item_key in seen: raise ValueError(f"duplicate item_id in SID input: {item_key}") @@ -575,7 +568,9 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: RawSidRow( item_id=item_id, item_key=item_key, - origin_codebook=origin_codebook, + origin_codebook=self._cell_to_code( + code_cell, self.args.code_delimiter + ), ) ) if candidate_field is not None and candidate_field in batch: @@ -585,23 +580,6 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: raise ValueError("SID input is empty.") return raw_rows, candidate_rows - def _origin_codes( - self, - batch: Dict[str, pa.Array], - code_field: Optional[str], - code_fields: Optional[List[str]], - n: int, - ) -> List[str]: - if code_field: - codes = self._array_to_pylist(batch, code_field) - return [self._cell_to_code(v, self.args.code_delimiter) for v in codes] - assert code_fields is not None - columns = [self._array_to_pylist(batch, f) for f in code_fields] - return [ - self.args.code_delimiter.join(str(col[i]) for col in columns) - for i in range(n) - ] - def _candidate_rows( self, batch: Dict[str, pa.Array], @@ -654,13 +632,6 @@ def _candidate_rows( ) return rows - def _code_fields(self) -> Optional[List[str]]: - if not self.args.code_fields: - return None - return [ - field.strip() for field in self.args.code_fields.split(",") if field.strip() - ] - @staticmethod def _item_id_array(rows: Sequence[AssignedSidRow]) -> pa.Array: values = [row.item_id for row in rows] @@ -777,11 +748,6 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--batch_size", type=int, default=4096) parser.add_argument("--item_id_field", default="item_id") parser.add_argument("--code_field", default="codes") - parser.add_argument( - "--code_fields", - default=None, - help="Comma-separated split code fields. Mutually exclusive with code_field.", - ) parser.add_argument("--code_delimiter", default=",") parser.add_argument("--candidate_codebook_field", default="candidate_codebook") parser.add_argument( diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index db5a2d80..88550698 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -239,15 +239,14 @@ def test_local_parquet_accepts_list_codes(self) -> None: self.assertIn("1,2", set(result["origin_codebook"].to_pylist())) self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) - def test_local_csv_accepts_split_codes_and_compact_candidates(self) -> None: - raw_path = os.path.join(self.test_dir, "raw_split.csv") - out_dir = os.path.join(self.test_dir, "out_split") + def test_local_csv_accepts_compact_candidates(self) -> None: + raw_path = os.path.join(self.test_dir, "raw_compact.csv") + out_dir = os.path.join(self.test_dir, "out_compact") csv.write_csv( pa.table( { "item_id": ["1", "2", "3"], - "code_0": ["A", "A", "A"], - "code_1": ["B", "B", "B"], + "codes": ["A,B", "A,B", "A,B"], "sorted_index": ["A|C", "A|C", "A|C"], } ), @@ -264,10 +263,6 @@ def test_local_csv_accepts_split_codes_and_compact_candidates(self) -> None: "CsvReader", "--writer_type", "CsvWriter", - "--code_field", - "", - "--code_fields", - "code_0,code_1", "--compact_candidate_field", "sorted_index", "--max_items_per_codebook", From 20515947017655a645128f70e74edd8a84fc2b8b Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 9 Jul 2026 02:14:33 +0000 Subject: [PATCH 09/29] [perf] SID collision tool: mechanical wins (slots, interning, single-pass stats) Byte-identical, deterministic-output-preserving cleanups from the perf review: - slots=True on the four row dataclasses (~2x per-object memory); item_key becomes a derived @property on RawSidRow/AssignedSidRow (was a duplicated str(item_id) field). - _dedup_candidates filters to overflow_items (only their candidates can ever be used), capping the dedup map + sort-key memo to the overflow fraction. - intern codebook strings (heavy repetition over a small distinct-SID space). - fold the stats passes (raw_collision_buckets from by_origin; reassigned + final counts in one loop); drop the throwaway dup-detection set (build the dict first); in-place assigned.sort(); heapq.nsmallest for the top-k trims. - _write_diagnostics via dataclasses.asdict; assign_sid_collisions collapsed to **kwargs (defaults live only on the class). Verified byte-identical by a 6-scenario CLI harness (compact / 1:1+score / score-order / drop / keep_original / random, with real multi-iteration reassignment) that hashes identically before and after; 15 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 128 ++++++++----------- tzrec/tools/sid/collision_prevention_test.py | 58 ++++----- 2 files changed, 81 insertions(+), 105 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index d6642a68..35b23751 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -20,9 +20,11 @@ import argparse import hashlib +import heapq import random +import sys from collections import defaultdict -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple import pyarrow as pa @@ -34,16 +36,20 @@ from tzrec.utils.logging_util import logger -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class RawSidRow: """One raw item -> SID row.""" item_id: Any - item_key: str origin_codebook: str + @property + def item_key(self) -> str: + """Stable dict/sort key for the item (its id as a string).""" + return str(self.item_id) + -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class CandidateSidRow: """One item -> candidate SID row.""" @@ -53,18 +59,22 @@ class CandidateSidRow: score: float -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class AssignedSidRow: """One final item -> SID assignment row.""" item_id: Any - item_key: str origin_codebook: str codebook: str index: int + @property + def item_key(self) -> str: + """Stable dict/sort key for the item (its id as a string).""" + return str(self.item_id) + -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class AssignmentStats: """Summary statistics for a collision-prevention run.""" @@ -128,10 +138,10 @@ def assign( candidate_rows: Sequence[CandidateSidRow], ) -> Tuple[List[AssignedSidRow], AssignmentStats]: """Assign overflow SID rows to non-full candidate codebooks.""" - if len({row.item_key for row in raw_rows}) != len(raw_rows): + raw_by_item = {row.item_key: row for row in raw_rows} + if len(raw_by_item) != len(raw_rows): raise ValueError("raw_rows contains duplicate item_id values.") - raw_by_item = {row.item_key: row for row in raw_rows} by_origin: Dict[str, List[RawSidRow]] = defaultdict(list) for row in raw_rows: by_origin[row.origin_codebook].append(row) @@ -142,13 +152,12 @@ def assign( for codebook, rows in by_origin.items(): for index, row in enumerate( - sorted(rows, key=self._assignment_sort_key)[: self.capacity], + heapq.nsmallest(self.capacity, rows, key=self._assignment_sort_key), start=1, ): assigned.append( AssignedSidRow( item_id=row.item_id, - item_key=row.item_key, origin_codebook=row.origin_codebook, codebook=codebook, index=index, @@ -167,7 +176,7 @@ def assign( "raw SID input has overflow rows, but no explicit candidate input was " "provided." ) - candidates_by_item = self._dedup_candidates(candidate_rows, raw_by_item) + candidates_by_item = self._dedup_candidates(candidate_rows, overflow_items) iteration_count = self._assign_candidates( raw_by_item, @@ -182,23 +191,30 @@ def assign( assigned_items, code_counts, ) - final_counts = self._final_counts(assigned) + + final_counts: Dict[str, int] = defaultdict(int) + reassigned = 0 + for row in assigned: + final_counts[row.codebook] += 1 + if row.origin_codebook != row.codebook: + reassigned += 1 stats = AssignmentStats( total_items=len(raw_rows), - raw_collision_buckets=self._raw_collision_buckets(raw_rows), + raw_collision_buckets=sum( + 1 for rows in by_origin.values() if len(rows) > self.capacity + ), final_collision_buckets=sum( 1 for count in final_counts.values() if count > self.capacity ), - reassigned_count=sum( - 1 for row in assigned if row.origin_codebook != row.codebook - ), + reassigned_count=reassigned, unassigned_count=( len(unassigned) if self.unassigned_policy != "keep_original" else 0 ), iteration_count=iteration_count, max_final_bucket_size=max(final_counts.values()) if final_counts else 0, ) - return sorted(assigned, key=lambda r: (r.codebook, r.index, r.item_key)), stats + assigned.sort(key=lambda r: (r.codebook, r.index, r.item_key)) + return assigned, stats @staticmethod def _stable_hash(*parts: Any) -> int: @@ -208,12 +224,6 @@ def _stable_hash(*parts: Any) -> int: h.update(b"\x1f") return int.from_bytes(h.digest(), byteorder="big", signed=False) - def _raw_collision_buckets(self, raw_rows: Sequence[RawSidRow]) -> int: - counts: Dict[str, int] = defaultdict(int) - for row in raw_rows: - counts[row.origin_codebook] += 1 - return sum(1 for count in counts.values() if count > self.capacity) - def _generate_random_candidates( self, overflow_items: Iterable[str], @@ -287,11 +297,13 @@ def _candidate_sort_key( def _dedup_candidates( self, candidate_rows: Sequence[CandidateSidRow], - raw_by_item: Dict[str, RawSidRow], + overflow_items: set, ) -> Dict[str, List[CandidateSidRow]]: + # Only overflow items' candidates can ever be used, so filtering here + # keeps the dedup map (and the sort-key memo) to the overflow fraction. dedup_candidates: Dict[Tuple[str, str], CandidateSidRow] = {} for row in candidate_rows: - if row.item_key not in raw_by_item: + if row.item_key not in overflow_items: continue key = (row.item_key, row.candidate_codebook) current = dedup_candidates.get(key) @@ -339,7 +351,6 @@ def _assign_candidates( assigned.append( AssignedSidRow( item_id=raw.item_id, - item_key=raw.item_key, origin_codebook=raw.origin_codebook, codebook=candidate.candidate_codebook, index=code_counts[candidate.candidate_codebook], @@ -381,7 +392,7 @@ def _select_candidates( if remaining <= 0: continue selected_by_codebook.extend( - sorted(rows, key=self._candidate_sort_key)[:remaining] + heapq.nsmallest(remaining, rows, key=self._candidate_sort_key) ) best_by_item: Dict[str, CandidateSidRow] = {} @@ -421,7 +432,6 @@ def _handle_unassigned( assigned.append( AssignedSidRow( item_id=raw.item_id, - item_key=raw.item_key, origin_codebook=raw.origin_codebook, codebook=raw.origin_codebook, index=code_counts[raw.origin_codebook], @@ -429,13 +439,6 @@ def _handle_unassigned( ) return unassigned - @staticmethod - def _final_counts(rows: Sequence[AssignedSidRow]) -> Dict[str, int]: - final_counts: Dict[str, int] = defaultdict(int) - for row in rows: - final_counts[row.codebook] += 1 - return final_counts - class CollisionRunner: """SID collision-prevention runner over the standard dataset reader/writer. @@ -472,11 +475,13 @@ def run(self) -> AssignmentStats: @staticmethod def _cell_to_code(value: Any, delimiter: str) -> str: + # intern: codebook strings repeat heavily over a small distinct-SID space, + # so they are stored once instead of once per cell. if value is None: raise ValueError("SID code value cannot be null.") if isinstance(value, (list, tuple)): - return delimiter.join(str(v) for v in value) - return str(value) + return sys.intern(delimiter.join(str(v) for v in value)) + return sys.intern(str(value)) @classmethod def _split_compact_candidates(cls, value: Any, delimiter: str) -> List[str]: @@ -487,7 +492,7 @@ def _split_compact_candidates(cls, value: Any, delimiter: str) -> List[str]: text = str(value).strip() if not text: return [] - return [part.strip() for part in text.split(delimiter) if part.strip()] + return [sys.intern(p.strip()) for p in text.split(delimiter) if p.strip()] @staticmethod def _array_to_pylist(batch: Dict[str, pa.Array], field: str) -> List[Any]: @@ -567,7 +572,6 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: raw_rows.append( RawSidRow( item_id=item_id, - item_key=item_key, origin_codebook=self._cell_to_code( code_cell, self.args.code_delimiter ), @@ -673,51 +677,23 @@ def _write_assignments(self, rows: Sequence[AssignedSidRow]) -> None: ) def _write_diagnostics(self, stats: AssignmentStats) -> None: + # asdict preserves field-declaration order, so column order is unchanged. self._write_table( self.args.diagnostics_output_path, - { - "total_items": pa.array([stats.total_items], type=pa.int64()), - "raw_collision_buckets": pa.array( - [stats.raw_collision_buckets], type=pa.int64() - ), - "final_collision_buckets": pa.array( - [stats.final_collision_buckets], type=pa.int64() - ), - "reassigned_count": pa.array([stats.reassigned_count], type=pa.int64()), - "unassigned_count": pa.array([stats.unassigned_count], type=pa.int64()), - "iteration_count": pa.array([stats.iteration_count], type=pa.int64()), - "max_final_bucket_size": pa.array( - [stats.max_final_bucket_size], type=pa.int64() - ), - }, + {k: pa.array([v], type=pa.int64()) for k, v in asdict(stats).items()}, ) def assign_sid_collisions( raw_rows: Sequence[RawSidRow], candidate_rows: Sequence[CandidateSidRow], - capacity: int, - max_iters: int = 50, - seed: int = 2026, - score_order: str = "lower", - unassigned_policy: str = "error", - strategy: str = "candidate", - code_delimiter: str = ",", - random_last_layer_size: Optional[int] = None, - random_num_candidates: int = 64, + **kwargs: Any, ) -> Tuple[List[AssignedSidRow], AssignmentStats]: - """Assign overflow SID rows to non-full candidate codebooks.""" - return SidCollisionAssigner( - capacity=capacity, - max_iters=max_iters, - seed=seed, - score_order=score_order, - unassigned_policy=unassigned_policy, - strategy=strategy, - code_delimiter=code_delimiter, - random_last_layer_size=random_last_layer_size, - random_num_candidates=random_num_candidates, - ).assign(raw_rows, candidate_rows) + """Assign overflow SID rows to non-full candidate codebooks. + + Thin functional entry point; ``kwargs`` are ``SidCollisionAssigner`` params. + """ + return SidCollisionAssigner(**kwargs).assign(raw_rows, candidate_rows) def run(args: argparse.Namespace) -> AssignmentStats: diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 88550698..dd2f9e1e 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -39,10 +39,10 @@ def tearDown(self) -> None: def test_assign_sid_collisions_respects_capacity(self) -> None: raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), - RawSidRow("item_2", "item_2", "A"), - RawSidRow("item_3", "item_3", "B"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), + RawSidRow("item_2", "A"), + RawSidRow("item_3", "B"), ] candidate_rows = [ CandidateSidRow("item_0", "C", 1, 0.1), @@ -68,9 +68,9 @@ def test_assign_sid_collisions_respects_capacity(self) -> None: def test_random_strategy_reassigns_within_band_without_candidates(self) -> None: # Three items collide on SID (lv1,lv2,lv3) = (1,2,3); capacity 1. raw_rows = [ - RawSidRow("item_0", "item_0", "1,2,3"), - RawSidRow("item_1", "item_1", "1,2,3"), - RawSidRow("item_2", "item_2", "1,2,3"), + RawSidRow("item_0", "1,2,3"), + RawSidRow("item_1", "1,2,3"), + RawSidRow("item_2", "1,2,3"), ] assigned, stats = assign_sid_collisions( @@ -96,7 +96,7 @@ def test_random_strategy_reassigns_within_band_without_candidates(self) -> None: self.assertEqual(stats.unassigned_count, 0) def test_random_strategy_is_deterministic_given_seed(self) -> None: - raw_rows = [RawSidRow(f"item_{i}", f"item_{i}", "0,0,0") for i in range(4)] + raw_rows = [RawSidRow(f"item_{i}", "0,0,0") for i in range(4)] kwargs = dict(capacity=1, strategy="random", random_last_layer_size=64, seed=11) first, _ = assign_sid_collisions(raw_rows, [], **kwargs) second, _ = assign_sid_collisions(raw_rows, [], **kwargs) @@ -106,23 +106,23 @@ def test_random_strategy_is_deterministic_given_seed(self) -> None: ) def test_random_strategy_requires_last_layer_size(self) -> None: - raw_rows = [RawSidRow("item_0", "item_0", "1,2,3")] + raw_rows = [RawSidRow("item_0", "1,2,3")] with self.assertRaisesRegex(ValueError, "random_last_layer_size"): assign_sid_collisions(raw_rows, [], capacity=1, strategy="random") def test_missing_candidates_errors_on_overflow(self) -> None: raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), ] with self.assertRaisesRegex(ValueError, "no explicit candidate input"): assign_sid_collisions(raw_rows, [], capacity=1) def test_duplicate_candidates_do_not_consume_capacity_twice(self) -> None: raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), - RawSidRow("item_2", "item_2", "A"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), + RawSidRow("item_2", "A"), ] candidate_rows = [ CandidateSidRow("item_0", "C", 1, 0.1), @@ -283,10 +283,10 @@ def test_local_reassignment_indices_are_unique_and_dense(self) -> None: # overflow A is decided by a seeded hash, so every item carries the same # B fallback.) raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), - RawSidRow("item_2", "item_2", "A"), - RawSidRow("item_3", "item_3", "A"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), + RawSidRow("item_2", "A"), + RawSidRow("item_3", "A"), ] candidate_rows = [ CandidateSidRow("item_0", "B", 1, 0.1), @@ -317,9 +317,9 @@ def test_local_drop_policy_omits_unplaceable_items(self) -> None: # which fits only one -> one item stays unplaceable and reaches # _handle_unassigned's drop branch. raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), - RawSidRow("item_2", "item_2", "A"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), + RawSidRow("item_2", "A"), ] candidate_rows = [ CandidateSidRow("item_0", "B", 1, 0.1), @@ -342,9 +342,9 @@ def test_local_drop_policy_omits_unplaceable_items(self) -> None: def test_local_keep_original_readds_over_capacity(self) -> None: raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), - RawSidRow("item_2", "item_2", "A"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), + RawSidRow("item_2", "A"), ] candidate_rows = [ CandidateSidRow("item_0", "B", 1, 0.1), @@ -369,9 +369,9 @@ def test_local_keep_original_readds_over_capacity(self) -> None: def test_local_error_policy_raises_on_unplaceable(self) -> None: raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), - RawSidRow("item_2", "item_2", "A"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), + RawSidRow("item_2", "A"), ] candidate_rows = [ CandidateSidRow("item_0", "B", 1, 0.1), @@ -396,8 +396,8 @@ def test_local_score_order_higher_prefers_high_score(self) -> None: # candidates so the choice is independent of which item the seeded hash # picks to overflow. raw_rows = [ - RawSidRow("item_0", "item_0", "A"), - RawSidRow("item_1", "item_1", "A"), + RawSidRow("item_0", "A"), + RawSidRow("item_1", "A"), ] candidate_rows = [ CandidateSidRow("item_0", "B", 1, 0.1), From d8b70b010532ba658ab2b91ae188490954b877a5 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 9 Jul 2026 02:22:10 +0000 Subject: [PATCH 10/29] [perf] SID collision tool: restructure the assignment loop (sort once) Rework SidCollisionAssigner's reassignment loop so candidates are grouped by codebook and sorted ONCE (the sort key is a total order) instead of re-grouping and re-sorting every codebook on each of up to max_iters passes. Each pass now scans the pre-sorted lists filtered by a live `unassigned` set (discard on assignment), rather than rebuilding `set(raw_by_item) - assigned_items` over all items every pass; _available_candidates is deleted and _handle_unassigned reuses the same set. The assignment phase drops from ~O(passes x N x K) to ~O(one sort + scans). Output is byte-identical: best-per-item and the final accepted sort are order-independent, and unique keys make filter-then-sort == sort-then-filter. Verified by a 6-scenario CLI harness (identical output hash) and 15 unit tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 96 +++++++++++-------------- 1 file changed, 41 insertions(+), 55 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 35b23751..fd616387 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -176,21 +176,17 @@ def assign( "raw SID input has overflow rows, but no explicit candidate input was " "provided." ) - candidates_by_item = self._dedup_candidates(candidate_rows, overflow_items) + sorted_by_codebook = self._dedup_candidates(candidate_rows, overflow_items) + unassigned = set(overflow_items) iteration_count = self._assign_candidates( raw_by_item, - candidates_by_item, + sorted_by_codebook, assigned, - assigned_items, - code_counts, - ) - unassigned = self._handle_unassigned( - raw_by_item, - assigned, - assigned_items, code_counts, + unassigned, ) + self._handle_unassigned(raw_by_item, assigned, code_counts, unassigned) final_counts: Dict[str, int] = defaultdict(int) reassigned = 0 @@ -312,37 +308,36 @@ def _dedup_candidates( ) < self._candidate_sort_key(current): dedup_candidates[key] = row - candidates_by_item: Dict[str, List[CandidateSidRow]] = defaultdict(list) + # Group by codebook and sort each list ONCE. The sort key is a total + # order, so the loop can scan these lists instead of re-sorting per pass. + sorted_by_codebook: Dict[str, List[CandidateSidRow]] = defaultdict(list) for row in dedup_candidates.values(): - candidates_by_item[row.item_key].append(row) - return candidates_by_item + sorted_by_codebook[row.candidate_codebook].append(row) + for rows in sorted_by_codebook.values(): + rows.sort(key=self._candidate_sort_key) + return sorted_by_codebook def _assign_candidates( self, raw_by_item: Dict[str, RawSidRow], - candidates_by_item: Dict[str, List[CandidateSidRow]], + sorted_by_codebook: Dict[str, List[CandidateSidRow]], assigned: List[AssignedSidRow], - assigned_items: set, code_counts: Dict[str, int], + unassigned: set, ) -> int: iteration_count = 0 for iteration in range(self.max_iters): - unassigned = set(raw_by_item) - assigned_items if not unassigned: break - - available = self._available_candidates( - unassigned, - candidates_by_item, - code_counts, + accepted = self._select_candidates( + sorted_by_codebook, unassigned, code_counts ) - if not available: + if not accepted: break - accepted = self._select_candidates(available, code_counts) progress = 0 for candidate in accepted: - if candidate.item_key in assigned_items: + if candidate.item_key not in unassigned: continue if code_counts[candidate.candidate_codebook] >= self.capacity: continue @@ -356,7 +351,7 @@ def _assign_candidates( index=code_counts[candidate.candidate_codebook], ) ) - assigned_items.add(candidate.item_key) + unassigned.discard(candidate.item_key) progress += 1 iteration_count = iteration + 1 @@ -364,36 +359,28 @@ def _assign_candidates( break return iteration_count - def _available_candidates( - self, - unassigned: Sequence[str], - candidates_by_item: Dict[str, List[CandidateSidRow]], - code_counts: Dict[str, int], - ) -> List[CandidateSidRow]: - available = [] - for item_key in unassigned: - for candidate in candidates_by_item.get(item_key, []): - if code_counts[candidate.candidate_codebook] < self.capacity: - available.append(candidate) - return available - def _select_candidates( self, - available: Sequence[CandidateSidRow], + sorted_by_codebook: Dict[str, List[CandidateSidRow]], + unassigned: set, code_counts: Dict[str, int], ) -> List[CandidateSidRow]: + # Per open codebook, take its `remaining` smallest still-unassigned + # candidates from the pre-sorted list; then keep each item's single best + # and return sorted -- identical to the old sort-per-pass path (unique + # keys make sort-then-filter == filter-then-sort). selected_by_codebook: List[CandidateSidRow] = [] - by_codebook: Dict[str, List[CandidateSidRow]] = defaultdict(list) - for candidate in available: - by_codebook[candidate.candidate_codebook].append(candidate) - - for codebook, rows in by_codebook.items(): + for codebook, rows in sorted_by_codebook.items(): remaining = self.capacity - code_counts[codebook] if remaining <= 0: continue - selected_by_codebook.extend( - heapq.nsmallest(remaining, rows, key=self._candidate_sort_key) - ) + taken = 0 + for row in rows: + if row.item_key in unassigned: + selected_by_codebook.append(row) + taken += 1 + if taken == remaining: + break best_by_item: Dict[str, CandidateSidRow] = {} for candidate in selected_by_codebook: @@ -412,21 +399,21 @@ def _handle_unassigned( self, raw_by_item: Dict[str, RawSidRow], assigned: List[AssignedSidRow], - assigned_items: set, code_counts: Dict[str, int], - ) -> List[str]: - unassigned = sorted(set(raw_by_item) - assigned_items) - if not unassigned: - return unassigned + unassigned: set, + ) -> None: + pending = sorted(unassigned) + if not pending: + return if self.unassigned_policy == "error": - preview = ",".join(unassigned[:10]) + preview = ",".join(pending[:10]) raise RuntimeError( - f"{len(unassigned)} items could not be assigned within capacity; " + f"{len(pending)} items could not be assigned within capacity; " f"first unassigned item_ids: {preview}" ) if self.unassigned_policy == "keep_original": - for item_key in unassigned: + for item_key in pending: raw = raw_by_item[item_key] code_counts[raw.origin_codebook] += 1 assigned.append( @@ -437,7 +424,6 @@ def _handle_unassigned( index=code_counts[raw.origin_codebook], ) ) - return unassigned class CollisionRunner: From 90f17d9910edad959e57902ce86121e1daadfe39 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 9 Jul 2026 02:26:14 +0000 Subject: [PATCH 11/29] [perf] SID collision tool: two-pass load (candidates only for overflow items) Split _load_rows into two passes over the single input table: pass 1 reads just [item_id, code] (projected) to build raw rows + per-origin counts; pass 2 reads the candidate columns and materializes a CandidateSidRow only for items in over-capacity buckets -- the only ones that can overflow. Candidate memory now scales with the overflow fraction instead of one row per candidate per item (the difference between OOM and feasible at 10M x 64). Also inlines the single-call _read_batches into _read (single-table -> paths come from self.args) and folds candidate-field resolution into _candidate_field. Output byte-identical: over-capacity buckets are a deterministic superset of the true overflow, which the assigner's dedup already narrows, so the assigner sees the same candidates. Verified by the 6-scenario CLI harness (identical hash) and 15 unit tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 119 +++++++++++++----------- 1 file changed, 65 insertions(+), 54 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index fd616387..488688fd 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -489,64 +489,60 @@ def _array_to_pylist(batch: Dict[str, pa.Array], field: str) -> List[Any]: ) return batch[field].to_pylist() - def _read_batches( - self, - input_path: str, - reader_type: Optional[str], - selected_cols: Optional[List[str]] = None, - capture_reader_cls: bool = False, + def _read( + self, selected_cols: Optional[List[str]], capture_reader_cls: bool = False ) -> Iterable[Dict[str, pa.Array]]: reader = create_reader( - input_path=input_path, + input_path=self.args.input_path, batch_size=self.args.batch_size, selected_cols=selected_cols, - reader_type=reader_type, + reader_type=self.args.reader_type, quota_name=self.args.odps_data_quota_name, ) if capture_reader_cls: self._input_reader_cls_name = reader.__class__.__name__ yield from reader.to_batches() - def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: - """Read raw SID rows and candidate rows from the single input table. + def _candidate_field(self) -> Tuple[Optional[str], bool]: + """Resolve the candidate column and whether it is compact. - The model emits ``codes`` and the candidate SIDs in the same rows, so - candidates come from the same table. They are read only for the - ``candidate`` strategy (``random`` synthesizes its own) and only when the - candidate column is present, so a plain SID table still runs (overflow - then follows ``--unassigned_policy``). + Returns ``(None, False)`` for the ``random`` strategy (it synthesizes its + own candidates and reads no candidate column). """ - candidate_field, is_compact = None, False - if self.args.strategy == "candidate": - codebook_field = ( - None - if self.args.compact_candidate_field - else self.args.candidate_codebook_field + if self.args.strategy != "candidate": + return None, False + codebook_field = ( + None + if self.args.compact_candidate_field + else self.args.candidate_codebook_field + ) + if bool(codebook_field) == bool(self.args.compact_candidate_field): + raise ValueError( + "Set exactly one of --candidate_codebook_field or " + "--compact_candidate_field." ) - if bool(codebook_field) == bool(self.args.compact_candidate_field): - raise ValueError( - "Set exactly one of --candidate_codebook_field or " - "--compact_candidate_field." - ) - is_compact = bool(self.args.compact_candidate_field) - candidate_field = self.args.compact_candidate_field or codebook_field - - # Project to the SID columns only when candidate columns aren't also read - # from the same table (their optional priority/score can't be projected). - selected_cols = ( - [self.args.item_id_field, self.args.code_field] - if candidate_field is None - else None + return ( + self.args.compact_candidate_field or codebook_field, + bool(self.args.compact_candidate_field), ) + def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: + """Read raw SID rows and candidate rows from the single input table. + + Two passes: pass 1 reads only the id + code columns (raw rows + per-origin + counts); pass 2 reads the candidate columns and builds a CandidateSidRow + only for items in over-capacity buckets -- the only ones that can overflow + -- so candidate memory scales with the overflow fraction, not the whole + table. Candidates are read only for the ``candidate`` strategy and only + when the column is present, so a plain SID table still runs. + """ + candidate_field, is_compact = self._candidate_field() + raw_rows: List[RawSidRow] = [] - candidate_rows: List[CandidateSidRow] = [] + origin_counts: Dict[str, int] = defaultdict(int) seen: set = set() - for batch in self._read_batches( - self.args.input_path, - self.args.reader_type, - selected_cols=selected_cols, - capture_reader_cls=True, + for batch in self._read( + [self.args.item_id_field, self.args.code_field], capture_reader_cls=True ): item_ids = self._array_to_pylist(batch, self.args.item_id_field) code_cells = self._array_to_pylist(batch, self.args.code_field) @@ -555,19 +551,27 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: if item_key in seen: raise ValueError(f"duplicate item_id in SID input: {item_key}") seen.add(item_key) - raw_rows.append( - RawSidRow( - item_id=item_id, - origin_codebook=self._cell_to_code( - code_cell, self.args.code_delimiter - ), - ) - ) - if candidate_field is not None and candidate_field in batch: - candidate_rows.extend(self._candidate_rows(batch, item_ids, is_compact)) - + origin = self._cell_to_code(code_cell, self.args.code_delimiter) + raw_rows.append(RawSidRow(item_id=item_id, origin_codebook=origin)) + origin_counts[origin] += 1 if not raw_rows: raise ValueError("SID input is empty.") + + cap = self.args.max_items_per_codebook + overflow = { + r.item_key for r in raw_rows if origin_counts[r.origin_codebook] > cap + } + if candidate_field is None or not overflow: + return raw_rows, [] + + candidate_rows: List[CandidateSidRow] = [] + for batch in self._read(None): + if candidate_field not in batch: + break + item_ids = self._array_to_pylist(batch, self.args.item_id_field) + candidate_rows.extend( + self._candidate_rows(batch, item_ids, is_compact, overflow) + ) return raw_rows, candidate_rows def _candidate_rows( @@ -575,6 +579,7 @@ def _candidate_rows( batch: Dict[str, pa.Array], item_ids: List[Any], is_compact: bool, + overflow: set, ) -> List[CandidateSidRow]: rows: List[CandidateSidRow] = [] if is_compact: @@ -582,6 +587,9 @@ def _candidate_rows( batch, self.args.compact_candidate_field ) for item_id, compact_value in zip(item_ids, compact_values): + item_key = str(item_id) + if item_key not in overflow: + continue for priority, candidate in enumerate( self._split_compact_candidates( compact_value, self.args.candidate_delimiter @@ -590,7 +598,7 @@ def _candidate_rows( ): rows.append( CandidateSidRow( - item_key=str(item_id), + item_key=item_key, candidate_codebook=candidate, priority=priority, score=0.0, @@ -610,9 +618,12 @@ def _candidate_rows( ) candidates = self._array_to_pylist(batch, self.args.candidate_codebook_field) for i, (item_id, candidate) in enumerate(zip(item_ids, candidates)): + item_key = str(item_id) + if item_key not in overflow: + continue rows.append( CandidateSidRow( - item_key=str(item_id), + item_key=item_key, candidate_codebook=self._cell_to_code( candidate, self.args.code_delimiter ), From 5d8f95f6d0a69fc199804fc641a2b61d872995be Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 9 Jul 2026 02:47:45 +0000 Subject: [PATCH 12/29] [feat] SID collision tool: tuple[int] bucket key + list I/O (CSV compat) Represent SID codebooks as tuple[int] (the codes are bigints) rather than a delimited string. The tool uses a codebook purely as an opaque hashable, ordered bucket identity, so a tuple is the faithful form. I/O: - Parquet/ODPS (primary): read/write list columns directly -- codes, candidate_codebook, and the compact-candidate field (list>); origin_codebook/codebook are written as list. - CSV (fallback): a compatibility layer parses delimited-int strings on read and joins codes with --code_delimiter on write (chosen by the writer type). Drops the string join/parse/intern path, makes the tie-break numeric (natural SID order), and hands downstream the integer codes with no re-parse; random candidate generation is now pure tuple ops. Determinism preserved. Verified that CSV and Parquet with the same logical data normalize to identical tuples and produce identical assignment decisions (a unit test plus a 3000-item cross-backend check); 16 tests pass; ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 100 ++-- tzrec/tools/sid/collision_prevention_test.py | 505 ++++++++++--------- 2 files changed, 344 insertions(+), 261 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 488688fd..f155dda0 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -22,7 +22,6 @@ import hashlib import heapq import random -import sys from collections import defaultdict from dataclasses import asdict, dataclass from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple @@ -35,13 +34,18 @@ from tzrec.datasets.dataset import create_reader, create_writer from tzrec.utils.logging_util import logger +# A SID bucket key: the tuple of per-layer integer codes. Used only as an opaque +# hashable, ordered identity (grouped / counted / compared), never for arithmetic. +# list on disk for Parquet/ODPS; a delimited string for the CSV fallback. +Codebook = Tuple[int, ...] + @dataclass(frozen=True, slots=True) class RawSidRow: """One raw item -> SID row.""" item_id: Any - origin_codebook: str + origin_codebook: Codebook @property def item_key(self) -> str: @@ -54,7 +58,7 @@ class CandidateSidRow: """One item -> candidate SID row.""" item_key: str - candidate_codebook: str + candidate_codebook: Codebook priority: int score: float @@ -64,8 +68,8 @@ class AssignedSidRow: """One final item -> SID assignment row.""" item_id: Any - origin_codebook: str - codebook: str + origin_codebook: Codebook + codebook: Codebook index: int @property @@ -129,7 +133,7 @@ def __init__( self.random_last_layer_size = random_last_layer_size self.random_num_candidates = random_num_candidates self._candidate_sort_keys: Dict[ - CandidateSidRow, Tuple[int, float, int, str, str] + CandidateSidRow, Tuple[int, float, int, str, Codebook] ] = {} def assign( @@ -238,12 +242,8 @@ def _generate_random_candidates( rows: List[CandidateSidRow] = [] for item_key in overflow_items: raw = raw_by_item[item_key] - parts = raw.origin_codebook.split(self.code_delimiter) - prefix = parts[:-1] - try: - origin_last: Optional[int] = int(parts[-1]) - except ValueError: - origin_last = None + prefix = raw.origin_codebook[:-1] + origin_last = raw.origin_codebook[-1] if raw.origin_codebook else None rng = random.Random(self._stable_hash(self.seed, item_key)) drawn: set = set() while len(drawn) < num_draws: @@ -254,9 +254,7 @@ def _generate_random_candidates( rows.append( CandidateSidRow( item_key=item_key, - candidate_codebook=self.code_delimiter.join( - [*prefix, str(value)] - ), + candidate_codebook=(*prefix, value), priority=1, score=0.0, ) @@ -272,7 +270,7 @@ def _assignment_sort_key(self, row: RawSidRow) -> Tuple[int, str]: def _candidate_sort_key( self, row: CandidateSidRow, - ) -> Tuple[int, float, int, str, str]: + ) -> Tuple[int, float, int, str, Codebook]: # Memoized: this key's blake2b tie-breaker is pure but is compared inside # the up-to-max_iters assignment loop, so recomputing it would dominate # the phase. CandidateSidRow is frozen/hashable, so cache on the row. @@ -438,6 +436,8 @@ def __init__(self, args: argparse.Namespace) -> None: # Class name of the main-input reader, captured during the raw read; used # to derive the writer type when --writer_type is unset. self._input_reader_cls_name: Optional[str] = None + # Dedup identical codebook tuples so a repeated SID is one object. + self._cb_cache: Dict[Codebook, Codebook] = {} def run(self) -> AssignmentStats: """Read inputs, assign, and write outputs via create_reader/create_writer.""" @@ -459,26 +459,41 @@ def run(self) -> AssignmentStats: logger.info("SID collision prevention finished: %s", stats) return stats - @staticmethod - def _cell_to_code(value: Any, delimiter: str) -> str: - # intern: codebook strings repeat heavily over a small distinct-SID space, - # so they are stored once instead of once per cell. + def _to_codebook(self, value: Any, delimiter: str) -> Codebook: + """Normalize a SID cell to a tuple of int codes (deduped). + + Parquet/ODPS give a list cell; CSV gives a delimited string. + """ if value is None: raise ValueError("SID code value cannot be null.") if isinstance(value, (list, tuple)): - return sys.intern(delimiter.join(str(v) for v in value)) - return sys.intern(str(value)) + codes: Codebook = tuple(int(v) for v in value) + else: + codes = tuple(int(p) for p in str(value).split(delimiter) if p.strip()) + return self._cb_cache.setdefault(codes, codes) + + def _split_compact(self, value: Any) -> List[Codebook]: + """Split a compact candidate cell into candidate codebooks. - @classmethod - def _split_compact_candidates(cls, value: Any, delimiter: str) -> List[str]: + Parquet/ODPS give a list> cell; CSV gives candidate SIDs + joined by --candidate_delimiter (e.g. ``"1,3|4,5"``). + """ if value is None: return [] if isinstance(value, (list, tuple)): - return [cls._cell_to_code(v, ",") for v in value if v is not None] + return [ + self._to_codebook(inner, self.args.code_delimiter) + for inner in value + if inner is not None + ] text = str(value).strip() if not text: return [] - return [sys.intern(p.strip()) for p in text.split(delimiter) if p.strip()] + return [ + self._to_codebook(part, self.args.code_delimiter) + for part in text.split(self.args.candidate_delimiter) + if part.strip() + ] @staticmethod def _array_to_pylist(batch: Dict[str, pa.Array], field: str) -> List[Any]: @@ -551,7 +566,7 @@ def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: if item_key in seen: raise ValueError(f"duplicate item_id in SID input: {item_key}") seen.add(item_key) - origin = self._cell_to_code(code_cell, self.args.code_delimiter) + origin = self._to_codebook(code_cell, self.args.code_delimiter) raw_rows.append(RawSidRow(item_id=item_id, origin_codebook=origin)) origin_counts[origin] += 1 if not raw_rows: @@ -591,9 +606,7 @@ def _candidate_rows( if item_key not in overflow: continue for priority, candidate in enumerate( - self._split_compact_candidates( - compact_value, self.args.candidate_delimiter - ), + self._split_compact(compact_value), start=1, ): rows.append( @@ -624,7 +637,7 @@ def _candidate_rows( rows.append( CandidateSidRow( item_key=item_key, - candidate_codebook=self._cell_to_code( + candidate_codebook=self._to_codebook( candidate, self.args.code_delimiter ), priority=int(priorities[i]) if priorities is not None else 1, @@ -660,15 +673,27 @@ def _write_table(self, output_path: str, columns: Dict[str, pa.Array]) -> None: writer.write(columns) writer.close() + def _is_csv_output(self) -> bool: + return self._writer_type() == "CsvWriter" + + def _codebook_array(self, codebooks: List[Codebook]) -> pa.Array: + # list for Parquet/ODPS; a --code_delimiter-joined string for CSV. + if self._is_csv_output(): + d = self.args.code_delimiter + return pa.array( + [d.join(str(c) for c in cb) for cb in codebooks], type=pa.string() + ) + return pa.array([list(cb) for cb in codebooks], type=pa.list_(pa.int64())) + def _write_assignments(self, rows: Sequence[AssignedSidRow]) -> None: self._write_table( self.args.output_path, { "item_id": self._item_id_array(rows), - "origin_codebook": pa.array( - [row.origin_codebook for row in rows], type=pa.string() + "origin_codebook": self._codebook_array( + [row.origin_codebook for row in rows] ), - "codebook": pa.array([row.codebook for row in rows], type=pa.string()), + "codebook": self._codebook_array([row.codebook for row in rows]), "index": pa.array([row.index for row in rows], type=pa.int64()), }, ) @@ -721,7 +746,12 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--batch_size", type=int, default=4096) parser.add_argument("--item_id_field", default="item_id") parser.add_argument("--code_field", default="codes") - parser.add_argument("--code_delimiter", default=",") + parser.add_argument( + "--code_delimiter", + default=",", + help="CSV-only: delimiter splitting/joining int codes in a codebook " + "string. Parquet/ODPS use list columns directly.", + ) parser.add_argument("--candidate_codebook_field", default="candidate_codebook") parser.add_argument( "--compact_candidate_field", diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index dd2f9e1e..203c5cfe 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -37,24 +37,23 @@ def tearDown(self) -> None: if os.path.exists(self.test_dir): shutil.rmtree(self.test_dir) + # ---- assigner (codebook is a tuple[int]) ---- + def test_assign_sid_collisions_respects_capacity(self) -> None: raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - RawSidRow("item_2", "A"), - RawSidRow("item_3", "B"), + RawSidRow("item_0", (1,)), + RawSidRow("item_1", (1,)), + RawSidRow("item_2", (1,)), + RawSidRow("item_3", (2,)), ] candidate_rows = [ - CandidateSidRow("item_0", "C", 1, 0.1), - CandidateSidRow("item_1", "C", 1, 0.1), - CandidateSidRow("item_2", "C", 1, 0.1), + CandidateSidRow("item_0", (3,), 1, 0.1), + CandidateSidRow("item_1", (3,), 1, 0.1), + CandidateSidRow("item_2", (3,), 1, 0.1), ] assigned, stats = assign_sid_collisions( - raw_rows, - candidate_rows, - capacity=2, - seed=7, + raw_rows, candidate_rows, capacity=2, seed=7 ) self.assertEqual(len(assigned), 4) @@ -66,11 +65,11 @@ def test_assign_sid_collisions_respects_capacity(self) -> None: self.assertEqual(stats.unassigned_count, 0) def test_random_strategy_reassigns_within_band_without_candidates(self) -> None: - # Three items collide on SID (lv1,lv2,lv3) = (1,2,3); capacity 1. + # Three items collide on SID (1,2,3); capacity 1. raw_rows = [ - RawSidRow("item_0", "1,2,3"), - RawSidRow("item_1", "1,2,3"), - RawSidRow("item_2", "1,2,3"), + RawSidRow("item_0", (1, 2, 3)), + RawSidRow("item_1", (1, 2, 3)), + RawSidRow("item_2", (1, 2, 3)), ] assigned, stats = assign_sid_collisions( @@ -84,19 +83,19 @@ def test_random_strategy_reassigns_within_band_without_candidates(self) -> None: self.assertEqual(len(assigned), 3) codebooks = {row.item_key: row.codebook for row in assigned} - # every item keeps its (lv1,lv2) band, only the last layer varies ... + # every item keeps its (1,2) band; only the last layer varies ... for codebook in codebooks.values(): - self.assertTrue(codebook.startswith("1,2,")) - # ... capacity 1 keeps exactly one item at the origin SID and reassigns the - # other two to distinct random last-layer codes -> near-injective. + self.assertEqual(codebook[:2], (1, 2)) + # ... capacity 1 keeps one item at the origin SID and reassigns the other + # two to distinct random last-layer codes -> near-injective. self.assertEqual(len(set(codebooks.values())), 3) - self.assertEqual(sum(1 for cb in codebooks.values() if cb == "1,2,3"), 1) + self.assertEqual(sum(1 for cb in codebooks.values() if cb == (1, 2, 3)), 1) self.assertEqual(stats.final_collision_buckets, 0) self.assertEqual(stats.reassigned_count, 2) self.assertEqual(stats.unassigned_count, 0) def test_random_strategy_is_deterministic_given_seed(self) -> None: - raw_rows = [RawSidRow(f"item_{i}", "0,0,0") for i in range(4)] + raw_rows = [RawSidRow(f"item_{i}", (0, 0, 0)) for i in range(4)] kwargs = dict(capacity=1, strategy="random", random_last_layer_size=64, seed=11) first, _ = assign_sid_collisions(raw_rows, [], **kwargs) second, _ = assign_sid_collisions(raw_rows, [], **kwargs) @@ -106,51 +105,213 @@ def test_random_strategy_is_deterministic_given_seed(self) -> None: ) def test_random_strategy_requires_last_layer_size(self) -> None: - raw_rows = [RawSidRow("item_0", "1,2,3")] + raw_rows = [RawSidRow("item_0", (1, 2, 3))] with self.assertRaisesRegex(ValueError, "random_last_layer_size"): assign_sid_collisions(raw_rows, [], capacity=1, strategy="random") def test_missing_candidates_errors_on_overflow(self) -> None: - raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - ] + raw_rows = [RawSidRow("item_0", (1,)), RawSidRow("item_1", (1,))] with self.assertRaisesRegex(ValueError, "no explicit candidate input"): assign_sid_collisions(raw_rows, [], capacity=1) def test_duplicate_candidates_do_not_consume_capacity_twice(self) -> None: raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - RawSidRow("item_2", "A"), + RawSidRow("item_0", (1,)), + RawSidRow("item_1", (1,)), + RawSidRow("item_2", (1,)), ] candidate_rows = [ - CandidateSidRow("item_0", "C", 1, 0.1), - CandidateSidRow("item_0", "C", 2, 0.2), - CandidateSidRow("item_1", "C", 1, 0.1), - CandidateSidRow("item_2", "C", 1, 0.1), + CandidateSidRow("item_0", (3,), 1, 0.1), + CandidateSidRow("item_0", (3,), 2, 0.2), + CandidateSidRow("item_1", (3,), 1, 0.1), + CandidateSidRow("item_2", (3,), 1, 0.1), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, candidate_rows, capacity=2, seed=7 + ) + + self.assertEqual(len(assigned), 3) + self.assertEqual(stats.unassigned_count, 0) + self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 2) + + def test_local_reassignment_indices_are_unique_and_dense(self) -> None: + # After reassignment every (codebook, index) pair must be unique and each + # bucket's indices must form a contiguous 1..N run. + raw_rows = [RawSidRow(f"item_{i}", (1,)) for i in range(4)] + candidate_rows = [ + CandidateSidRow("item_0", (2,), 1, 0.1), + CandidateSidRow("item_1", (2,), 1, 0.2), + CandidateSidRow("item_2", (2,), 1, 0.3), + CandidateSidRow("item_3", (2,), 1, 0.4), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, candidate_rows, capacity=2, seed=7 + ) + + self.assertEqual(len(assigned), 4) + self.assertEqual(stats.unassigned_count, 0) + pairs = [(row.codebook, row.index) for row in assigned] + self.assertEqual(len(pairs), len(set(pairs))) + by_codebook = defaultdict(list) + for codebook, index in pairs: + by_codebook[codebook].append(index) + for indices in by_codebook.values(): + self.assertEqual(sorted(indices), list(range(1, len(indices) + 1))) + + def test_local_drop_policy_omits_unplaceable_items(self) -> None: + # (1,) (capacity 1) keeps one item; the other two overflow and both want + # (2,), which fits only one -> one item is unplaceable (drop branch). + raw_rows = [ + RawSidRow("item_0", (1,)), + RawSidRow("item_1", (1,)), + RawSidRow("item_2", (1,)), + ] + candidate_rows = [ + CandidateSidRow("item_0", (2,), 1, 0.1), + CandidateSidRow("item_1", (2,), 1, 0.2), + CandidateSidRow("item_2", (2,), 1, 0.3), + ] + + assigned, stats = assign_sid_collisions( + raw_rows, candidate_rows, capacity=1, seed=7, unassigned_policy="drop" + ) + + self.assertEqual(stats.unassigned_count, 1) + self.assertEqual(len(assigned), 2) + self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 1) + + def test_local_keep_original_readds_over_capacity(self) -> None: + raw_rows = [ + RawSidRow("item_0", (1,)), + RawSidRow("item_1", (1,)), + RawSidRow("item_2", (1,)), + ] + candidate_rows = [ + CandidateSidRow("item_0", (2,), 1, 0.1), + CandidateSidRow("item_1", (2,), 1, 0.2), + CandidateSidRow("item_2", (2,), 1, 0.3), ] assigned, stats = assign_sid_collisions( raw_rows, candidate_rows, - capacity=2, + capacity=1, seed=7, + unassigned_policy="keep_original", ) - self.assertEqual(len(assigned), 3) self.assertEqual(stats.unassigned_count, 0) - self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 2) + self.assertEqual(len(assigned), 3) + # keep_original is the only policy allowed to exceed capacity. + self.assertEqual(Counter(row.codebook for row in assigned)[(1,)], 2) + self.assertEqual(stats.max_final_bucket_size, 2) + + def test_local_error_policy_raises_on_unplaceable(self) -> None: + raw_rows = [ + RawSidRow("item_0", (1,)), + RawSidRow("item_1", (1,)), + RawSidRow("item_2", (1,)), + ] + candidate_rows = [ + CandidateSidRow("item_0", (2,), 1, 0.1), + CandidateSidRow("item_1", (2,), 1, 0.2), + CandidateSidRow("item_2", (2,), 1, 0.3), + ] + + with self.assertRaisesRegex(RuntimeError, "could not be assigned"): + assign_sid_collisions( + raw_rows, + candidate_rows, + capacity=1, + seed=7, + unassigned_policy="error", + ) + + def test_local_score_order_higher_prefers_high_score(self) -> None: + # One item overflows (1,) (capacity 1); it can go to (2,) (score 0.1) or + # (3,) (score 0.9). score_order flips which score wins. + raw_rows = [RawSidRow("item_0", (1,)), RawSidRow("item_1", (1,))] + candidate_rows = [ + CandidateSidRow("item_0", (2,), 1, 0.1), + CandidateSidRow("item_0", (3,), 1, 0.9), + CandidateSidRow("item_1", (2,), 1, 0.1), + CandidateSidRow("item_1", (3,), 1, 0.9), + ] - def test_local_csv_outputs_codebooks_as_strings(self) -> None: + lower, _ = assign_sid_collisions( + raw_rows, candidate_rows, capacity=1, seed=7, score_order="lower" + ) + higher, _ = assign_sid_collisions( + raw_rows, candidate_rows, capacity=1, seed=7, score_order="higher" + ) + + reassigned_lower = next( + row for row in lower if row.origin_codebook != row.codebook + ) + reassigned_higher = next( + row for row in higher if row.origin_codebook != row.codebook + ) + self.assertEqual(reassigned_lower.codebook, (2,)) + self.assertEqual(reassigned_higher.codebook, (3,)) + + # ---- file backends: Parquet/ODPS list, CSV string fallback ---- + + def test_local_parquet_writes_list_codebooks(self) -> None: + raw_path = os.path.join(self.test_dir, "raw.parquet") + out_dir = os.path.join(self.test_dir, "out_parquet") + parquet.write_table( + pa.table( + { + "item_id": pa.array([1, 2, 3], type=pa.int64()), + "codes": pa.array( + [[1, 2], [1, 2], [1, 2]], type=pa.list_(pa.int64()) + ), + "candidate_codebook": pa.array( + [[1, 3], [1, 3], [1, 3]], type=pa.list_(pa.int64()) + ), + "priority": [1, 1, 1], + "score": [0.1, 0.1, 0.1], + } + ), + raw_path, + ) + + args = build_parser().parse_args( + [ + "--input_path", + raw_path, + "--output_path", + out_dir, + "--writer_type", + "ParquetWriter", + "--max_items_per_codebook", + "2", + ] + ) + stats = run(args) + + self.assertEqual(stats.reassigned_count, 1) + result = parquet.read_table(os.path.join(out_dir, "part-0.parquet")) + # Parquet/ODPS keep codebooks as list. + self.assertEqual(result.schema.field("codebook").type, pa.list_(pa.int64())) + origins = {tuple(x) for x in result["origin_codebook"].to_pylist()} + self.assertIn((1, 2), origins) + codebooks = [tuple(x) for x in result["codebook"].to_pylist()] + self.assertLessEqual(max(Counter(codebooks).values()), 2) + + def test_local_csv_writes_codebooks_as_strings(self) -> None: raw_path = os.path.join(self.test_dir, "raw.csv") out_dir = os.path.join(self.test_dir, "out") csv.write_csv( pa.table( { + # multi-layer so the codebook string keeps a comma and CSV + # read-back doesn't re-infer it as int. "item_id": ["1", "2", "3"], - "codes": ["A", "A", "A"], - "candidate_codebook": ["C", "C", "C"], + "codes": ["1,2", "1,2", "1,2"], + "candidate_codebook": ["3,4", "3,4", "3,4"], "priority": [1, 1, 1], "score": [0.1, 0.1, 0.1], } @@ -176,17 +337,16 @@ def test_local_csv_outputs_codebooks_as_strings(self) -> None: self.assertEqual(stats.reassigned_count, 1) result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) + # CSV fallback encodes codebooks as delimited strings. self.assertEqual(result.schema.field("origin_codebook").type, pa.string()) self.assertEqual(result.schema.field("codebook").type, pa.string()) self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) def test_writer_type_defaults_to_matching_the_reader(self) -> None: - # --writer_type unset: the writer is derived from the input reader - # (CsvReader -> CsvWriter), so CSV in yields CSV out. raw_path = os.path.join(self.test_dir, "raw_derive.csv") out_dir = os.path.join(self.test_dir, "out_derive") csv.write_csv( - pa.table({"item_id": ["1", "2", "3"], "codes": ["A", "A", "A"]}), + pa.table({"item_id": ["1", "2", "3"], "codes": ["1", "1", "1"]}), raw_path, ) args = build_parser().parse_args( @@ -200,45 +360,9 @@ def test_writer_type_defaults_to_matching_the_reader(self) -> None: ] ) run(args) - # A CSV part file the CSV reader can parse => CsvWriter was derived. result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) self.assertEqual(result.num_rows, 3) - def test_local_parquet_accepts_list_codes(self) -> None: - raw_path = os.path.join(self.test_dir, "raw.parquet") - out_dir = os.path.join(self.test_dir, "out_parquet") - parquet.write_table( - pa.table( - { - "item_id": pa.array([1, 2, 3], type=pa.int64()), - "codes": pa.array([[1, 2], [1, 2], [1, 2]]), - "candidate_codebook": ["1,3", "1,3", "1,3"], - "priority": [1, 1, 1], - "score": [0.1, 0.1, 0.1], - } - ), - raw_path, - ) - - args = build_parser().parse_args( - [ - "--input_path", - raw_path, - "--output_path", - out_dir, - "--writer_type", - "ParquetWriter", - "--max_items_per_codebook", - "2", - ] - ) - stats = run(args) - - self.assertEqual(stats.reassigned_count, 1) - result = parquet.read_table(os.path.join(out_dir, "part-0.parquet")) - self.assertIn("1,2", set(result["origin_codebook"].to_pylist())) - self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) - def test_local_csv_accepts_compact_candidates(self) -> None: raw_path = os.path.join(self.test_dir, "raw_compact.csv") out_dir = os.path.join(self.test_dir, "out_compact") @@ -246,8 +370,8 @@ def test_local_csv_accepts_compact_candidates(self) -> None: pa.table( { "item_id": ["1", "2", "3"], - "codes": ["A,B", "A,B", "A,B"], - "sorted_index": ["A|C", "A|C", "A|C"], + "codes": ["1,2", "1,2", "1,2"], + "candidates": ["1|3", "1|3", "1|3"], } ), raw_path, @@ -264,7 +388,7 @@ def test_local_csv_accepts_compact_candidates(self) -> None: "--writer_type", "CsvWriter", "--compact_candidate_field", - "sorted_index", + "candidates", "--max_items_per_codebook", "2", ] @@ -273,154 +397,83 @@ def test_local_csv_accepts_compact_candidates(self) -> None: self.assertEqual(stats.reassigned_count, 1) result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) - self.assertIn("A,B", set(result["origin_codebook"].to_pylist())) - self.assertIn("A", set(result["codebook"].to_pylist())) + self.assertIn("1,2", set(result["origin_codebook"].to_pylist())) + self.assertTrue({"1", "3"} & set(result["codebook"].to_pylist())) self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) - def test_local_reassignment_indices_are_unique_and_dense(self) -> None: - # After reassignment every (codebook, index) pair must be unique and each - # bucket's indices must form a contiguous 1..N run. (Which two items - # overflow A is decided by a seeded hash, so every item carries the same - # B fallback.) - raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - RawSidRow("item_2", "A"), - RawSidRow("item_3", "A"), - ] - candidate_rows = [ - CandidateSidRow("item_0", "B", 1, 0.1), - CandidateSidRow("item_1", "B", 1, 0.2), - CandidateSidRow("item_2", "B", 1, 0.3), - CandidateSidRow("item_3", "B", 1, 0.4), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, - candidate_rows, - capacity=2, - seed=7, - ) - - self.assertEqual(len(assigned), 4) - self.assertEqual(stats.unassigned_count, 0) - pairs = [(row.codebook, row.index) for row in assigned] - self.assertEqual(len(pairs), len(set(pairs))) - by_codebook = defaultdict(list) - for codebook, index in pairs: - by_codebook[codebook].append(index) - for indices in by_codebook.values(): - self.assertEqual(sorted(indices), list(range(1, len(indices) + 1))) - - def test_local_drop_policy_omits_unplaceable_items(self) -> None: - # A (capacity 1) keeps one item; the other two overflow and both want B, - # which fits only one -> one item stays unplaceable and reaches - # _handle_unassigned's drop branch. - raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - RawSidRow("item_2", "A"), - ] - candidate_rows = [ - CandidateSidRow("item_0", "B", 1, 0.1), - CandidateSidRow("item_1", "B", 1, 0.2), - CandidateSidRow("item_2", "B", 1, 0.3), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, - candidate_rows, - capacity=1, - seed=7, - unassigned_policy="drop", - ) - - self.assertEqual(stats.unassigned_count, 1) - # The unplaceable item is dropped: A keeps 1, B keeps 1, nothing else. - self.assertEqual(len(assigned), 2) - self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 1) - - def test_local_keep_original_readds_over_capacity(self) -> None: - raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - RawSidRow("item_2", "A"), - ] - candidate_rows = [ - CandidateSidRow("item_0", "B", 1, 0.1), - CandidateSidRow("item_1", "B", 1, 0.2), - CandidateSidRow("item_2", "B", 1, 0.3), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, - candidate_rows, - capacity=1, - seed=7, - unassigned_policy="keep_original", - ) - - self.assertEqual(stats.unassigned_count, 0) - self.assertEqual(len(assigned), 3) - # The unplaceable item is re-added at its origin: keep_original is the - # only policy allowed to exceed capacity (A now holds 2 > capacity 1). - self.assertEqual(Counter(row.codebook for row in assigned)["A"], 2) - self.assertEqual(stats.max_final_bucket_size, 2) - - def test_local_error_policy_raises_on_unplaceable(self) -> None: - raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - RawSidRow("item_2", "A"), - ] - candidate_rows = [ - CandidateSidRow("item_0", "B", 1, 0.1), - CandidateSidRow("item_1", "B", 1, 0.2), - CandidateSidRow("item_2", "B", 1, 0.3), - ] - - # Distinct from the "no explicit candidate input" ValueError: candidates - # are provided but cannot place every overflow item. - with self.assertRaisesRegex(RuntimeError, "could not be assigned"): - assign_sid_collisions( - raw_rows, - candidate_rows, - capacity=1, - seed=7, - unassigned_policy="error", + def test_csv_and_parquet_yield_identical_decisions(self) -> None: + # The same logical data as CSV (delimited strings) and Parquet + # (list) must normalize to the same tuple codebooks, so the + # assignment decisions are identical across backends. + item_ids = list(range(20)) + origin = [[i % 3, i % 2] for i in item_ids] # heavy collision on few buckets + cand = [[9, i % 5] for i in item_ids] + + def decisions(mode: str) -> dict: + if mode == "parquet": + path = os.path.join(self.test_dir, "eq.parquet") + parquet.write_table( + pa.table( + { + "item_id": pa.array(item_ids, type=pa.int64()), + "codes": pa.array(origin, type=pa.list_(pa.int64())), + "candidate_codebook": pa.array( + cand, type=pa.list_(pa.int64()) + ), + } + ), + path, + ) + reader, writer = "ParquetReader", "ParquetWriter" + else: + path = os.path.join(self.test_dir, "eq.csv") + csv.write_csv( + pa.table( + { + "item_id": [str(i) for i in item_ids], + "codes": [",".join(map(str, c)) for c in origin], + "candidate_codebook": [",".join(map(str, c)) for c in cand], + } + ), + path, + ) + reader, writer = "CsvReader", "CsvWriter" + out = os.path.join(self.test_dir, f"eq_out_{mode}") + run( + build_parser().parse_args( + [ + "--input_path", + path, + "--output_path", + out, + "--reader_type", + reader, + "--writer_type", + writer, + "--max_items_per_codebook", + "2", + "--unassigned_policy", + "keep_original", + "--seed", + "5", + ] + ) ) - - def test_local_score_order_higher_prefers_high_score(self) -> None: - # One item overflows A (capacity 1); it can go to B (score 0.1) or - # C (score 0.9). score_order flips which score wins. Both items carry both - # candidates so the choice is independent of which item the seeded hash - # picks to overflow. - raw_rows = [ - RawSidRow("item_0", "A"), - RawSidRow("item_1", "A"), - ] - candidate_rows = [ - CandidateSidRow("item_0", "B", 1, 0.1), - CandidateSidRow("item_0", "C", 1, 0.9), - CandidateSidRow("item_1", "B", 1, 0.1), - CandidateSidRow("item_1", "C", 1, 0.9), - ] - - lower, _ = assign_sid_collisions( - raw_rows, candidate_rows, capacity=1, seed=7, score_order="lower" - ) - higher, _ = assign_sid_collisions( - raw_rows, candidate_rows, capacity=1, seed=7, score_order="higher" - ) - - reassigned_lower = next( - row for row in lower if row.origin_codebook != row.codebook - ) - reassigned_higher = next( - row for row in higher if row.origin_codebook != row.codebook - ) - self.assertEqual(reassigned_lower.codebook, "B") - self.assertEqual(reassigned_higher.codebook, "C") + if mode == "parquet": + d = parquet.read_table(os.path.join(out, "part-0.parquet")).to_pydict() + return { + int(d["item_id"][i]): tuple(d["codebook"][i]) + for i in range(len(d["item_id"])) + } + d = csv.read_csv(os.path.join(out, "part-0.csv")).to_pydict() + return { + int(d["item_id"][i]): tuple( + int(x) for x in str(d["codebook"][i]).split(",") + ) + for i in range(len(d["item_id"])) + } + + self.assertEqual(decisions("parquet"), decisions("csv")) if __name__ == "__main__": From cd3f5036ef7927ed4e44bd9bd588d56efb702e7c Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 9 Jul 2026 11:22:58 +0000 Subject: [PATCH 13/29] [perf] SID collision tool: vectorized within-band greedy reassignment Replace the per-row object-based iterative allocator with a numpy/Arrow vectorized single greedy pass that reassigns only the last SID layer within an item's band. The object path materialized a dataclass plus candidate objects per row and drove an up-to-50-pass allocation, which does not scale to the hundred-million-row maps this tool targets; the vectorized path packs band keys to int64, ranks with lexsort/unique, and runs one first-fit pass, keeping the multi-backend reader/writer and adding chunked writes plus --rate_only. The greedy single pass replaces ALGR's iterative multi-pass allocation, trading a small placement-quality loss in near-saturated bands for a ~15x faster, vectorizable, contention-insensitive pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 1119 ++++++++---------- tzrec/tools/sid/collision_prevention_test.py | 659 ++++------- 2 files changed, 715 insertions(+), 1063 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index f155dda0..a868f449 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -9,24 +9,31 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Offline SID codebook collision-prevention tool. - -The default ('candidate') allocator is a deterministic post-process over predicted -SID rows plus explicit candidate SID rows. The opt-in 'random' strategy instead -generates random within-band last-layer candidates (no candidate input required) as -a baseline that ignores semantic nearest-neighbor proximity; it is still fully -reproducible given ``seed``. +"""Offline SID codebook collision-prevention tool (vectorized). + +Caps every SID bucket at ``--max_items_per_codebook`` and reassigns overflow +items to a free code in the SAME band -- keeping every layer but the last and +varying only the last SID layer, so an item never leaves its ``(prefix)`` band: + +- ``--strategy candidate`` walks the item's model-provided candidate SIDs + (``candidate_codes``), taking the last code of each, best-first; +- ``--strategy random`` draws random last-layer codes (excluding the origin), + a baseline that ignores semantic proximity. + +Items with no free slot fall back per ``--unassigned_policy``. The hot path is +numpy/Arrow-vectorized so a hundred-million-row map fits in one pass; results +are deterministic given ``--seed`` and independent of input row order. I/O goes +through ``create_reader`` / ``create_writer`` so CSV, Parquet, and ODPS all work; +outputs are written in chunks to stay under Arrow's 2^31 list-offset limit. """ import argparse -import hashlib -import heapq -import random -from collections import defaultdict from dataclasses import asdict, dataclass -from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple +from typing import Any, Dict, Iterable, List, Optional, Tuple +import numpy as np import pyarrow as pa +import pyarrow.compute as pc # Register the local reader/writer classes used through create_reader/writer. import tzrec.datasets.csv_dataset # noqa: F401 @@ -34,51 +41,20 @@ from tzrec.datasets.dataset import create_reader, create_writer from tzrec.utils.logging_util import logger -# A SID bucket key: the tuple of per-layer integer codes. Used only as an opaque -# hashable, ordered identity (grouped / counted / compared), never for arithmetic. -# list on disk for Parquet/ODPS; a delimited string for the CSV fallback. -Codebook = Tuple[int, ...] - - -@dataclass(frozen=True, slots=True) -class RawSidRow: - """One raw item -> SID row.""" - - item_id: Any - origin_codebook: Codebook - - @property - def item_key(self) -> str: - """Stable dict/sort key for the item (its id as a string).""" - return str(self.item_id) - +_MASK64 = np.uint64((1 << 64) - 1) -@dataclass(frozen=True, slots=True) -class CandidateSidRow: - """One item -> candidate SID row.""" +# CSV-fallback delimiters (Parquet/ODPS carry codes as native list, so +# these are unused there). Both are non-comma so a codebook cell never collides +# with CSV's own field separator, and distinct so a compact cell parses cleanly. +_CODE_SEP = "|" # separates the int codes within one SID, e.g. "1|2|3" +_CAND_SEP = ";" # separates candidate SIDs in a compact cell, e.g. "1|2;3|4" - item_key: str - candidate_codebook: Codebook - priority: int - score: float +# Chunk rows per output write: keeps each Arrow array's int32 offsets under 2^31 +# (a single 255M-row string/list column would overflow one array). +_WRITE_CHUNK = 20_000_000 -@dataclass(frozen=True, slots=True) -class AssignedSidRow: - """One final item -> SID assignment row.""" - - item_id: Any - origin_codebook: Codebook - codebook: Codebook - index: int - - @property - def item_key(self) -> str: - """Stable dict/sort key for the item (its id as a string).""" - return str(self.item_id) - - -@dataclass(frozen=True, slots=True) +@dataclass class AssignmentStats: """Summary statistics for a collision-prevention run.""" @@ -87,345 +63,55 @@ class AssignmentStats: final_collision_buckets: int reassigned_count: int unassigned_count: int - iteration_count: int max_final_bucket_size: int -class SidCollisionAssigner: - """Deterministically assign overflow SID rows to explicit candidates.""" - - def __init__( - self, - capacity: int, - max_iters: int = 50, - seed: int = 2026, - score_order: str = "lower", - unassigned_policy: str = "error", - strategy: str = "candidate", - code_delimiter: str = ",", - random_last_layer_size: Optional[int] = None, - random_num_candidates: int = 64, - ) -> None: - if capacity < 1: - raise ValueError(f"capacity must be >= 1, got {capacity}") - if score_order not in ("lower", "higher"): - raise ValueError("score_order must be 'lower' or 'higher'.") - if unassigned_policy not in ("error", "drop", "keep_original"): - raise ValueError( - "unassigned_policy must be one of: error, drop, keep_original." - ) - if strategy not in ("candidate", "random"): - raise ValueError("strategy must be 'candidate' or 'random'.") - if strategy == "random": - if random_last_layer_size is None or random_last_layer_size < 2: - raise ValueError( - "strategy='random' requires random_last_layer_size >= 2." - ) - if random_num_candidates < 1: - raise ValueError("random_num_candidates must be >= 1.") - self.capacity = capacity - self.max_iters = max_iters - self.seed = seed - self.score_order = score_order - self.unassigned_policy = unassigned_policy - self.strategy = strategy - self.code_delimiter = code_delimiter - self.random_last_layer_size = random_last_layer_size - self.random_num_candidates = random_num_candidates - self._candidate_sort_keys: Dict[ - CandidateSidRow, Tuple[int, float, int, str, Codebook] - ] = {} - - def assign( - self, - raw_rows: Sequence[RawSidRow], - candidate_rows: Sequence[CandidateSidRow], - ) -> Tuple[List[AssignedSidRow], AssignmentStats]: - """Assign overflow SID rows to non-full candidate codebooks.""" - raw_by_item = {row.item_key: row for row in raw_rows} - if len(raw_by_item) != len(raw_rows): - raise ValueError("raw_rows contains duplicate item_id values.") - - by_origin: Dict[str, List[RawSidRow]] = defaultdict(list) - for row in raw_rows: - by_origin[row.origin_codebook].append(row) - - assigned: List[AssignedSidRow] = [] - assigned_items = set() - code_counts: Dict[str, int] = defaultdict(int) - - for codebook, rows in by_origin.items(): - for index, row in enumerate( - heapq.nsmallest(self.capacity, rows, key=self._assignment_sort_key), - start=1, - ): - assigned.append( - AssignedSidRow( - item_id=row.item_id, - origin_codebook=row.origin_codebook, - codebook=codebook, - index=index, - ) - ) - assigned_items.add(row.item_key) - code_counts[codebook] += 1 - - overflow_items = set(raw_by_item) - assigned_items - if self.strategy == "random": - candidate_rows = self._generate_random_candidates( - overflow_items, raw_by_item - ) - elif overflow_items and not candidate_rows: - raise ValueError( - "raw SID input has overflow rows, but no explicit candidate input was " - "provided." - ) - sorted_by_codebook = self._dedup_candidates(candidate_rows, overflow_items) - - unassigned = set(overflow_items) - iteration_count = self._assign_candidates( - raw_by_item, - sorted_by_codebook, - assigned, - code_counts, - unassigned, - ) - self._handle_unassigned(raw_by_item, assigned, code_counts, unassigned) - - final_counts: Dict[str, int] = defaultdict(int) - reassigned = 0 - for row in assigned: - final_counts[row.codebook] += 1 - if row.origin_codebook != row.codebook: - reassigned += 1 - stats = AssignmentStats( - total_items=len(raw_rows), - raw_collision_buckets=sum( - 1 for rows in by_origin.values() if len(rows) > self.capacity - ), - final_collision_buckets=sum( - 1 for count in final_counts.values() if count > self.capacity - ), - reassigned_count=reassigned, - unassigned_count=( - len(unassigned) if self.unassigned_policy != "keep_original" else 0 - ), - iteration_count=iteration_count, - max_final_bucket_size=max(final_counts.values()) if final_counts else 0, - ) - assigned.sort(key=lambda r: (r.codebook, r.index, r.item_key)) - return assigned, stats +def _splitmix64(x: np.ndarray, seed: int) -> np.ndarray: + """Vectorized order-independent SplitMix64 hash of a uint64 array. - @staticmethod - def _stable_hash(*parts: Any) -> int: - h = hashlib.blake2b(digest_size=8) - for part in parts: - h.update(str(part).encode("utf-8")) - h.update(b"\x1f") - return int.from_bytes(h.digest(), byteorder="big", signed=False) - - def _generate_random_candidates( - self, - overflow_items: Iterable[str], - raw_by_item: Dict[str, RawSidRow], - ) -> List[CandidateSidRow]: - """Generate random within-band last-layer candidates for overflow items. - - Keeps every layer except the last and draws distinct random codes for the - last layer in ``[0, random_last_layer_size)``, excluding the origin code, so - each item stays in its own ``(prefix)`` band. Deterministic given ``seed``; - unlike the 'candidate' strategy it ignores nearest-neighbor proximity. - """ - assert self.random_last_layer_size is not None - size = self.random_last_layer_size - num_draws = min(self.random_num_candidates, size - 1) - rows: List[CandidateSidRow] = [] - for item_key in overflow_items: - raw = raw_by_item[item_key] - prefix = raw.origin_codebook[:-1] - origin_last = raw.origin_codebook[-1] if raw.origin_codebook else None - rng = random.Random(self._stable_hash(self.seed, item_key)) - drawn: set = set() - while len(drawn) < num_draws: - value = rng.randrange(size) - if value in drawn or value == origin_last: - continue - drawn.add(value) - rows.append( - CandidateSidRow( - item_key=item_key, - candidate_codebook=(*prefix, value), - priority=1, - score=0.0, - ) - ) - return rows + A pure function of the input values (not their position), so it gives a + stable, seedable tie-break that is invariant to input row order -- unlike a + read-order index -- while staying fully vectorized. - def _assignment_sort_key(self, row: RawSidRow) -> Tuple[int, str]: - return ( - self._stable_hash(self.seed, row.origin_codebook, row.item_key), - row.item_key, - ) + Args: + x (np.ndarray): uint64 values to hash. + seed (int): mixing seed. - def _candidate_sort_key( - self, - row: CandidateSidRow, - ) -> Tuple[int, float, int, str, Codebook]: - # Memoized: this key's blake2b tie-breaker is pure but is compared inside - # the up-to-max_iters assignment loop, so recomputing it would dominate - # the phase. CandidateSidRow is frozen/hashable, so cache on the row. - cached = self._candidate_sort_keys.get(row) - if cached is not None: - return cached - score = row.score if self.score_order == "lower" else -row.score - key = ( - row.priority, - score, - self._stable_hash(self.seed, row.item_key, row.candidate_codebook), - row.item_key, - row.candidate_codebook, - ) - self._candidate_sort_keys[row] = key - return key + Returns: + np.ndarray: uint64 hashes, same shape as ``x``. + """ + with np.errstate(over="ignore"): + z = x.astype(np.uint64) + np.uint64((seed * 0x9E3779B97F4A7C15) & int(_MASK64)) + z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) + z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) + return z ^ (z >> np.uint64(31)) - def _dedup_candidates( - self, - candidate_rows: Sequence[CandidateSidRow], - overflow_items: set, - ) -> Dict[str, List[CandidateSidRow]]: - # Only overflow items' candidates can ever be used, so filtering here - # keeps the dedup map (and the sort-key memo) to the overflow fraction. - dedup_candidates: Dict[Tuple[str, str], CandidateSidRow] = {} - for row in candidate_rows: - if row.item_key not in overflow_items: - continue - key = (row.item_key, row.candidate_codebook) - current = dedup_candidates.get(key) - if current is None or self._candidate_sort_key( - row - ) < self._candidate_sort_key(current): - dedup_candidates[key] = row - - # Group by codebook and sort each list ONCE. The sort key is a total - # order, so the loop can scan these lists instead of re-sorting per pass. - sorted_by_codebook: Dict[str, List[CandidateSidRow]] = defaultdict(list) - for row in dedup_candidates.values(): - sorted_by_codebook[row.candidate_codebook].append(row) - for rows in sorted_by_codebook.values(): - rows.sort(key=self._candidate_sort_key) - return sorted_by_codebook - - def _assign_candidates( - self, - raw_by_item: Dict[str, RawSidRow], - sorted_by_codebook: Dict[str, List[CandidateSidRow]], - assigned: List[AssignedSidRow], - code_counts: Dict[str, int], - unassigned: set, - ) -> int: - iteration_count = 0 - for iteration in range(self.max_iters): - if not unassigned: - break - accepted = self._select_candidates( - sorted_by_codebook, unassigned, code_counts - ) - if not accepted: - break - progress = 0 - for candidate in accepted: - if candidate.item_key not in unassigned: - continue - if code_counts[candidate.candidate_codebook] >= self.capacity: - continue - raw = raw_by_item[candidate.item_key] - code_counts[candidate.candidate_codebook] += 1 - assigned.append( - AssignedSidRow( - item_id=raw.item_id, - origin_codebook=raw.origin_codebook, - codebook=candidate.candidate_codebook, - index=code_counts[candidate.candidate_codebook], - ) - ) - unassigned.discard(candidate.item_key) - progress += 1 +def _order_hash(item_ids: np.ndarray, seed: int) -> np.ndarray: + """Per-item order-independent tie-break hash (uint64), vectorized. - iteration_count = iteration + 1 - if progress == 0: - break - return iteration_count + Integer ids hash directly; string/object ids are folded to uint64 via + ``pandas.util.hash_array`` first. Both then pass through :func:`_splitmix64` + so the ``seed`` is mixed in uniformly. - def _select_candidates( - self, - sorted_by_codebook: Dict[str, List[CandidateSidRow]], - unassigned: set, - code_counts: Dict[str, int], - ) -> List[CandidateSidRow]: - # Per open codebook, take its `remaining` smallest still-unassigned - # candidates from the pre-sorted list; then keep each item's single best - # and return sorted -- identical to the old sort-per-pass path (unique - # keys make sort-then-filter == filter-then-sort). - selected_by_codebook: List[CandidateSidRow] = [] - for codebook, rows in sorted_by_codebook.items(): - remaining = self.capacity - code_counts[codebook] - if remaining <= 0: - continue - taken = 0 - for row in rows: - if row.item_key in unassigned: - selected_by_codebook.append(row) - taken += 1 - if taken == remaining: - break - - best_by_item: Dict[str, CandidateSidRow] = {} - for candidate in selected_by_codebook: - current = best_by_item.get(candidate.item_key) - if current is None or self._candidate_sort_key( - candidate - ) < self._candidate_sort_key(current): - best_by_item[candidate.item_key] = candidate - - return sorted( - best_by_item.values(), - key=lambda r: (r.candidate_codebook, self._candidate_sort_key(r)), - ) + Args: + item_ids (np.ndarray): item id values. + seed (int): mixing seed. - def _handle_unassigned( - self, - raw_by_item: Dict[str, RawSidRow], - assigned: List[AssignedSidRow], - code_counts: Dict[str, int], - unassigned: set, - ) -> None: - pending = sorted(unassigned) - if not pending: - return + Returns: + np.ndarray: uint64 per-item hashes. + """ + if np.issubdtype(item_ids.dtype, np.integer): + base = item_ids.astype(np.uint64) + else: + import pandas as pd - if self.unassigned_policy == "error": - preview = ",".join(pending[:10]) - raise RuntimeError( - f"{len(pending)} items could not be assigned within capacity; " - f"first unassigned item_ids: {preview}" - ) - if self.unassigned_policy == "keep_original": - for item_key in pending: - raw = raw_by_item[item_key] - code_counts[raw.origin_codebook] += 1 - assigned.append( - AssignedSidRow( - item_id=raw.item_id, - origin_codebook=raw.origin_codebook, - codebook=raw.origin_codebook, - index=code_counts[raw.origin_codebook], - ) - ) + base = pd.util.hash_array(np.asarray(item_ids, dtype=object)) + return _splitmix64(base, seed) class CollisionRunner: - """SID collision-prevention runner over the standard dataset reader/writer. + """Vectorized SID collision-prevention runner over the dataset reader/writer. The backend (CSV / Parquet / ODPS) is chosen by the reader/writer type, so the same path serves local files and MaxCompute tables. @@ -433,77 +119,21 @@ class CollisionRunner: def __init__(self, args: argparse.Namespace) -> None: self.args = args - # Class name of the main-input reader, captured during the raw read; used - # to derive the writer type when --writer_type is unset. + # Class name of the input reader, captured on the first read; used to + # derive the writer type when --writer_type is unset. self._input_reader_cls_name: Optional[str] = None - # Dedup identical codebook tuples so a repeated SID is one object. - self._cb_cache: Dict[Codebook, Codebook] = {} def run(self) -> AssignmentStats: - """Read inputs, assign, and write outputs via create_reader/create_writer.""" - raw_rows, candidate_rows = self._load_rows() - rows, stats = SidCollisionAssigner( - capacity=self.args.max_items_per_codebook, - max_iters=self.args.max_iters, - seed=self.args.seed, - score_order=self.args.score_order, - unassigned_policy=self.args.unassigned_policy, - strategy=self.args.strategy, - code_delimiter=self.args.code_delimiter, - random_last_layer_size=self.args.random_last_layer_size, - random_num_candidates=self.args.random_num_candidates, - ).assign(raw_rows, candidate_rows) - self._write_assignments(rows) - if self.args.diagnostics_output_path: - self._write_diagnostics(stats) + """Read the SID map, assign, and write the reassigned map + stats.""" + item_ids, codes = self._load_codes() + final_codes, index, keep_mask, stats = self._assign(item_ids, codes) + if not self.args.rate_only: + self._write(item_ids, codes, final_codes, index, keep_mask, stats) + else: + logger.info("rate_only: skipping map write") logger.info("SID collision prevention finished: %s", stats) return stats - def _to_codebook(self, value: Any, delimiter: str) -> Codebook: - """Normalize a SID cell to a tuple of int codes (deduped). - - Parquet/ODPS give a list cell; CSV gives a delimited string. - """ - if value is None: - raise ValueError("SID code value cannot be null.") - if isinstance(value, (list, tuple)): - codes: Codebook = tuple(int(v) for v in value) - else: - codes = tuple(int(p) for p in str(value).split(delimiter) if p.strip()) - return self._cb_cache.setdefault(codes, codes) - - def _split_compact(self, value: Any) -> List[Codebook]: - """Split a compact candidate cell into candidate codebooks. - - Parquet/ODPS give a list> cell; CSV gives candidate SIDs - joined by --candidate_delimiter (e.g. ``"1,3|4,5"``). - """ - if value is None: - return [] - if isinstance(value, (list, tuple)): - return [ - self._to_codebook(inner, self.args.code_delimiter) - for inner in value - if inner is not None - ] - text = str(value).strip() - if not text: - return [] - return [ - self._to_codebook(part, self.args.code_delimiter) - for part in text.split(self.args.candidate_delimiter) - if part.strip() - ] - - @staticmethod - def _array_to_pylist(batch: Dict[str, pa.Array], field: str) -> List[Any]: - if field not in batch: - raise ValueError( - f"required field {field!r} not found; available fields: " - f"{sorted(batch.keys())}" - ) - return batch[field].to_pylist() - def _read( self, selected_cols: Optional[List[str]], capture_reader_cls: bool = False ) -> Iterable[Dict[str, pa.Array]]: @@ -518,140 +148,353 @@ def _read( self._input_reader_cls_name = reader.__class__.__name__ yield from reader.to_batches() - def _candidate_field(self) -> Tuple[Optional[str], bool]: - """Resolve the candidate column and whether it is compact. - - Returns ``(None, False)`` for the ``random`` strategy (it synthesizes its - own candidates and reads no candidate column). - """ - if self.args.strategy != "candidate": - return None, False - codebook_field = ( - None - if self.args.compact_candidate_field - else self.args.candidate_codebook_field - ) - if bool(codebook_field) == bool(self.args.compact_candidate_field): - raise ValueError( - "Set exactly one of --candidate_codebook_field or " - "--compact_candidate_field." - ) - return ( - self.args.compact_candidate_field or codebook_field, - bool(self.args.compact_candidate_field), - ) - - def _load_rows(self) -> Tuple[List[RawSidRow], List[CandidateSidRow]]: - """Read raw SID rows and candidate rows from the single input table. + @staticmethod + def _codes_matrix(arr: pa.Array) -> np.ndarray: + """Decode a SID code column into an (N, n_layers) int64 matrix. - Two passes: pass 1 reads only the id + code columns (raw rows + per-origin - counts); pass 2 reads the candidate columns and builds a CandidateSidRow - only for items in over-capacity buckets -- the only ones that can overflow - -- so candidate memory scales with the overflow fraction, not the whole - table. Candidates are read only for the ``candidate`` strategy and only - when the column is present, so a plain SID table still runs. + Parquet/ODPS give a ``list`` cell; CSV gives a ``_CODE_SEP`` + string; a single-layer numeric CSV column may arrive already as ints. """ - candidate_field, is_compact = self._candidate_field() - - raw_rows: List[RawSidRow] = [] - origin_counts: Dict[str, int] = defaultdict(int) - seen: set = set() + if pa.types.is_list(arr.type) or pa.types.is_large_list(arr.type): + n = len(arr) + flat = arr.flatten().to_numpy(zero_copy_only=False).astype(np.int64) + elif pa.types.is_integer(arr.type): + return arr.to_numpy(zero_copy_only=False).astype(np.int64).reshape(-1, 1) + else: + parts = pc.split_pattern(arr, _CODE_SEP) + n = len(parts) + flat = pc.cast(parts.flatten(), pa.int64()).to_numpy(zero_copy_only=False) + if n == 0: + return np.empty((0, 0), dtype=np.int64) + n_layers = flat.shape[0] // n + if flat.shape[0] != n * n_layers: + raise ValueError("ragged SID codes: all items must share n_layers.") + return flat.reshape(n, n_layers) + + def _load_codes(self) -> Tuple[np.ndarray, np.ndarray]: + """Stream the map into an item-id array and an (N, n_layers) code matrix.""" + id_chunks: List[np.ndarray] = [] + code_chunks: List[np.ndarray] = [] for batch in self._read( [self.args.item_id_field, self.args.code_field], capture_reader_cls=True ): - item_ids = self._array_to_pylist(batch, self.args.item_id_field) - code_cells = self._array_to_pylist(batch, self.args.code_field) - for item_id, code_cell in zip(item_ids, code_cells): - item_key = str(item_id) - if item_key in seen: - raise ValueError(f"duplicate item_id in SID input: {item_key}") - seen.add(item_key) - origin = self._to_codebook(code_cell, self.args.code_delimiter) - raw_rows.append(RawSidRow(item_id=item_id, origin_codebook=origin)) - origin_counts[origin] += 1 - if not raw_rows: + id_chunks.append( + batch[self.args.item_id_field].to_numpy(zero_copy_only=False) + ) + code_chunks.append(self._codes_matrix(batch[self.args.code_field])) + if not id_chunks: raise ValueError("SID input is empty.") - - cap = self.args.max_items_per_codebook - overflow = { - r.item_key for r in raw_rows if origin_counts[r.origin_codebook] > cap - } - if candidate_field is None or not overflow: - return raw_rows, [] - - candidate_rows: List[CandidateSidRow] = [] - for batch in self._read(None): - if candidate_field not in batch: + item_ids = np.concatenate(id_chunks) + codes = np.concatenate(code_chunks, axis=0) + if codes.shape[1] < 1: + raise ValueError("SID codes must have at least one layer.") + return item_ids, codes + + def _load_candidate_last( + self, overflow_id_arr: np.ndarray + ) -> Dict[Any, np.ndarray]: + """Map each overflow item id to its ordered candidate last-layer codes. + + Reads ``candidate_codes`` (list> for Parquet/ODPS, a + ``_CAND_SEP``/``_CODE_SEP`` string for CSV) but keeps only the last code + of each candidate SID and only for overflow items, so memory scales with + the overflow fraction, not the whole table. + """ + depth = self.args.candidate_depth + field = self.args.candidate_codes_field + cand: Dict[Any, np.ndarray] = {} + for batch in self._read([self.args.item_id_field, field]): + if field not in batch: break - item_ids = self._array_to_pylist(batch, self.args.item_id_field) - candidate_rows.extend( - self._candidate_rows(batch, item_ids, is_compact, overflow) - ) - return raw_rows, candidate_rows + ids = batch[self.args.item_id_field].to_numpy(zero_copy_only=False) + keep = np.where(np.isin(ids, overflow_id_arr))[0] + if keep.size == 0: + continue + col = batch[field] + if pa.types.is_list(col.type) or pa.types.is_large_list(col.type): + self._collect_candidate_lists(ids, col, keep, depth, cand) + else: + self._collect_candidate_strings(ids, col, keep, depth, cand) + return cand - def _candidate_rows( + @staticmethod + def _collect_candidate_lists( + ids: np.ndarray, + col: pa.Array, + keep: np.ndarray, + depth: Optional[int], + cand: Dict[Any, np.ndarray], + ) -> None: + """Vectorized last-code extraction from a list> batch.""" + n = len(col) + inner = col.flatten() # list, one per candidate SID + if len(inner) == 0: + return + k = len(inner) // n + if len(inner) != n * k: + raise ValueError("ragged candidate_codes: all items must share topk.") + flat = inner.flatten().to_numpy(zero_copy_only=False).astype(np.int64) + n_layers = flat.shape[0] // (n * k) + last = flat.reshape(n, k, n_layers)[:, :, n_layers - 1] + if depth is not None: + last = last[:, :depth] + for i in keep.tolist(): + cand[ids[i]] = last[i] + + def _collect_candidate_strings( self, - batch: Dict[str, pa.Array], - item_ids: List[Any], - is_compact: bool, - overflow: set, - ) -> List[CandidateSidRow]: - rows: List[CandidateSidRow] = [] - if is_compact: - compact_values = self._array_to_pylist( - batch, self.args.compact_candidate_field - ) - for item_id, compact_value in zip(item_ids, compact_values): - item_key = str(item_id) - if item_key not in overflow: + ids: np.ndarray, + col: pa.Array, + keep: np.ndarray, + depth: Optional[int], + cand: Dict[Any, np.ndarray], + ) -> None: + """Per-row last-code extraction from a CSV compact-candidate batch.""" + values = col.to_pylist() + for i in keep.tolist(): + text = values[i] + if not text: + continue + last_codes = [] + for part in str(text).split(_CAND_SEP): + part = part.strip() + if not part: continue - for priority, candidate in enumerate( - self._split_compact(compact_value), - start=1, - ): - rows.append( - CandidateSidRow( - item_key=item_key, - candidate_codebook=candidate, - priority=priority, - score=0.0, - ) - ) - return rows - - priorities = ( - self._array_to_pylist(batch, self.args.priority_field) - if self.args.priority_field and self.args.priority_field in batch - else None + codes = [int(p) for p in part.split(_CODE_SEP) if p.strip()] + if codes: + last_codes.append(codes[-1]) + cand[ids[i]] = np.asarray( + last_codes[:depth] if depth else last_codes, dtype=np.int64 + ) + + def _assign( + self, item_ids: np.ndarray, codes: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], AssignmentStats]: + """Cap buckets and reassign overflow within-band; return the final map.""" + cap = self.args.max_items_per_codebook + if cap < 1: + raise ValueError(f"max_items_per_codebook must be >= 1, got {cap}") + n, n_layers = codes.shape + last = codes[:, n_layers - 1].astype(np.int64) + + band_id = self._band_ids(codes) + order = _order_hash(item_ids, self.args.seed) + + # rank within bucket by (bucket, order-hash); collision RATE is invariant + # to which cap items are kept, but the choice is made deterministic here. + group_radix = int(last.max()) + 1 if n else 1 + group_key = band_id * group_radix + last + rank, counts_by_row = self._within_bucket_rank(group_key, order) + overflow_mask = rank >= cap + overflow_order = np.where(overflow_mask)[0] + overflow_order = overflow_order[ + np.lexsort((order[overflow_order], group_key[overflow_order])) + ] + + cand, last_size = self._candidates(item_ids, overflow_order) + last_size = max(last_size, group_radix) + + sid_key = band_id * last_size + last + slot_count: Dict[int, int] = {} + first = np.unique(group_key, return_index=True)[1] + for i in first.tolist(): + slot_count[int(band_id[i]) * last_size + int(last[i])] = int( + min(counts_by_row[i], cap) + ) + + final_key = sid_key.copy() + final_index = (rank + 1).astype(np.int64) + reassigned, unassigned = self._place_overflow( + overflow_order, + item_ids, + band_id, + last, + sid_key, + cand, + last_size, + cap, + slot_count, + final_key, + final_index, ) - scores = ( - self._array_to_pylist(batch, self.args.score_field) - if self.args.score_field and self.args.score_field in batch - else None + + keep_mask: Optional[np.ndarray] = None + if self.args.unassigned_policy == "drop" and unassigned: + keep_mask = np.ones(n, dtype=bool) + keep_mask[unassigned] = False + + final_last = (final_key % last_size).astype(np.int64) + final_codes = codes.copy() + final_codes[:, n_layers - 1] = final_last + + stats = self._stats( + n, + counts_by_row, + first, + final_key, + keep_mask, + cap, + reassigned, + len(unassigned), ) - candidates = self._array_to_pylist(batch, self.args.candidate_codebook_field) - for i, (item_id, candidate) in enumerate(zip(item_ids, candidates)): - item_key = str(item_id) - if item_key not in overflow: - continue - rows.append( - CandidateSidRow( - item_key=item_key, - candidate_codebook=self._to_codebook( - candidate, self.args.code_delimiter - ), - priority=int(priorities[i]) if priorities is not None else 1, - score=float(scores[i]) if scores is not None else 0.0, + return final_codes, final_index, keep_mask, stats + + @staticmethod + def _band_ids(codes: np.ndarray) -> np.ndarray: + """Dense integer id per distinct ``(prefix)`` band (all layers but last).""" + n, n_layers = codes.shape + if n_layers == 1: + return np.zeros(n, dtype=np.int64) + _, band_id = np.unique(codes[:, : n_layers - 1], axis=0, return_inverse=True) + return band_id.astype(np.int64).reshape(-1) + + @staticmethod + def _within_bucket_rank( + group_key: np.ndarray, order: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """Per-row within-bucket rank (0-based) and its bucket size.""" + n = group_key.shape[0] + sort_idx = np.lexsort((order, group_key)) + sk = group_key[sort_idx] + _, first, counts = np.unique(sk, return_index=True, return_counts=True) + rank_sorted = np.arange(n) - np.repeat(first, counts) + counts_sorted = np.repeat(counts, counts) + rank = np.empty(n, dtype=np.int64) + size = np.empty(n, dtype=np.int64) + rank[sort_idx] = rank_sorted + size[sort_idx] = counts_sorted + return rank, size + + def _candidates( + self, + item_ids: np.ndarray, + overflow_order: np.ndarray, + ) -> Tuple[Dict[Any, np.ndarray], int]: + """Build per-overflow-item last-code candidate lists and the code space.""" + if self.args.strategy == "random": + size = self.args.random_last_layer_size + if size is None or size < 2: + raise ValueError( + "strategy='random' requires --random_last_layer_size >= 2." ) + num = min(self.args.random_num_candidates, size - 1) + draws = self._random_draws(overflow_order, item_ids, num, size) + cand = {item_ids[i]: draws[j] for j, i in enumerate(overflow_order)} + return cand, size + if overflow_order.size == 0: + return {}, 1 + cand = self._load_candidate_last(item_ids[overflow_order]) + if not cand: + raise ValueError( + "map has overflow items but candidate_codes yielded no candidates." + ) + max_last = max((int(a.max()) for a in cand.values() if a.size), default=0) + return cand, max_last + 1 + + def _random_draws( + self, overflow_order: np.ndarray, item_ids: np.ndarray, num: int, size: int + ) -> np.ndarray: + """Order-independent random last-layer draws per overflow item, (M, num).""" + h = _order_hash(item_ids[overflow_order], self.args.seed) + k = np.arange(num, dtype=np.uint64) + with np.errstate(over="ignore"): + mixed = _splitmix64( + h[:, None] + k[None, :] * np.uint64(0x9E3779B97F4A7C15), self.args.seed + ) + return (mixed % np.uint64(size)).astype(np.int64) + + def _place_overflow( + self, + overflow_order: np.ndarray, + item_ids: np.ndarray, + band_id: np.ndarray, + last: np.ndarray, + sid_key: np.ndarray, + cand: Dict[Any, np.ndarray], + last_size: int, + cap: int, + slot_count: Dict[int, int], + final_key: np.ndarray, + final_index: np.ndarray, + ) -> Tuple[int, List[int]]: + """Greedy single pass: place each overflow item in the first free slot.""" + reassigned = 0 + unassigned: List[int] = [] + for i in overflow_order: + base = int(band_id[i]) * last_size + origin_last = int(last[i]) + placed = False + for code in cand.get(item_ids[i], ()): # ordered best-first + code = int(code) + if code == origin_last: + continue + ck = base + code + count = slot_count.get(ck, 0) + if count < cap: + slot_count[ck] = count + 1 + final_key[i] = ck + final_index[i] = count + 1 + reassigned += 1 + placed = True + break + if not placed: + unassigned.append(int(i)) + self._resolve_unassigned(unassigned, sid_key, cap, slot_count, final_index) + return reassigned, unassigned + + def _resolve_unassigned( + self, + unassigned: List[int], + sid_key: np.ndarray, + cap: int, + slot_count: Dict[int, int], + final_index: np.ndarray, + ) -> None: + """Apply --unassigned_policy to items that found no free slot.""" + if not unassigned: + return + policy = self.args.unassigned_policy + if policy == "error": + preview = ",".join(str(i) for i in unassigned[:10]) + raise RuntimeError( + f"{len(unassigned)} items could not be assigned within capacity; " + f"first unassigned row indices: {preview}" ) - return rows + if policy == "keep_original": + for i in unassigned: + key = int(sid_key[i]) + slot_count[key] = slot_count.get(key, 0) + 1 + final_index[i] = slot_count[key] @staticmethod - def _item_id_array(rows: Sequence[AssignedSidRow]) -> pa.Array: - values = [row.item_id for row in rows] - if all(isinstance(v, int) and not isinstance(v, bool) for v in values): - return pa.array(values, type=pa.int64()) - return pa.array([str(v) for v in values], type=pa.string()) + def _stats( + n: int, + counts_by_row: np.ndarray, + first: np.ndarray, + final_key: np.ndarray, + keep_mask: Optional[np.ndarray], + cap: int, + reassigned: int, + unassigned_count: int, + ) -> AssignmentStats: + """Summarize raw vs final bucket occupancy.""" + raw_counts = counts_by_row[first] + kept_key = final_key if keep_mask is None else final_key[keep_mask] + if kept_key.size: + _, fc = np.unique(kept_key, return_counts=True) + final_collision = int((fc > cap).sum()) + max_final = int(fc.max()) + else: + final_collision = 0 + max_final = 0 + return AssignmentStats( + total_items=n, + raw_collision_buckets=int((raw_counts > cap).sum()), + final_collision_buckets=final_collision, + reassigned_count=reassigned, + unassigned_count=unassigned_count, + max_final_bucket_size=max_final, + ) def _writer_type(self) -> Optional[str]: # Derive the writer from the input reader (hitrate.py idiom) when unset, @@ -663,70 +506,79 @@ def _writer_type(self) -> Optional[str]: return self._input_reader_cls_name.replace("Reader", "Writer") return None - def _write_table(self, output_path: str, columns: Dict[str, pa.Array]) -> None: - writer = create_writer( - output_path, - writer_type=self._writer_type(), - quota_name=self.args.odps_data_quota_name, - world_size=1, - ) - writer.write(columns) - writer.close() - def _is_csv_output(self) -> bool: return self._writer_type() == "CsvWriter" - def _codebook_array(self, codebooks: List[Codebook]) -> pa.Array: - # list for Parquet/ODPS; a --code_delimiter-joined string for CSV. + @staticmethod + def _item_id_array(values: np.ndarray) -> pa.Array: + if np.issubdtype(values.dtype, np.integer): + return pa.array(values, type=pa.int64()) + return pa.array([str(v) for v in values], type=pa.string()) + + def _codes_column(self, codes: np.ndarray) -> pa.Array: + """Encode an (M, n_layers) matrix as list (or a CSV string).""" + m, n_layers = codes.shape if self._is_csv_output(): - d = self.args.code_delimiter - return pa.array( - [d.join(str(c) for c in cb) for cb in codebooks], type=pa.string() - ) - return pa.array([list(cb) for cb in codebooks], type=pa.list_(pa.int64())) + cols = [ + pc.cast(pa.array(codes[:, j]), pa.string()) for j in range(n_layers) + ] + return pc.binary_join_element_wise(*cols, _CODE_SEP) + values = pa.array(codes.reshape(-1)) + offsets = pa.array(np.arange(0, (m + 1) * n_layers, n_layers, dtype=np.int32)) + return pa.ListArray.from_arrays(offsets, values) - def _write_assignments(self, rows: Sequence[AssignedSidRow]) -> None: - self._write_table( + def _write( + self, + item_ids: np.ndarray, + origin_codes: np.ndarray, + final_codes: np.ndarray, + index: np.ndarray, + keep_mask: Optional[np.ndarray], + stats: AssignmentStats, + ) -> None: + """Chunked, vectorized write of the reassigned map (and diagnostics).""" + if keep_mask is not None: + item_ids = item_ids[keep_mask] + origin_codes = origin_codes[keep_mask] + final_codes = final_codes[keep_mask] + index = index[keep_mask] + writer = create_writer( self.args.output_path, - { - "item_id": self._item_id_array(rows), - "origin_codebook": self._codebook_array( - [row.origin_codebook for row in rows] - ), - "codebook": self._codebook_array([row.codebook for row in rows]), - "index": pa.array([row.index for row in rows], type=pa.int64()), - }, + writer_type=self._writer_type(), + quota_name=self.args.odps_data_quota_name, + world_size=1, ) + n = len(item_ids) + for s in range(0, n, _WRITE_CHUNK): + e = min(s + _WRITE_CHUNK, n) + writer.write( + { + "item_id": self._item_id_array(item_ids[s:e]), + "origin_codebook": self._codes_column(origin_codes[s:e]), + "codebook": self._codes_column(final_codes[s:e]), + "index": pa.array(index[s:e], type=pa.int64()), + } + ) + writer.close() + if self.args.diagnostics_output_path: + self._write_diagnostics(stats) def _write_diagnostics(self, stats: AssignmentStats) -> None: - # asdict preserves field-declaration order, so column order is unchanged. - self._write_table( + writer = create_writer( self.args.diagnostics_output_path, - {k: pa.array([v], type=pa.int64()) for k, v in asdict(stats).items()}, + writer_type=self._writer_type(), + quota_name=self.args.odps_data_quota_name, + world_size=1, ) + writer.write( + {k: pa.array([v], type=pa.int64()) for k, v in asdict(stats).items()} + ) + writer.close() -def assign_sid_collisions( - raw_rows: Sequence[RawSidRow], - candidate_rows: Sequence[CandidateSidRow], - **kwargs: Any, -) -> Tuple[List[AssignedSidRow], AssignmentStats]: - """Assign overflow SID rows to non-full candidate codebooks. - - Thin functional entry point; ``kwargs`` are ``SidCollisionAssigner`` params. - """ - return SidCollisionAssigner(**kwargs).assign(raw_rows, candidate_rows) - - -def run(args: argparse.Namespace) -> AssignmentStats: - """Run collision prevention over the configured reader/writer backend.""" - return CollisionRunner(args).run() - - -def build_parser() -> argparse.ArgumentParser: - """Build the command line argument parser.""" +if __name__ == "__main__": parser = argparse.ArgumentParser( - description="Prevent SID codebook collisions with explicit candidates." + description="Prevent SID codebook collisions (vectorized, within-band)." ) parser.add_argument("--input_path", required=True) parser.add_argument("--output_path", required=True) @@ -743,31 +595,17 @@ def build_parser() -> argparse.ArgumentParser: help="Output writer; defaults to matching the input reader " "(CsvReader -> CsvWriter, etc.).", ) - parser.add_argument("--batch_size", type=int, default=4096) + parser.add_argument("--batch_size", type=int, default=100000) parser.add_argument("--item_id_field", default="item_id") parser.add_argument("--code_field", default="codes") + parser.add_argument("--candidate_codes_field", default="candidate_codes") parser.add_argument( - "--code_delimiter", - default=",", - help="CSV-only: delimiter splitting/joining int codes in a codebook " - "string. Parquet/ODPS use list columns directly.", - ) - parser.add_argument("--candidate_codebook_field", default="candidate_codebook") - parser.add_argument( - "--compact_candidate_field", + "--candidate_depth", + type=int, default=None, - help="Compact string/list candidate field.", - ) - parser.add_argument( - "--candidate_delimiter", - default="|", - help="Delimiter for compact string candidate lists.", + help="Cap on candidate last-codes tried per overflow item (default: all).", ) - parser.add_argument("--priority_field", default="priority") - parser.add_argument("--score_field", default="score") - parser.add_argument("--score_order", choices=["lower", "higher"], default="lower") parser.add_argument("--max_items_per_codebook", type=int, required=True) - parser.add_argument("--max_iters", type=int, default=50) parser.add_argument("--seed", type=int, default=2026) parser.add_argument( "--unassigned_policy", @@ -778,9 +616,8 @@ def build_parser() -> argparse.ArgumentParser: "--strategy", choices=["candidate", "random"], default="candidate", - help="Reassignment strategy: 'candidate' uses explicit nearest-neighbor " - "candidate rows; 'random' draws random within-band last-layer codes and " - "needs no candidate input.", + help="Reassignment strategy: 'candidate' uses model candidate_codes; " + "'random' draws random within-band last-layer codes (no candidates).", ) parser.add_argument( "--random_last_layer_size", @@ -794,15 +631,11 @@ def build_parser() -> argparse.ArgumentParser: default=64, help="Random last-layer codes drawn per overflow item for --strategy random.", ) + parser.add_argument( + "--rate_only", + action="store_true", + help="Compute + log stats only; skip writing the map (avoids the map-write " + "cost when only the collision rate is needed).", + ) parser.add_argument("--odps_data_quota_name", default="pay-as-you-go") - return parser - - -def main() -> None: - """Command line entrypoint.""" - args = build_parser().parse_args() - run(args) - - -if __name__ == "__main__": - main() + CollisionRunner(parser.parse_args()).run() diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 203c5cfe..120a0dfc 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -10,23 +10,53 @@ # limitations under the License. import os +import random import shutil import tempfile import unittest -from collections import Counter, defaultdict +from argparse import Namespace +from collections import Counter import pyarrow as pa from pyarrow import csv, parquet -from tzrec.tools.sid.collision_prevention import ( - CandidateSidRow, - RawSidRow, - assign_sid_collisions, - build_parser, - run, +from tzrec.tools.sid.collision_prevention import CollisionRunner + +# Defaults mirror the __main__ argparse; tests override only what they exercise. +_DEFAULTS = dict( + input_path=None, + output_path=None, + diagnostics_output_path=None, + reader_type=None, + writer_type=None, + batch_size=100000, + item_id_field="item_id", + code_field="codes", + candidate_codes_field="candidate_codes", + candidate_depth=None, + max_items_per_codebook=2, + seed=2026, + unassigned_policy="error", + strategy="candidate", + random_last_layer_size=None, + random_num_candidates=64, + rate_only=False, + odps_data_quota_name="pay-as-you-go", ) +def _parquet(path, item_ids, codes, candidate_codes=None): + cols = { + "item_id": pa.array(item_ids, type=pa.int64()), + "codes": pa.array(codes, type=pa.list_(pa.int64())), + } + if candidate_codes is not None: + cols["candidate_codes"] = pa.array( + candidate_codes, type=pa.list_(pa.list_(pa.int64())) + ) + parquet.write_table(pa.table(cols), path) + + class SidCollisionPreventionTest(unittest.TestCase): def setUp(self) -> None: if not os.path.exists("./tmp"): @@ -37,443 +67,232 @@ def tearDown(self) -> None: if os.path.exists(self.test_dir): shutil.rmtree(self.test_dir) - # ---- assigner (codebook is a tuple[int]) ---- - - def test_assign_sid_collisions_respects_capacity(self) -> None: - raw_rows = [ - RawSidRow("item_0", (1,)), - RawSidRow("item_1", (1,)), - RawSidRow("item_2", (1,)), - RawSidRow("item_3", (2,)), - ] - candidate_rows = [ - CandidateSidRow("item_0", (3,), 1, 0.1), - CandidateSidRow("item_1", (3,), 1, 0.1), - CandidateSidRow("item_2", (3,), 1, 0.1), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, candidate_rows, capacity=2, seed=7 - ) - - self.assertEqual(len(assigned), 4) - self.assertEqual(len({row.item_key for row in assigned}), 4) - self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 2) - self.assertEqual(stats.raw_collision_buckets, 1) - self.assertEqual(stats.final_collision_buckets, 0) - self.assertEqual(stats.reassigned_count, 1) - self.assertEqual(stats.unassigned_count, 0) - - def test_random_strategy_reassigns_within_band_without_candidates(self) -> None: - # Three items collide on SID (1,2,3); capacity 1. - raw_rows = [ - RawSidRow("item_0", (1, 2, 3)), - RawSidRow("item_1", (1, 2, 3)), - RawSidRow("item_2", (1, 2, 3)), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, - [], # random needs no candidate input - capacity=1, - strategy="random", - random_last_layer_size=16, - seed=7, - ) + def _run(self, inp, out, **kw): + args = dict(_DEFAULTS) + args.update(input_path=inp, output_path=out) + args.update(kw) + return CollisionRunner(Namespace(**args)).run() - self.assertEqual(len(assigned), 3) - codebooks = {row.item_key: row.codebook for row in assigned} - # every item keeps its (1,2) band; only the last layer varies ... - for codebook in codebooks.values(): - self.assertEqual(codebook[:2], (1, 2)) - # ... capacity 1 keeps one item at the origin SID and reassigns the other - # two to distinct random last-layer codes -> near-injective. - self.assertEqual(len(set(codebooks.values())), 3) - self.assertEqual(sum(1 for cb in codebooks.values() if cb == (1, 2, 3)), 1) - self.assertEqual(stats.final_collision_buckets, 0) - self.assertEqual(stats.reassigned_count, 2) - self.assertEqual(stats.unassigned_count, 0) + def _read_parquet(self, out_dir): + return parquet.read_table(os.path.join(out_dir, "part-0.parquet")).to_pydict() - def test_random_strategy_is_deterministic_given_seed(self) -> None: - raw_rows = [RawSidRow(f"item_{i}", (0, 0, 0)) for i in range(4)] - kwargs = dict(capacity=1, strategy="random", random_last_layer_size=64, seed=11) - first, _ = assign_sid_collisions(raw_rows, [], **kwargs) - second, _ = assign_sid_collisions(raw_rows, [], **kwargs) - self.assertEqual( - {r.item_key: r.codebook for r in first}, - {r.item_key: r.codebook for r in second}, - ) + # ---- candidate strategy ---- - def test_random_strategy_requires_last_layer_size(self) -> None: - raw_rows = [RawSidRow("item_0", (1, 2, 3))] - with self.assertRaisesRegex(ValueError, "random_last_layer_size"): - assign_sid_collisions(raw_rows, [], capacity=1, strategy="random") - - def test_missing_candidates_errors_on_overflow(self) -> None: - raw_rows = [RawSidRow("item_0", (1,)), RawSidRow("item_1", (1,))] - with self.assertRaisesRegex(ValueError, "no explicit candidate input"): - assign_sid_collisions(raw_rows, [], capacity=1) - - def test_duplicate_candidates_do_not_consume_capacity_twice(self) -> None: - raw_rows = [ - RawSidRow("item_0", (1,)), - RawSidRow("item_1", (1,)), - RawSidRow("item_2", (1,)), - ] - candidate_rows = [ - CandidateSidRow("item_0", (3,), 1, 0.1), - CandidateSidRow("item_0", (3,), 2, 0.2), - CandidateSidRow("item_1", (3,), 1, 0.1), - CandidateSidRow("item_2", (3,), 1, 0.1), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, candidate_rows, capacity=2, seed=7 + def test_candidate_reassigns_within_band(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + # 5 items in bucket (0,0); candidates vary only the last layer within band 0. + _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1], [0, 2], [0, 3]]] * 5) + stats = self._run( + inp, out, max_items_per_codebook=2, unassigned_policy="keep_original" ) - - self.assertEqual(len(assigned), 3) + self.assertEqual(stats.total_items, 5) + self.assertEqual(stats.reassigned_count, 3) self.assertEqual(stats.unassigned_count, 0) - self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 2) - - def test_local_reassignment_indices_are_unique_and_dense(self) -> None: - # After reassignment every (codebook, index) pair must be unique and each - # bucket's indices must form a contiguous 1..N run. - raw_rows = [RawSidRow(f"item_{i}", (1,)) for i in range(4)] - candidate_rows = [ - CandidateSidRow("item_0", (2,), 1, 0.1), - CandidateSidRow("item_1", (2,), 1, 0.2), - CandidateSidRow("item_2", (2,), 1, 0.3), - CandidateSidRow("item_3", (2,), 1, 0.4), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, candidate_rows, capacity=2, seed=7 + self.assertEqual(stats.max_final_bucket_size, 2) + d = self._read_parquet(out) + # every final SID keeps prefix 0 (stays in band) ... + self.assertTrue(all(cb[0] == 0 for cb in d["codebook"])) + # ... and no bucket exceeds the cap. + self.assertLessEqual(max(Counter(map(tuple, d["codebook"])).values()), 2) + + def test_output_is_list_int64(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0, 1, 2], [[0, 0]] * 3, [[[0, 1]]] * 3) + self._run(inp, out, max_items_per_codebook=2, unassigned_policy="keep_original") + schema = parquet.read_table(os.path.join(out, "part-0.parquet")).schema + self.assertEqual(schema.field("codebook").type, pa.list_(pa.int64())) + self.assertEqual(schema.field("origin_codebook").type, pa.list_(pa.int64())) + + def test_keep_original_when_only_origin_candidate(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + # candidate == origin last -> skipped -> nothing to place. + _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) + stats = self._run( + inp, out, max_items_per_codebook=2, unassigned_policy="keep_original" ) - - self.assertEqual(len(assigned), 4) - self.assertEqual(stats.unassigned_count, 0) - pairs = [(row.codebook, row.index) for row in assigned] - self.assertEqual(len(pairs), len(set(pairs))) - by_codebook = defaultdict(list) - for codebook, index in pairs: - by_codebook[codebook].append(index) - for indices in by_codebook.values(): - self.assertEqual(sorted(indices), list(range(1, len(indices) + 1))) - - def test_local_drop_policy_omits_unplaceable_items(self) -> None: - # (1,) (capacity 1) keeps one item; the other two overflow and both want - # (2,), which fits only one -> one item is unplaceable (drop branch). - raw_rows = [ - RawSidRow("item_0", (1,)), - RawSidRow("item_1", (1,)), - RawSidRow("item_2", (1,)), - ] - candidate_rows = [ - CandidateSidRow("item_0", (2,), 1, 0.1), - CandidateSidRow("item_1", (2,), 1, 0.2), - CandidateSidRow("item_2", (2,), 1, 0.3), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, candidate_rows, capacity=1, seed=7, unassigned_policy="drop" + self.assertEqual(stats.reassigned_count, 0) + self.assertEqual(stats.unassigned_count, 3) + d = self._read_parquet(out) + self.assertTrue(all(tuple(cb) == (0, 0) for cb in d["codebook"])) + + def test_error_policy_raises_on_overflow(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) + with self.assertRaises(RuntimeError): + self._run(inp, out, max_items_per_codebook=2) + + def test_drop_policy_excludes_unassigned(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) + stats = self._run(inp, out, max_items_per_codebook=2, unassigned_policy="drop") + self.assertEqual(stats.unassigned_count, 3) + d = self._read_parquet(out) + self.assertEqual(len(d["item_id"]), 2) + + def test_raises_when_overflow_but_no_candidates(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(5)), [[0, 0]] * 5) # no candidate_codes column + with self.assertRaises(ValueError): + self._run(inp, out, max_items_per_codebook=2) + + # ---- random strategy ---- + + def test_random_reassigns_within_band(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(5)), [[0, 0]] * 5) + stats = self._run( + inp, + out, + max_items_per_codebook=2, + strategy="random", + random_last_layer_size=8, + unassigned_policy="keep_original", ) - - self.assertEqual(stats.unassigned_count, 1) - self.assertEqual(len(assigned), 2) - self.assertLessEqual(max(Counter(row.codebook for row in assigned).values()), 1) - - def test_local_keep_original_readds_over_capacity(self) -> None: - raw_rows = [ - RawSidRow("item_0", (1,)), - RawSidRow("item_1", (1,)), - RawSidRow("item_2", (1,)), - ] - candidate_rows = [ - CandidateSidRow("item_0", (2,), 1, 0.1), - CandidateSidRow("item_1", (2,), 1, 0.2), - CandidateSidRow("item_2", (2,), 1, 0.3), - ] - - assigned, stats = assign_sid_collisions( - raw_rows, - candidate_rows, - capacity=1, - seed=7, + self.assertEqual(stats.reassigned_count, 3) + d = self._read_parquet(out) + self.assertTrue(all(cb[0] == 0 for cb in d["codebook"])) + self.assertLessEqual(max(Counter(map(tuple, d["codebook"])).values()), 2) + + def test_random_requires_last_layer_size(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(5)), [[0, 0]] * 5) + with self.assertRaises(ValueError): + self._run(inp, out, max_items_per_codebook=2, strategy="random") + + def test_single_layer_sid(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(4)), [[0]] * 4) + stats = self._run( + inp, + out, + max_items_per_codebook=1, + strategy="random", + random_last_layer_size=16, unassigned_policy="keep_original", ) - - self.assertEqual(stats.unassigned_count, 0) - self.assertEqual(len(assigned), 3) - # keep_original is the only policy allowed to exceed capacity. - self.assertEqual(Counter(row.codebook for row in assigned)[(1,)], 2) - self.assertEqual(stats.max_final_bucket_size, 2) - - def test_local_error_policy_raises_on_unplaceable(self) -> None: - raw_rows = [ - RawSidRow("item_0", (1,)), - RawSidRow("item_1", (1,)), - RawSidRow("item_2", (1,)), - ] - candidate_rows = [ - CandidateSidRow("item_0", (2,), 1, 0.1), - CandidateSidRow("item_1", (2,), 1, 0.2), - CandidateSidRow("item_2", (2,), 1, 0.3), - ] - - with self.assertRaisesRegex(RuntimeError, "could not be assigned"): - assign_sid_collisions( - raw_rows, - candidate_rows, - capacity=1, + self.assertEqual(stats.reassigned_count, 3) + d = self._read_parquet(out) + self.assertEqual(len({tuple(cb) for cb in d["codebook"]}), 4) + + # ---- determinism ---- + + def test_deterministic_and_order_independent(self) -> None: + rng = random.Random(0) + n = 200 + item_ids = list(range(n)) + codes = [[rng.randrange(3), rng.randrange(2)] for _ in item_ids] + cands = [[[c[0], j] for j in range(8)] for c in codes] + + def decisions(order): + inp = os.path.join(self.test_dir, f"in_{order[0]}_{len(order)}.parquet") + out = os.path.join(self.test_dir, f"out_{order[0]}_{len(order)}") + _parquet( + inp, + [item_ids[i] for i in order], + [codes[i] for i in order], + [cands[i] for i in order], + ) + self._run( + inp, + out, + max_items_per_codebook=2, + unassigned_policy="keep_original", seed=7, - unassigned_policy="error", ) + d = self._read_parquet(out) + return { + d["item_id"][i]: tuple(d["codebook"][i]) + for i in range(len(d["item_id"])) + } - def test_local_score_order_higher_prefers_high_score(self) -> None: - # One item overflows (1,) (capacity 1); it can go to (2,) (score 0.1) or - # (3,) (score 0.9). score_order flips which score wins. - raw_rows = [RawSidRow("item_0", (1,)), RawSidRow("item_1", (1,))] - candidate_rows = [ - CandidateSidRow("item_0", (2,), 1, 0.1), - CandidateSidRow("item_0", (3,), 1, 0.9), - CandidateSidRow("item_1", (2,), 1, 0.1), - CandidateSidRow("item_1", (3,), 1, 0.9), - ] - - lower, _ = assign_sid_collisions( - raw_rows, candidate_rows, capacity=1, seed=7, score_order="lower" - ) - higher, _ = assign_sid_collisions( - raw_rows, candidate_rows, capacity=1, seed=7, score_order="higher" - ) - - reassigned_lower = next( - row for row in lower if row.origin_codebook != row.codebook - ) - reassigned_higher = next( - row for row in higher if row.origin_codebook != row.codebook - ) - self.assertEqual(reassigned_lower.codebook, (2,)) - self.assertEqual(reassigned_higher.codebook, (3,)) - - # ---- file backends: Parquet/ODPS list, CSV string fallback ---- + base = decisions(list(range(n))) + self.assertEqual(base, decisions(list(range(n)))) # run-twice + shuffled = list(range(n)) + rng.shuffle(shuffled) + self.assertEqual(base, decisions(shuffled)) # order-independent - def test_local_parquet_writes_list_codebooks(self) -> None: - raw_path = os.path.join(self.test_dir, "raw.parquet") - out_dir = os.path.join(self.test_dir, "out_parquet") - parquet.write_table( - pa.table( - { - "item_id": pa.array([1, 2, 3], type=pa.int64()), - "codes": pa.array( - [[1, 2], [1, 2], [1, 2]], type=pa.list_(pa.int64()) - ), - "candidate_codebook": pa.array( - [[1, 3], [1, 3], [1, 3]], type=pa.list_(pa.int64()) - ), - "priority": [1, 1, 1], - "score": [0.1, 0.1, 0.1], - } - ), - raw_path, - ) + # ---- CSV backend ---- - args = build_parser().parse_args( - [ - "--input_path", - raw_path, - "--output_path", - out_dir, - "--writer_type", - "ParquetWriter", - "--max_items_per_codebook", - "2", - ] - ) - stats = run(args) - - self.assertEqual(stats.reassigned_count, 1) - result = parquet.read_table(os.path.join(out_dir, "part-0.parquet")) - # Parquet/ODPS keep codebooks as list. - self.assertEqual(result.schema.field("codebook").type, pa.list_(pa.int64())) - origins = {tuple(x) for x in result["origin_codebook"].to_pylist()} - self.assertIn((1, 2), origins) - codebooks = [tuple(x) for x in result["codebook"].to_pylist()] - self.assertLessEqual(max(Counter(codebooks).values()), 2) - - def test_local_csv_writes_codebooks_as_strings(self) -> None: - raw_path = os.path.join(self.test_dir, "raw.csv") - out_dir = os.path.join(self.test_dir, "out") + def test_csv_backend_within_band(self) -> None: + inp = os.path.join(self.test_dir, "in.csv") + out = os.path.join(self.test_dir, "out") csv.write_csv( pa.table( { - # multi-layer so the codebook string keeps a comma and CSV - # read-back doesn't re-infer it as int. - "item_id": ["1", "2", "3"], - "codes": ["1,2", "1,2", "1,2"], - "candidate_codebook": ["3,4", "3,4", "3,4"], - "priority": [1, 1, 1], - "score": [0.1, 0.1, 0.1], + "item_id": ["0", "1", "2", "3", "4"], + "codes": ["0|0"] * 5, + "candidate_codes": ["0|1;0|2;0|3"] * 5, } ), - raw_path, + inp, ) - - args = build_parser().parse_args( - [ - "--input_path", - raw_path, - "--output_path", - out_dir, - "--reader_type", - "CsvReader", - "--writer_type", - "CsvWriter", - "--max_items_per_codebook", - "2", - ] + stats = self._run( + inp, + out, + reader_type="CsvReader", + writer_type="CsvWriter", + max_items_per_codebook=2, + unassigned_policy="keep_original", ) - stats = run(args) - - self.assertEqual(stats.reassigned_count, 1) - result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) - # CSV fallback encodes codebooks as delimited strings. - self.assertEqual(result.schema.field("origin_codebook").type, pa.string()) + self.assertEqual(stats.reassigned_count, 3) + result = csv.read_csv(os.path.join(out, "part-0.csv")) self.assertEqual(result.schema.field("codebook").type, pa.string()) - self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) - - def test_writer_type_defaults_to_matching_the_reader(self) -> None: - raw_path = os.path.join(self.test_dir, "raw_derive.csv") - out_dir = os.path.join(self.test_dir, "out_derive") - csv.write_csv( - pa.table({"item_id": ["1", "2", "3"], "codes": ["1", "1", "1"]}), - raw_path, - ) - args = build_parser().parse_args( - [ - "--input_path", - raw_path, - "--output_path", - out_dir, - "--max_items_per_codebook", - "3", - ] - ) - run(args) - result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) - self.assertEqual(result.num_rows, 3) - - def test_local_csv_accepts_compact_candidates(self) -> None: - raw_path = os.path.join(self.test_dir, "raw_compact.csv") - out_dir = os.path.join(self.test_dir, "out_compact") - csv.write_csv( - pa.table( - { - "item_id": ["1", "2", "3"], - "codes": ["1,2", "1,2", "1,2"], - "candidates": ["1|3", "1|3", "1|3"], - } - ), - raw_path, + self.assertTrue(all(s.startswith("0|") for s in result["codebook"].to_pylist())) + + def test_writer_defaults_to_reader(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0, 1, 2], [[0, 0]] * 3, [[[0, 1]]] * 3) + self._run(inp, out, max_items_per_codebook=3) + self.assertTrue(os.path.exists(os.path.join(out, "part-0.parquet"))) + + # ---- misc ---- + + def test_rate_only_skips_write(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1], [0, 2], [0, 3]]] * 5) + stats = self._run( + inp, + out, + max_items_per_codebook=2, + unassigned_policy="keep_original", + rate_only=True, ) - - args = build_parser().parse_args( - [ - "--input_path", - raw_path, - "--output_path", - out_dir, - "--reader_type", - "CsvReader", - "--writer_type", - "CsvWriter", - "--compact_candidate_field", - "candidates", - "--max_items_per_codebook", - "2", - ] + self.assertEqual(stats.reassigned_count, 3) + self.assertFalse(os.path.exists(os.path.join(out, "part-0.parquet"))) + + def test_diagnostics_output(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + diag = os.path.join(self.test_dir, "diag") + _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1], [0, 2], [0, 3]]] * 5) + self._run( + inp, + out, + max_items_per_codebook=2, + unassigned_policy="keep_original", + diagnostics_output_path=diag, ) - stats = run(args) - - self.assertEqual(stats.reassigned_count, 1) - result = csv.read_csv(os.path.join(out_dir, "part-0.csv")) - self.assertIn("1,2", set(result["origin_codebook"].to_pylist())) - self.assertTrue({"1", "3"} & set(result["codebook"].to_pylist())) - self.assertLessEqual(max(Counter(result["codebook"].to_pylist()).values()), 2) - - def test_csv_and_parquet_yield_identical_decisions(self) -> None: - # The same logical data as CSV (delimited strings) and Parquet - # (list) must normalize to the same tuple codebooks, so the - # assignment decisions are identical across backends. - item_ids = list(range(20)) - origin = [[i % 3, i % 2] for i in item_ids] # heavy collision on few buckets - cand = [[9, i % 5] for i in item_ids] - - def decisions(mode: str) -> dict: - if mode == "parquet": - path = os.path.join(self.test_dir, "eq.parquet") - parquet.write_table( - pa.table( - { - "item_id": pa.array(item_ids, type=pa.int64()), - "codes": pa.array(origin, type=pa.list_(pa.int64())), - "candidate_codebook": pa.array( - cand, type=pa.list_(pa.int64()) - ), - } - ), - path, - ) - reader, writer = "ParquetReader", "ParquetWriter" - else: - path = os.path.join(self.test_dir, "eq.csv") - csv.write_csv( - pa.table( - { - "item_id": [str(i) for i in item_ids], - "codes": [",".join(map(str, c)) for c in origin], - "candidate_codebook": [",".join(map(str, c)) for c in cand], - } - ), - path, - ) - reader, writer = "CsvReader", "CsvWriter" - out = os.path.join(self.test_dir, f"eq_out_{mode}") - run( - build_parser().parse_args( - [ - "--input_path", - path, - "--output_path", - out, - "--reader_type", - reader, - "--writer_type", - writer, - "--max_items_per_codebook", - "2", - "--unassigned_policy", - "keep_original", - "--seed", - "5", - ] - ) - ) - if mode == "parquet": - d = parquet.read_table(os.path.join(out, "part-0.parquet")).to_pydict() - return { - int(d["item_id"][i]): tuple(d["codebook"][i]) - for i in range(len(d["item_id"])) - } - d = csv.read_csv(os.path.join(out, "part-0.csv")).to_pydict() - return { - int(d["item_id"][i]): tuple( - int(x) for x in str(d["codebook"][i]).split(",") - ) - for i in range(len(d["item_id"])) - } - - self.assertEqual(decisions("parquet"), decisions("csv")) + d = parquet.read_table(os.path.join(diag, "part-0.parquet")).to_pydict() + self.assertEqual(d["total_items"][0], 5) + self.assertEqual(d["reassigned_count"][0], 3) + + def test_empty_input_raises(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [], []) + with self.assertRaises(ValueError): + self._run(inp, out, max_items_per_codebook=2) if __name__ == "__main__": From ba2bdf1d7d08eeddf7c034e048f0d75ced5cefeb Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Fri, 10 Jul 2026 03:08:50 +0000 Subject: [PATCH 14/29] [refactor] SID collision tool: declare shape via --codebook; private config Declare the SID shape with a required --codebook (comma-separated per-layer sizes) instead of inferring n_layers, band-packing radices, and the last-layer code space from the data; this removes the "batch missing the max code" edge case and subsumes --random_last_layer_size (now codebook[-1]). --seed becomes a fixed module constant since the collision rate is seed-invariant. Config args are parsed into private attributes and validated at construction; the vectorized paths pick up review cleanups (in-place band packing, Python-native candidate keys, a param-less _read). Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 458 ++++++++++--------- tzrec/tools/sid/collision_prevention_test.py | 13 +- 2 files changed, 254 insertions(+), 217 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index a868f449..0b655e8d 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -22,7 +22,7 @@ Items with no free slot fall back per ``--unassigned_policy``. The hot path is numpy/Arrow-vectorized so a hundred-million-row map fits in one pass; results -are deterministic given ``--seed`` and independent of input row order. I/O goes +are deterministic and independent of input row order. I/O goes through ``create_reader`` / ``create_writer`` so CSV, Parquet, and ODPS all work; outputs are written in chunks to stay under Arrow's 2^31 list-offset limit. """ @@ -38,10 +38,14 @@ # Register the local reader/writer classes used through create_reader/writer. import tzrec.datasets.csv_dataset # noqa: F401 import tzrec.datasets.parquet_dataset # noqa: F401 -from tzrec.datasets.dataset import create_reader, create_writer +from tzrec.datasets.dataset import BaseReader, create_reader, create_writer from tzrec.utils.logging_util import logger -_MASK64 = np.uint64((1 << 64) - 1) +_MASK64 = (1 << 64) - 1 + +# Fixed hash seed for the order-independent tie-break. Not a CLI knob: the +# collision rate is seed-invariant, so runs stay reproducible without one. +_SEED = 2026 # CSV-fallback delimiters (Parquet/ODPS carry codes as native list, so # these are unused there). Both are non-comma so a codebook cell never collides @@ -66,37 +70,34 @@ class AssignmentStats: max_final_bucket_size: int -def _splitmix64(x: np.ndarray, seed: int) -> np.ndarray: +def _splitmix64(x: np.ndarray) -> np.ndarray: """Vectorized order-independent SplitMix64 hash of a uint64 array. - A pure function of the input values (not their position), so it gives a - stable, seedable tie-break that is invariant to input row order -- unlike a - read-order index -- while staying fully vectorized. + A pure function of the input values (not their position), mixed with the + module ``_SEED``, so it gives a stable tie-break invariant to input row order + -- unlike a read-order index -- while staying fully vectorized. Args: x (np.ndarray): uint64 values to hash. - seed (int): mixing seed. Returns: np.ndarray: uint64 hashes, same shape as ``x``. """ with np.errstate(over="ignore"): - z = x.astype(np.uint64) + np.uint64((seed * 0x9E3779B97F4A7C15) & int(_MASK64)) + z = x.astype(np.uint64) + np.uint64((_SEED * 0x9E3779B97F4A7C15) & _MASK64) z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) return z ^ (z >> np.uint64(31)) -def _order_hash(item_ids: np.ndarray, seed: int) -> np.ndarray: +def _order_hash(item_ids: np.ndarray) -> np.ndarray: """Per-item order-independent tie-break hash (uint64), vectorized. Integer ids hash directly; string/object ids are folded to uint64 via - ``pandas.util.hash_array`` first. Both then pass through :func:`_splitmix64` - so the ``seed`` is mixed in uniformly. + ``pandas.util.hash_array`` first. Both then pass through :func:`_splitmix64`. Args: item_ids (np.ndarray): item id values. - seed (int): mixing seed. Returns: np.ndarray: uint64 per-item hashes. @@ -107,7 +108,7 @@ def _order_hash(item_ids: np.ndarray, seed: int) -> np.ndarray: import pandas as pd base = pd.util.hash_array(np.asarray(item_ids, dtype=object)) - return _splitmix64(base, seed) + return _splitmix64(base) class CollisionRunner: @@ -118,34 +119,68 @@ class CollisionRunner: """ def __init__(self, args: argparse.Namespace) -> None: - self.args = args - # Class name of the input reader, captured on the first read; used to - # derive the writer type when --writer_type is unset. - self._input_reader_cls_name: Optional[str] = None + self._input_path = args.input_path + self._output_path = args.output_path + self._diagnostics_output_path = args.diagnostics_output_path + self._reader_type = args.reader_type + self._writer_type = args.writer_type + self._batch_size = args.batch_size + self._item_id_field = args.item_id_field + self._code_field = args.code_field + self._candidate_codes_field = args.candidate_codes_field + self._candidate_depth = args.candidate_depth + self._max_items_per_codebook = args.max_items_per_codebook + if self._max_items_per_codebook < 1: + raise ValueError( + "--max_items_per_codebook must be >= 1, got " + f"{self._max_items_per_codebook}." + ) + self._unassigned_policy = args.unassigned_policy + self._strategy = args.strategy + self._codebook = [int(x) for x in args.codebook.split(",") if x.strip()] + if not self._codebook: + raise ValueError("--codebook must list at least one per-layer size.") + self._random_num_candidates = args.random_num_candidates + self._rate_only = args.rate_only + self._odps_data_quota_name = args.odps_data_quota_name + # The writer type (and its is-CSV flag) can depend on the input reader + # class when --writer_type is unset, so they are finalized once at the + # first read rather than here; declared now so the attributes always exist. + self._writer_type_str: Optional[str] = None + self._is_csv: bool = False def run(self) -> AssignmentStats: """Read the SID map, assign, and write the reassigned map + stats.""" item_ids, codes = self._load_codes() final_codes, index, keep_mask, stats = self._assign(item_ids, codes) - if not self.args.rate_only: + if not self._rate_only: self._write(item_ids, codes, final_codes, index, keep_mask, stats) else: logger.info("rate_only: skipping map write") logger.info("SID collision prevention finished: %s", stats) return stats - def _read( - self, selected_cols: Optional[List[str]], capture_reader_cls: bool = False - ) -> Iterable[Dict[str, pa.Array]]: - reader = create_reader( - input_path=self.args.input_path, - batch_size=self.args.batch_size, + def _make_reader(self, selected_cols: List[str]) -> BaseReader: + """Open the input reader projecting ``selected_cols``.""" + return create_reader( + input_path=self._input_path, + batch_size=self._batch_size, selected_cols=selected_cols, - reader_type=self.args.reader_type, - quota_name=self.args.odps_data_quota_name, + reader_type=self._reader_type, + quota_name=self._odps_data_quota_name, ) - if capture_reader_cls: - self._input_reader_cls_name = reader.__class__.__name__ + + def _read(self) -> Iterable[Dict[str, pa.Array]]: + """Stream item-id + code batches, caching the resolved output backend.""" + reader = self._make_reader([self._item_id_field, self._code_field]) + # Writer defaults to matching the input reader (hitrate.py idiom) when + # --writer_type is unset; ODPS still resolves to OdpsWriter via + # create_writer's path detection regardless. + reader_cls_name = reader.__class__.__name__ + self._writer_type_str = self._writer_type or ( + reader_cls_name.replace("Reader", "Writer") + ) + self._is_csv = self._writer_type_str == "CsvWriter" yield from reader.to_batches() @staticmethod @@ -157,9 +192,14 @@ def _codes_matrix(arr: pa.Array) -> np.ndarray: """ if pa.types.is_list(arr.type) or pa.types.is_large_list(arr.type): n = len(arr) - flat = arr.flatten().to_numpy(zero_copy_only=False).astype(np.int64) + flat = ( + arr.flatten() + .to_numpy(zero_copy_only=False) + .astype(np.int64, copy=False) + ) elif pa.types.is_integer(arr.type): - return arr.to_numpy(zero_copy_only=False).astype(np.int64).reshape(-1, 1) + arr_np = arr.to_numpy(zero_copy_only=False).astype(np.int64, copy=False) + return arr_np.reshape(-1, 1) else: parts = pc.split_pattern(arr, _CODE_SEP) n = len(parts) @@ -175,24 +215,20 @@ def _load_codes(self) -> Tuple[np.ndarray, np.ndarray]: """Stream the map into an item-id array and an (N, n_layers) code matrix.""" id_chunks: List[np.ndarray] = [] code_chunks: List[np.ndarray] = [] - for batch in self._read( - [self.args.item_id_field, self.args.code_field], capture_reader_cls=True - ): - id_chunks.append( - batch[self.args.item_id_field].to_numpy(zero_copy_only=False) - ) - code_chunks.append(self._codes_matrix(batch[self.args.code_field])) + for batch in self._read(): + id_chunks.append(batch[self._item_id_field].to_numpy(zero_copy_only=False)) + code_chunks.append(self._codes_matrix(batch[self._code_field])) if not id_chunks: raise ValueError("SID input is empty.") item_ids = np.concatenate(id_chunks) + del id_chunks codes = np.concatenate(code_chunks, axis=0) + del code_chunks if codes.shape[1] < 1: raise ValueError("SID codes must have at least one layer.") return item_ids, codes - def _load_candidate_last( - self, overflow_id_arr: np.ndarray - ) -> Dict[Any, np.ndarray]: + def _load_candidate_last(self, overflow_id_arr: np.ndarray) -> Dict[Any, List[int]]: """Map each overflow item id to its ordered candidate last-layer codes. Reads ``candidate_codes`` (list> for Parquet/ODPS, a @@ -200,30 +236,33 @@ def _load_candidate_last( of each candidate SID and only for overflow items, so memory scales with the overflow fraction, not the whole table. """ - depth = self.args.candidate_depth - field = self.args.candidate_codes_field - cand: Dict[Any, np.ndarray] = {} - for batch in self._read([self.args.item_id_field, field]): + field = self._candidate_codes_field + # Sort the overflow ids once; per-batch membership is then a bisect + # instead of re-sorting the whole (potentially huge) overflow set. + ov = np.sort(overflow_id_arr) + cand: Dict[Any, List[int]] = {} + for batch in self._make_reader([self._item_id_field, field]).to_batches(): if field not in batch: break - ids = batch[self.args.item_id_field].to_numpy(zero_copy_only=False) - keep = np.where(np.isin(ids, overflow_id_arr))[0] + ids = batch[self._item_id_field].to_numpy(zero_copy_only=False) + pos = np.searchsorted(ov, ids) + np.clip(pos, 0, len(ov) - 1, out=pos) + keep = np.where(ov[pos] == ids)[0] if keep.size == 0: continue col = batch[field] if pa.types.is_list(col.type) or pa.types.is_large_list(col.type): - self._collect_candidate_lists(ids, col, keep, depth, cand) + self._collect_candidate_lists(ids, col, keep, cand) else: - self._collect_candidate_strings(ids, col, keep, depth, cand) + self._collect_candidate_strings(ids, col, keep, cand) return cand - @staticmethod def _collect_candidate_lists( + self, ids: np.ndarray, col: pa.Array, keep: np.ndarray, - depth: Optional[int], - cand: Dict[Any, np.ndarray], + cand: Dict[Any, List[int]], ) -> None: """Vectorized last-code extraction from a list> batch.""" n = len(col) @@ -233,25 +272,28 @@ def _collect_candidate_lists( k = len(inner) // n if len(inner) != n * k: raise ValueError("ragged candidate_codes: all items must share topk.") - flat = inner.flatten().to_numpy(zero_copy_only=False).astype(np.int64) + flat = ( + inner.flatten().to_numpy(zero_copy_only=False).astype(np.int64, copy=False) + ) n_layers = flat.shape[0] // (n * k) last = flat.reshape(n, k, n_layers)[:, :, n_layers - 1] - if depth is not None: - last = last[:, :depth] - for i in keep.tolist(): - cand[ids[i]] = last[i] + if self._candidate_depth is not None: + last = last[:, : self._candidate_depth] + last_lists = last.tolist() # per-item Python lists: cheap to scan in the loop + for iid, i in zip(ids[keep].tolist(), keep.tolist()): + cand[iid] = last_lists[i] def _collect_candidate_strings( self, ids: np.ndarray, col: pa.Array, keep: np.ndarray, - depth: Optional[int], - cand: Dict[Any, np.ndarray], + cand: Dict[Any, List[int]], ) -> None: """Per-row last-code extraction from a CSV compact-candidate batch.""" + depth = self._candidate_depth values = col.to_pylist() - for i in keep.tolist(): + for iid, i in zip(ids[keep].tolist(), keep.tolist()): text = values[i] if not text: continue @@ -263,53 +305,50 @@ def _collect_candidate_strings( codes = [int(p) for p in part.split(_CODE_SEP) if p.strip()] if codes: last_codes.append(codes[-1]) - cand[ids[i]] = np.asarray( - last_codes[:depth] if depth else last_codes, dtype=np.int64 - ) + cand[iid] = last_codes[:depth] if depth is not None else last_codes def _assign( self, item_ids: np.ndarray, codes: np.ndarray ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], AssignmentStats]: """Cap buckets and reassign overflow within-band; return the final map.""" - cap = self.args.max_items_per_codebook - if cap < 1: - raise ValueError(f"max_items_per_codebook must be >= 1, got {cap}") + cap = self._max_items_per_codebook n, n_layers = codes.shape - last = codes[:, n_layers - 1].astype(np.int64) + if n_layers != len(self._codebook): + raise ValueError( + f"codes have {n_layers} layers but --codebook has " + f"{len(self._codebook)}." + ) + last = codes[:, n_layers - 1].astype(np.int64, copy=False) band_id = self._band_ids(codes) - order = _order_hash(item_ids, self.args.seed) - - # rank within bucket by (bucket, order-hash); collision RATE is invariant - # to which cap items are kept, but the choice is made deterministic here. - group_radix = int(last.max()) + 1 if n else 1 - group_key = band_id * group_radix + last - rank, counts_by_row = self._within_bucket_rank(group_key, order) - overflow_mask = rank >= cap - overflow_order = np.where(overflow_mask)[0] + order = _order_hash(item_ids) + + # Rank rows within their (band_id, last) bucket by order-hash: the RATE is + # invariant to which cap items are kept, but the choice is deterministic. + # One lexsort yields the rank, each bucket's representative row, and size. + rank, rep_rows, counts = self._within_bucket_rank(band_id, last, order) + overflow_order = np.flatnonzero(rank >= cap) overflow_order = overflow_order[ - np.lexsort((order[overflow_order], group_key[overflow_order])) + np.lexsort( + (order[overflow_order], last[overflow_order], band_id[overflow_order]) + ) ] - cand, last_size = self._candidates(item_ids, overflow_order) - last_size = max(last_size, group_radix) + cand = self._candidates(item_ids, overflow_order) + last_size = self._codebook[-1] sid_key = band_id * last_size + last - slot_count: Dict[int, int] = {} - first = np.unique(group_key, return_index=True)[1] - for i in first.tolist(): - slot_count[int(band_id[i]) * last_size + int(last[i])] = int( - min(counts_by_row[i], cap) - ) + slot_count: Dict[int, int] = dict( + zip(sid_key[rep_rows].tolist(), np.minimum(counts, cap).tolist()) + ) final_key = sid_key.copy() - final_index = (rank + 1).astype(np.int64) + final_index = rank + 1 reassigned, unassigned = self._place_overflow( overflow_order, item_ids, band_id, last, - sid_key, cand, last_size, cap, @@ -317,90 +356,88 @@ def _assign( final_key, final_index, ) + keep_mask = self._resolve_unassigned( + unassigned, sid_key, slot_count, final_index, n + ) - keep_mask: Optional[np.ndarray] = None - if self.args.unassigned_policy == "drop" and unassigned: - keep_mask = np.ones(n, dtype=bool) - keep_mask[unassigned] = False - - final_last = (final_key % last_size).astype(np.int64) final_codes = codes.copy() - final_codes[:, n_layers - 1] = final_last + final_codes[:, n_layers - 1] = final_key % last_size - stats = self._stats( - n, - counts_by_row, - first, - final_key, - keep_mask, - cap, - reassigned, - len(unassigned), - ) + stats = self._stats(n, counts, slot_count, cap, reassigned, len(unassigned)) return final_codes, final_index, keep_mask, stats - @staticmethod - def _band_ids(codes: np.ndarray) -> np.ndarray: - """Dense integer id per distinct ``(prefix)`` band (all layers but last).""" + def _band_ids(self, codes: np.ndarray) -> np.ndarray: + """Dense integer id per distinct ``(prefix)`` band (all layers but last). + + Packs the prefix layers into one int64 key using the declared per-layer + codebook sizes as radices, then factorizes -- far faster than + ``np.unique`` over 2-D void rows at scale, and independent of the values a + given batch happens to contain. + """ n, n_layers = codes.shape if n_layers == 1: return np.zeros(n, dtype=np.int64) - _, band_id = np.unique(codes[:, : n_layers - 1], axis=0, return_inverse=True) - return band_id.astype(np.int64).reshape(-1) + key = codes[:, 0].astype(np.int64) # fresh copy: safe to pack in place below + for j in range(1, n_layers - 1): + key *= self._codebook[j] + key += codes[:, j] + inverse = np.unique(key, return_inverse=True)[1] + return inverse.astype(np.int64, copy=False).reshape(-1) @staticmethod def _within_bucket_rank( - group_key: np.ndarray, order: np.ndarray - ) -> Tuple[np.ndarray, np.ndarray]: - """Per-row within-bucket rank (0-based) and its bucket size.""" - n = group_key.shape[0] - sort_idx = np.lexsort((order, group_key)) - sk = group_key[sort_idx] - _, first, counts = np.unique(sk, return_index=True, return_counts=True) - rank_sorted = np.arange(n) - np.repeat(first, counts) - counts_sorted = np.repeat(counts, counts) + band_id: np.ndarray, last: np.ndarray, order: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Rank rows within their ``(band_id, last)`` bucket by order-hash. + + Returns the per-row 0-based rank, each bucket's representative original + row index, and each bucket's size -- all from a single column lexsort, so + the grouping is computed once (no packed key, no per-row size scatter). + """ + n = band_id.shape[0] + sort_idx = np.lexsort((order, last, band_id)) + sb = band_id[sort_idx] + sl = last[sort_idx] + is_first = np.empty(n, dtype=bool) + is_first[0] = True + is_first[1:] = (sb[1:] != sb[:-1]) | (sl[1:] != sl[:-1]) + first_sorted = np.flatnonzero(is_first) + counts = np.diff(np.append(first_sorted, n)) rank = np.empty(n, dtype=np.int64) - size = np.empty(n, dtype=np.int64) - rank[sort_idx] = rank_sorted - size[sort_idx] = counts_sorted - return rank, size + rank[sort_idx] = np.arange(n) - np.repeat(first_sorted, counts) + return rank, sort_idx[first_sorted], counts def _candidates( self, item_ids: np.ndarray, overflow_order: np.ndarray, - ) -> Tuple[Dict[Any, np.ndarray], int]: - """Build per-overflow-item last-code candidate lists and the code space.""" - if self.args.strategy == "random": - size = self.args.random_last_layer_size - if size is None or size < 2: - raise ValueError( - "strategy='random' requires --random_last_layer_size >= 2." - ) - num = min(self.args.random_num_candidates, size - 1) + ) -> Dict[Any, List[int]]: + """Build per-overflow-item last-code candidate lists.""" + if self._strategy == "random": + size = self._codebook[-1] + if size < 2: + raise ValueError("strategy='random' requires a last codebook >= 2.") + num = min(self._random_num_candidates, size - 1) draws = self._random_draws(overflow_order, item_ids, num, size) - cand = {item_ids[i]: draws[j] for j, i in enumerate(overflow_order)} - return cand, size + ov_ids = item_ids[overflow_order].tolist() # Python-native dict keys + return {ov_ids[j]: draws[j].tolist() for j in range(len(ov_ids))} if overflow_order.size == 0: - return {}, 1 + return {} cand = self._load_candidate_last(item_ids[overflow_order]) if not cand: raise ValueError( "map has overflow items but candidate_codes yielded no candidates." ) - max_last = max((int(a.max()) for a in cand.values() if a.size), default=0) - return cand, max_last + 1 + return cand def _random_draws( self, overflow_order: np.ndarray, item_ids: np.ndarray, num: int, size: int ) -> np.ndarray: """Order-independent random last-layer draws per overflow item, (M, num).""" - h = _order_hash(item_ids[overflow_order], self.args.seed) + h = _order_hash(item_ids[overflow_order]) k = np.arange(num, dtype=np.uint64) with np.errstate(over="ignore"): - mixed = _splitmix64( - h[:, None] + k[None, :] * np.uint64(0x9E3779B97F4A7C15), self.args.seed - ) + mixed = _splitmix64(h[:, None] + k[None, :] * np.uint64(0x9E3779B97F4A7C15)) return (mixed % np.uint64(size)).astype(np.int64) def _place_overflow( @@ -409,116 +446,117 @@ def _place_overflow( item_ids: np.ndarray, band_id: np.ndarray, last: np.ndarray, - sid_key: np.ndarray, - cand: Dict[Any, np.ndarray], + cand: Dict[Any, List[int]], last_size: int, cap: int, slot_count: Dict[int, int], final_key: np.ndarray, final_index: np.ndarray, ) -> Tuple[int, List[int]]: - """Greedy single pass: place each overflow item in the first free slot.""" - reassigned = 0 + """Greedy single pass: place each overflow item in the first free slot. + + Overflow-row inputs are pre-materialized to Python scalars so the loop + avoids per-iteration 0-d numpy indexing; placements are scattered back + into ``final_key`` / ``final_index`` vectorized at the end. + """ + rows = overflow_order.tolist() + bases = (band_id[overflow_order] * last_size).tolist() + origin_lasts = last[overflow_order].tolist() + ids = item_ids[overflow_order].tolist() + placed_rows: List[int] = [] + placed_key: List[int] = [] + placed_idx: List[int] = [] unassigned: List[int] = [] - for i in overflow_order: - base = int(band_id[i]) * last_size - origin_last = int(last[i]) - placed = False - for code in cand.get(item_ids[i], ()): # ordered best-first - code = int(code) + get = slot_count.get + for k, base in enumerate(bases): + origin_last = origin_lasts[k] + for code in cand.get(ids[k], ()): # ordered best-first if code == origin_last: continue ck = base + code - count = slot_count.get(ck, 0) + count = get(ck, 0) if count < cap: slot_count[ck] = count + 1 - final_key[i] = ck - final_index[i] = count + 1 - reassigned += 1 - placed = True + placed_rows.append(rows[k]) + placed_key.append(ck) + placed_idx.append(count + 1) break - if not placed: - unassigned.append(int(i)) - self._resolve_unassigned(unassigned, sid_key, cap, slot_count, final_index) - return reassigned, unassigned + else: + unassigned.append(rows[k]) + if placed_rows: + final_key[placed_rows] = placed_key + final_index[placed_rows] = placed_idx + return len(placed_rows), unassigned def _resolve_unassigned( self, unassigned: List[int], sid_key: np.ndarray, - cap: int, slot_count: Dict[int, int], final_index: np.ndarray, - ) -> None: - """Apply --unassigned_policy to items that found no free slot.""" + n: int, + ) -> Optional[np.ndarray]: + """Apply --unassigned_policy; return the drop mask (None unless dropping).""" if not unassigned: - return - policy = self.args.unassigned_policy + return None + policy = self._unassigned_policy if policy == "error": preview = ",".join(str(i) for i in unassigned[:10]) raise RuntimeError( f"{len(unassigned)} items could not be assigned within capacity; " f"first unassigned row indices: {preview}" ) - if policy == "keep_original": - for i in unassigned: - key = int(sid_key[i]) - slot_count[key] = slot_count.get(key, 0) + 1 - final_index[i] = slot_count[key] + if policy == "drop": + keep_mask = np.ones(n, dtype=bool) + keep_mask[unassigned] = False + return keep_mask + for i in unassigned: # keep_original + key = int(sid_key[i]) + slot_count[key] = slot_count.get(key, 0) + 1 + final_index[i] = slot_count[key] + return None @staticmethod def _stats( n: int, - counts_by_row: np.ndarray, - first: np.ndarray, - final_key: np.ndarray, - keep_mask: Optional[np.ndarray], + counts: np.ndarray, + slot_count: Dict[int, int], cap: int, reassigned: int, unassigned_count: int, ) -> AssignmentStats: - """Summarize raw vs final bucket occupancy.""" - raw_counts = counts_by_row[first] - kept_key = final_key if keep_mask is None else final_key[keep_mask] - if kept_key.size: - _, fc = np.unique(kept_key, return_counts=True) - final_collision = int((fc > cap).sum()) - max_final = int(fc.max()) + """Summarize raw vs final bucket occupancy. + + Final occupancy is read straight from ``slot_count`` (maintained during + placement) instead of re-grouping the final keys. + """ + if slot_count: + vals = slot_count.values() + final_collision = sum(1 for v in vals if v > cap) + max_final = max(vals) else: final_collision = 0 max_final = 0 return AssignmentStats( total_items=n, - raw_collision_buckets=int((raw_counts > cap).sum()), + raw_collision_buckets=int((counts > cap).sum()), final_collision_buckets=final_collision, reassigned_count=reassigned, unassigned_count=unassigned_count, max_final_bucket_size=max_final, ) - def _writer_type(self) -> Optional[str]: - # Derive the writer from the input reader (hitrate.py idiom) when unset, - # so the output backend matches the input's. ODPS output still resolves - # to OdpsWriter via create_writer's path detection regardless. - if self.args.writer_type: - return self.args.writer_type - if self._input_reader_cls_name: - return self._input_reader_cls_name.replace("Reader", "Writer") - return None - - def _is_csv_output(self) -> bool: - return self._writer_type() == "CsvWriter" - @staticmethod def _item_id_array(values: np.ndarray) -> pa.Array: + """Encode item ids as int64 when integral, else as strings.""" if np.issubdtype(values.dtype, np.integer): return pa.array(values, type=pa.int64()) - return pa.array([str(v) for v in values], type=pa.string()) + return pa.array(values, type=pa.string()) def _codes_column(self, codes: np.ndarray) -> pa.Array: """Encode an (M, n_layers) matrix as list (or a CSV string).""" m, n_layers = codes.shape - if self._is_csv_output(): + if self._is_csv: cols = [ pc.cast(pa.array(codes[:, j]), pa.string()) for j in range(n_layers) ] @@ -543,9 +581,9 @@ def _write( final_codes = final_codes[keep_mask] index = index[keep_mask] writer = create_writer( - self.args.output_path, - writer_type=self._writer_type(), - quota_name=self.args.odps_data_quota_name, + self._output_path, + writer_type=self._writer_type_str, + quota_name=self._odps_data_quota_name, world_size=1, ) n = len(item_ids) @@ -560,14 +598,15 @@ def _write( } ) writer.close() - if self.args.diagnostics_output_path: + if self._diagnostics_output_path: self._write_diagnostics(stats) def _write_diagnostics(self, stats: AssignmentStats) -> None: + """Write the one-row AssignmentStats table to the diagnostics path.""" writer = create_writer( - self.args.diagnostics_output_path, - writer_type=self._writer_type(), - quota_name=self.args.odps_data_quota_name, + self._diagnostics_output_path, + writer_type=self._writer_type_str, + quota_name=self._odps_data_quota_name, world_size=1, ) writer.write( @@ -598,6 +637,12 @@ def _write_diagnostics(self, stats: AssignmentStats) -> None: parser.add_argument("--batch_size", type=int, default=100000) parser.add_argument("--item_id_field", default="item_id") parser.add_argument("--code_field", default="codes") + parser.add_argument( + "--codebook", + required=True, + help="Comma-separated per-layer codebook sizes, e.g. '8192,8192,8192'. " + "Declares n_layers and the last-layer code space used for reassignment.", + ) parser.add_argument("--candidate_codes_field", default="candidate_codes") parser.add_argument( "--candidate_depth", @@ -606,7 +651,6 @@ def _write_diagnostics(self, stats: AssignmentStats) -> None: help="Cap on candidate last-codes tried per overflow item (default: all).", ) parser.add_argument("--max_items_per_codebook", type=int, required=True) - parser.add_argument("--seed", type=int, default=2026) parser.add_argument( "--unassigned_policy", choices=["error", "drop", "keep_original"], @@ -619,12 +663,6 @@ def _write_diagnostics(self, stats: AssignmentStats) -> None: help="Reassignment strategy: 'candidate' uses model candidate_codes; " "'random' draws random within-band last-layer codes (no candidates).", ) - parser.add_argument( - "--random_last_layer_size", - type=int, - default=None, - help="Code-space size of the last SID layer; required for --strategy random.", - ) parser.add_argument( "--random_num_candidates", type=int, diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 120a0dfc..8653c01f 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -35,10 +35,9 @@ candidate_codes_field="candidate_codes", candidate_depth=None, max_items_per_codebook=2, - seed=2026, unassigned_policy="error", strategy="candidate", - random_last_layer_size=None, + codebook="8,8", random_num_candidates=64, rate_only=False, odps_data_quota_name="pay-as-you-go", @@ -152,7 +151,6 @@ def test_random_reassigns_within_band(self) -> None: out, max_items_per_codebook=2, strategy="random", - random_last_layer_size=8, unassigned_policy="keep_original", ) self.assertEqual(stats.reassigned_count, 3) @@ -160,12 +158,14 @@ def test_random_reassigns_within_band(self) -> None: self.assertTrue(all(cb[0] == 0 for cb in d["codebook"])) self.assertLessEqual(max(Counter(map(tuple, d["codebook"])).values()), 2) - def test_random_requires_last_layer_size(self) -> None: + def test_random_requires_last_codebook_ge_2(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") _parquet(inp, list(range(5)), [[0, 0]] * 5) with self.assertRaises(ValueError): - self._run(inp, out, max_items_per_codebook=2, strategy="random") + self._run( + inp, out, max_items_per_codebook=2, strategy="random", codebook="8,1" + ) def test_single_layer_sid(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") @@ -176,7 +176,7 @@ def test_single_layer_sid(self) -> None: out, max_items_per_codebook=1, strategy="random", - random_last_layer_size=16, + codebook="16", unassigned_policy="keep_original", ) self.assertEqual(stats.reassigned_count, 3) @@ -206,7 +206,6 @@ def decisions(order): out, max_items_per_codebook=2, unassigned_policy="keep_original", - seed=7, ) d = self._read_parquet(out) return { From 5c67d6d6e0d6bd0e17748631ee19edc55dfc7ac1 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Mon, 13 Jul 2026 09:45:03 +0000 Subject: [PATCH 15/29] [refactor] add SID collision grouping outputs --- tzrec/tools/sid/collision_prevention.py | 1193 ++++++++++-------- tzrec/tools/sid/collision_prevention_test.py | 473 ++++++- tzrec/tools/sid/collision_resolution.py | 659 ++++++++++ tzrec/tools/sid/collision_resolution_test.py | 300 +++++ 4 files changed, 2075 insertions(+), 550 deletions(-) create mode 100644 tzrec/tools/sid/collision_resolution.py create mode 100644 tzrec/tools/sid/collision_resolution_test.py diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 0b655e8d..de0fb324 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -9,27 +9,26 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Offline SID codebook collision-prevention tool (vectorized). - -Caps every SID bucket at ``--max_items_per_codebook`` and reassigns overflow -items to a free code in the SAME band -- keeping every layer but the last and -varying only the last SID layer, so an item never leaves its ``(prefix)`` band: - -- ``--strategy candidate`` walks the item's model-provided candidate SIDs - (``candidate_codes``), taking the last code of each, best-first; -- ``--strategy random`` draws random last-layer codes (excluding the origin), - a baseline that ignores semantic proximity. - -Items with no free slot fall back per ``--unassigned_policy``. The hot path is -numpy/Arrow-vectorized so a hundred-million-row map fits in one pass; results -are deterministic and independent of input row order. I/O goes -through ``create_reader`` / ``create_writer`` so CSV, Parquet, and ODPS all work; -outputs are written in chunks to stay under Arrow's 2^31 list-offset limit. +"""Offline SID collision prevention with TorchEasyRec-native I/O. + +The runner caps each SID bucket, loads fixed-width last-layer candidates only +for overflow items, delegates collision resolution to the pure NumPy core, and +writes item-level and grouped SID results through TorchEasyRec readers and +writers. CSV uses compact SID strings and JSON item-ID arrays because Arrow's +CSV writer cannot serialize list columns. + +The random strategy intentionally preserves the legacy deterministic baseline: +it draws with replacement from the full last-layer space. Placement skips an +item's origin, so an origin draw or a duplicate draw is not replaced. """ +from __future__ import annotations + import argparse -from dataclasses import asdict, dataclass -from typing import Any, Dict, Iterable, List, Optional, Tuple +import json +import os +from dataclasses import dataclass +from typing import Dict, Iterable, List, Optional, Tuple, Union import numpy as np import pyarrow as pa @@ -38,589 +37,699 @@ # Register the local reader/writer classes used through create_reader/writer. import tzrec.datasets.csv_dataset # noqa: F401 import tzrec.datasets.parquet_dataset # noqa: F401 -from tzrec.datasets.dataset import BaseReader, create_reader, create_writer +from tzrec.datasets.dataset import ( + BaseReader, + BaseWriter, + create_reader, + create_writer, +) +from tzrec.datasets.odps_dataset import _parse_table_path +from tzrec.tools.sid.collision_resolution import ( + CodebookItemGrouping, + CollisionPlan, + CollisionResolutionConfig, + CollisionResolutionResult, + CollisionResolutionStats, + build_original_item_grouping, + build_resolved_item_grouping, + generate_random_candidate_last_codes, + prepare_collision_plan, + resolve_sid_collisions, +) from tzrec.utils.logging_util import logger -_MASK64 = (1 << 64) - 1 - -# Fixed hash seed for the order-independent tie-break. Not a CLI knob: the -# collision rate is seed-invariant, so runs stay reproducible without one. -_SEED = 2026 - -# CSV-fallback delimiters (Parquet/ODPS carry codes as native list, so -# these are unused there). Both are non-comma so a codebook cell never collides -# with CSV's own field separator, and distinct so a compact cell parses cleanly. -_CODE_SEP = "|" # separates the int codes within one SID, e.g. "1|2|3" -_CAND_SEP = ";" # separates candidate SIDs in a compact cell, e.g. "1|2;3|4" - -# Chunk rows per output write: keeps each Arrow array's int32 offsets under 2^31 -# (a single 255M-row string/list column would overflow one array). +_CODE_SEP = "|" +_CAND_SEP = ";" _WRITE_CHUNK = 20_000_000 +_CSV_GROUP_WRITE_CHUNK = 1_000_000 +_ARROW_LIST_OFFSET_MAX = int(np.iinfo(np.int32).max) + + +def _output_path_identity( + path: str, +) -> Union[str, Tuple[str, str, Optional[str], str, Optional[str]]]: + """Return the destination identity used by the repository writer.""" + if not path.startswith("odps://"): + return os.path.realpath(path) + + project, table_name, partitions, schema = _parse_table_path(path) + partition_spec = partitions[0] if partitions and partitions[0] else None + return "odps", project, schema, table_name, partition_spec + + +@dataclass(frozen=True) +class CollisionPreventionConfig: + """Validated configuration for collision-prevention orchestration.""" + + input_path: str + output_path: str + original_sid_groups_output_path: Optional[str] + resolved_sid_groups_output_path: Optional[str] + diagnostics_output_path: Optional[str] + reader_type: Optional[str] + writer_type: Optional[str] + batch_size: int + item_id_field: str + code_field: str + candidate_codes_field: str + layer_sizes: Tuple[int, ...] + max_items_per_codebook: int + unassigned_policy: str + strategy: str + random_num_candidates: int + rate_only: bool + odps_data_quota_name: str + + def __post_init__(self) -> None: + if self.batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {self.batch_size}.") + if self.random_num_candidates < 1: + raise ValueError( + f"random_num_candidates must be >= 1, got {self.random_num_candidates}." + ) + if self.strategy not in {"candidate", "random"}: + raise ValueError(f"unsupported strategy: {self.strategy!r}.") + has_original_groups = bool(self.original_sid_groups_output_path) + has_resolved_groups = bool(self.resolved_sid_groups_output_path) + if has_original_groups != has_resolved_groups: + raise ValueError( + "original_sid_groups_output_path and " + "resolved_sid_groups_output_path must be supplied together." + ) + if not self.rate_only and not has_original_groups: + raise ValueError( + "both SID group output paths are required unless rate_only is set." + ) + named_paths = { + "output_path": self.output_path, + "original_sid_groups_output_path": self.original_sid_groups_output_path, + "resolved_sid_groups_output_path": self.resolved_sid_groups_output_path, + "diagnostics_output_path": self.diagnostics_output_path, + } + seen_paths: Dict[ + Union[str, Tuple[str, str, Optional[str], str, Optional[str]]], str + ] = {} + for name, path in named_paths.items(): + if not path: + continue + path_identity = _output_path_identity(path) + previous_name = seen_paths.get(path_identity) + if previous_name is not None: + raise ValueError(f"{name} must differ from {previous_name}: {path!r}.") + seen_paths[path_identity] = name + _ = self.resolution_config + + @classmethod + def from_namespace(cls, args: argparse.Namespace) -> "CollisionPreventionConfig": + """Build a validated configuration from parsed CLI arguments.""" + layer_sizes = tuple( + int(value) for value in args.codebook.split(",") if value.strip() + ) + return cls( + input_path=args.input_path, + output_path=args.output_path, + original_sid_groups_output_path=getattr( + args, "original_sid_groups_output_path", None + ), + resolved_sid_groups_output_path=getattr( + args, "resolved_sid_groups_output_path", None + ), + diagnostics_output_path=args.diagnostics_output_path, + reader_type=args.reader_type, + writer_type=args.writer_type, + batch_size=args.batch_size, + item_id_field=args.item_id_field, + code_field=args.code_field, + candidate_codes_field=args.candidate_codes_field, + layer_sizes=layer_sizes, + max_items_per_codebook=args.max_items_per_codebook, + unassigned_policy=args.unassigned_policy, + strategy=args.strategy, + random_num_candidates=args.random_num_candidates, + rate_only=args.rate_only, + odps_data_quota_name=args.odps_data_quota_name, + ) -@dataclass -class AssignmentStats: - """Summary statistics for a collision-prevention run.""" - - total_items: int - raw_collision_buckets: int - final_collision_buckets: int - reassigned_count: int - unassigned_count: int - max_final_bucket_size: int - - -def _splitmix64(x: np.ndarray) -> np.ndarray: - """Vectorized order-independent SplitMix64 hash of a uint64 array. - - A pure function of the input values (not their position), mixed with the - module ``_SEED``, so it gives a stable tie-break invariant to input row order - -- unlike a read-order index -- while staying fully vectorized. - - Args: - x (np.ndarray): uint64 values to hash. - - Returns: - np.ndarray: uint64 hashes, same shape as ``x``. - """ - with np.errstate(over="ignore"): - z = x.astype(np.uint64) + np.uint64((_SEED * 0x9E3779B97F4A7C15) & _MASK64) - z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) - z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) - return z ^ (z >> np.uint64(31)) - - -def _order_hash(item_ids: np.ndarray) -> np.ndarray: - """Per-item order-independent tie-break hash (uint64), vectorized. - - Integer ids hash directly; string/object ids are folded to uint64 via - ``pandas.util.hash_array`` first. Both then pass through :func:`_splitmix64`. - - Args: - item_ids (np.ndarray): item id values. + @property + def resolution_config(self) -> CollisionResolutionConfig: + """Return the pure-core configuration.""" + return CollisionResolutionConfig( + layer_sizes=self.layer_sizes, + capacity=self.max_items_per_codebook, + fallback_policy=self.unassigned_policy, + ) - Returns: - np.ndarray: uint64 per-item hashes. - """ - if np.issubdtype(item_ids.dtype, np.integer): - base = item_ids.astype(np.uint64) - else: - import pandas as pd - base = pd.util.hash_array(np.asarray(item_ids, dtype=object)) - return _splitmix64(base) +class CollisionRunner: + """Run SID collision prevention over TorchEasyRec readers and writers.""" + def __init__( + self, + config: Union[CollisionPreventionConfig, argparse.Namespace], + ) -> None: + self._config = ( + config + if isinstance(config, CollisionPreventionConfig) + else CollisionPreventionConfig.from_namespace(config) + ) + self._default_writer_type: Optional[str] = None + self._item_id_type: Optional[pa.DataType] = None -class CollisionRunner: - """Vectorized SID collision-prevention runner over the dataset reader/writer. - - The backend (CSV / Parquet / ODPS) is chosen by the reader/writer type, so - the same path serves local files and MaxCompute tables. - """ - - def __init__(self, args: argparse.Namespace) -> None: - self._input_path = args.input_path - self._output_path = args.output_path - self._diagnostics_output_path = args.diagnostics_output_path - self._reader_type = args.reader_type - self._writer_type = args.writer_type - self._batch_size = args.batch_size - self._item_id_field = args.item_id_field - self._code_field = args.code_field - self._candidate_codes_field = args.candidate_codes_field - self._candidate_depth = args.candidate_depth - self._max_items_per_codebook = args.max_items_per_codebook - if self._max_items_per_codebook < 1: - raise ValueError( - "--max_items_per_codebook must be >= 1, got " - f"{self._max_items_per_codebook}." - ) - self._unassigned_policy = args.unassigned_policy - self._strategy = args.strategy - self._codebook = [int(x) for x in args.codebook.split(",") if x.strip()] - if not self._codebook: - raise ValueError("--codebook must list at least one per-layer size.") - self._random_num_candidates = args.random_num_candidates - self._rate_only = args.rate_only - self._odps_data_quota_name = args.odps_data_quota_name - # The writer type (and its is-CSV flag) can depend on the input reader - # class when --writer_type is unset, so they are finalized once at the - # first read rather than here; declared now so the attributes always exist. - self._writer_type_str: Optional[str] = None - self._is_csv: bool = False - - def run(self) -> AssignmentStats: - """Read the SID map, assign, and write the reassigned map + stats.""" + def run(self) -> CollisionResolutionStats: + """Read, resolve collisions, and write the resulting SID map.""" item_ids, codes = self._load_codes() - final_codes, index, keep_mask, stats = self._assign(item_ids, codes) - if not self._rate_only: - self._write(item_ids, codes, final_codes, index, keep_mask, stats) + plan = prepare_collision_plan( + item_ids, + codes, + self._config.resolution_config, + ) + candidate_last_codes = self._candidate_last_codes(plan) + result = resolve_sid_collisions( + plan, + candidate_last_codes, + collect_grouping=( + not self._config.rate_only and plan.overflow_rows.size > 0 + ), + ) + del candidate_last_codes + + if self._config.rate_only: + logger.info("rate_only: skipping map and SID group writes") + del plan else: - logger.info("rate_only: skipping map write") - logger.info("SID collision prevention finished: %s", stats) - return stats + self._write_group_outputs(item_ids, codes, plan, result) + del plan + self._write_map(item_ids, codes, result) + if self._config.diagnostics_output_path: + self._write_diagnostics(result.stats) + + logger.info("SID collision prevention finished: %s", result.stats) + return result.stats def _make_reader(self, selected_cols: List[str]) -> BaseReader: - """Open the input reader projecting ``selected_cols``.""" + """Open a repository reader projecting ``selected_cols``.""" return create_reader( - input_path=self._input_path, - batch_size=self._batch_size, + input_path=self._config.input_path, + batch_size=self._config.batch_size, selected_cols=selected_cols, - reader_type=self._reader_type, - quota_name=self._odps_data_quota_name, + reader_type=self._config.reader_type, + quota_name=self._config.odps_data_quota_name, ) - def _read(self) -> Iterable[Dict[str, pa.Array]]: - """Stream item-id + code batches, caching the resolved output backend.""" - reader = self._make_reader([self._item_id_field, self._code_field]) - # Writer defaults to matching the input reader (hitrate.py idiom) when - # --writer_type is unset; ODPS still resolves to OdpsWriter via - # create_writer's path detection regardless. - reader_cls_name = reader.__class__.__name__ - self._writer_type_str = self._writer_type or ( - reader_cls_name.replace("Reader", "Writer") + def _read_codes(self) -> Iterable[Dict[str, pa.Array]]: + """Stream item IDs and SIDs while resolving the default writer type.""" + reader = self._make_reader( + [self._config.item_id_field, self._config.code_field] + ) + reader_name = reader.__class__.__name__ + self._default_writer_type = self._config.writer_type or reader_name.replace( + "Reader", "Writer" ) - self._is_csv = self._writer_type_str == "CsvWriter" yield from reader.to_batches() @staticmethod - def _codes_matrix(arr: pa.Array) -> np.ndarray: - """Decode a SID code column into an (N, n_layers) int64 matrix. + def _is_list_type(data_type: pa.DataType) -> bool: + """Return whether ``data_type`` is an Arrow list representation.""" + return ( + pa.types.is_list(data_type) + or pa.types.is_large_list(data_type) + or pa.types.is_fixed_size_list(data_type) + ) - Parquet/ODPS give a ``list`` cell; CSV gives a ``_CODE_SEP`` - string; a single-layer numeric CSV column may arrive already as ints. - """ - if pa.types.is_list(arr.type) or pa.types.is_large_list(arr.type): - n = len(arr) + @classmethod + def _codes_matrix(cls, values: pa.Array) -> np.ndarray: + """Decode an SID column into an ``(N, n_layers)`` int64 matrix.""" + if cls._is_list_type(values.type): + row_count = len(values) flat = ( - arr.flatten() + values.flatten() .to_numpy(zero_copy_only=False) .astype(np.int64, copy=False) ) - elif pa.types.is_integer(arr.type): - arr_np = arr.to_numpy(zero_copy_only=False).astype(np.int64, copy=False) - return arr_np.reshape(-1, 1) + elif pa.types.is_integer(values.type): + array = values.to_numpy(zero_copy_only=False).astype(np.int64, copy=False) + return array.reshape(-1, 1) else: - parts = pc.split_pattern(arr, _CODE_SEP) - n = len(parts) + parts = pc.split_pattern(values, _CODE_SEP) + row_count = len(parts) flat = pc.cast(parts.flatten(), pa.int64()).to_numpy(zero_copy_only=False) - if n == 0: + + if row_count == 0: return np.empty((0, 0), dtype=np.int64) - n_layers = flat.shape[0] // n - if flat.shape[0] != n * n_layers: + layer_count = flat.shape[0] // row_count + if flat.shape[0] != row_count * layer_count: raise ValueError("ragged SID codes: all items must share n_layers.") - return flat.reshape(n, n_layers) + return flat.reshape(row_count, layer_count) + + def _validate_item_ids(self, item_ids: pa.Array) -> None: + """Validate one item ID batch and retain its Arrow type. + + Args: + item_ids: Item IDs from one input batch. + + Raises: + ValueError: If IDs contain nulls or their type changes between batches. + """ + if item_ids.null_count: + raise ValueError("item IDs must not contain null values.") + if self._item_id_type is None: + self._item_id_type = item_ids.type + elif self._item_id_type != item_ids.type: + raise ValueError( + "item ID type changed between batches: " + f"{self._item_id_type} vs {item_ids.type}." + ) def _load_codes(self) -> Tuple[np.ndarray, np.ndarray]: - """Stream the map into an item-id array and an (N, n_layers) code matrix.""" + """Load item IDs and SIDs into arrays used by the NumPy core.""" id_chunks: List[np.ndarray] = [] code_chunks: List[np.ndarray] = [] - for batch in self._read(): - id_chunks.append(batch[self._item_id_field].to_numpy(zero_copy_only=False)) - code_chunks.append(self._codes_matrix(batch[self._code_field])) + for batch in self._read_codes(): + item_ids = batch[self._config.item_id_field] + self._validate_item_ids(item_ids) + id_chunks.append(item_ids.to_numpy(zero_copy_only=False)) + code_chunks.append(self._codes_matrix(batch[self._config.code_field])) + if not id_chunks: raise ValueError("SID input is empty.") - item_ids = np.concatenate(id_chunks) - del id_chunks - codes = np.concatenate(code_chunks, axis=0) - del code_chunks - if codes.shape[1] < 1: + item_id_array = np.concatenate(id_chunks) + code_matrix = np.concatenate(code_chunks, axis=0) + if code_matrix.shape[1] < 1: raise ValueError("SID codes must have at least one layer.") - return item_ids, codes - - def _load_candidate_last(self, overflow_id_arr: np.ndarray) -> Dict[Any, List[int]]: - """Map each overflow item id to its ordered candidate last-layer codes. + return item_id_array, code_matrix - Reads ``candidate_codes`` (list> for Parquet/ODPS, a - ``_CAND_SEP``/``_CODE_SEP`` string for CSV) but keeps only the last code - of each candidate SID and only for overflow items, so memory scales with - the overflow fraction, not the whole table. - """ - field = self._candidate_codes_field - # Sort the overflow ids once; per-batch membership is then a bisect - # instead of re-sorting the whole (potentially huge) overflow set. - ov = np.sort(overflow_id_arr) - cand: Dict[Any, List[int]] = {} - for batch in self._make_reader([self._item_id_field, field]).to_batches(): - if field not in batch: - break - ids = batch[self._item_id_field].to_numpy(zero_copy_only=False) - pos = np.searchsorted(ov, ids) - np.clip(pos, 0, len(ov) - 1, out=pos) - keep = np.where(ov[pos] == ids)[0] - if keep.size == 0: - continue - col = batch[field] - if pa.types.is_list(col.type) or pa.types.is_large_list(col.type): - self._collect_candidate_lists(ids, col, keep, cand) - else: - self._collect_candidate_strings(ids, col, keep, cand) - return cand - - def _collect_candidate_lists( - self, - ids: np.ndarray, - col: pa.Array, - keep: np.ndarray, - cand: Dict[Any, List[int]], - ) -> None: - """Vectorized last-code extraction from a list> batch.""" - n = len(col) - inner = col.flatten() # list, one per candidate SID - if len(inner) == 0: - return - k = len(inner) // n - if len(inner) != n * k: + def _candidate_last_codes(self, plan: CollisionPlan) -> np.ndarray: + """Provide row-aligned candidate last codes for the overflow plan.""" + if plan.overflow_rows.size == 0: + return np.empty((0, 0), dtype=np.int64) + if self._config.strategy == "random": + return generate_random_candidate_last_codes( + plan.overflow_item_ids, + self._config.random_num_candidates, + self._config.layer_sizes[-1], + ) + return self._load_candidate_last_codes(plan.overflow_item_ids) + + def _candidate_matrix(self, values: pa.Array) -> np.ndarray: + """Decode candidate SIDs and retain only their last-layer codes.""" + if self._is_list_type(values.type): + return self._candidate_list_matrix(values) + return self._candidate_string_matrix(values) + + def _candidate_list_matrix(self, values: pa.Array) -> np.ndarray: + """Decode a nested Arrow candidate column into an ``(N, K)`` matrix.""" + row_count = len(values) + if row_count == 0: + return np.empty((0, 0), dtype=np.int64) + inner = values.flatten() + candidate_count = len(inner) // row_count + if len(inner) != row_count * candidate_count: raise ValueError("ragged candidate_codes: all items must share topk.") + if candidate_count == 0: + return np.empty((row_count, 0), dtype=np.int64) flat = ( inner.flatten().to_numpy(zero_copy_only=False).astype(np.int64, copy=False) ) - n_layers = flat.shape[0] // (n * k) - last = flat.reshape(n, k, n_layers)[:, :, n_layers - 1] - if self._candidate_depth is not None: - last = last[:, : self._candidate_depth] - last_lists = last.tolist() # per-item Python lists: cheap to scan in the loop - for iid, i in zip(ids[keep].tolist(), keep.tolist()): - cand[iid] = last_lists[i] - - def _collect_candidate_strings( - self, - ids: np.ndarray, - col: pa.Array, - keep: np.ndarray, - cand: Dict[Any, List[int]], - ) -> None: - """Per-row last-code extraction from a CSV compact-candidate batch.""" - depth = self._candidate_depth - values = col.to_pylist() - for iid, i in zip(ids[keep].tolist(), keep.tolist()): - text = values[i] - if not text: - continue - last_codes = [] - for part in str(text).split(_CAND_SEP): - part = part.strip() - if not part: - continue - codes = [int(p) for p in part.split(_CODE_SEP) if p.strip()] - if codes: - last_codes.append(codes[-1]) - cand[iid] = last_codes[:depth] if depth is not None else last_codes - - def _assign( - self, item_ids: np.ndarray, codes: np.ndarray - ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], AssignmentStats]: - """Cap buckets and reassign overflow within-band; return the final map.""" - cap = self._max_items_per_codebook - n, n_layers = codes.shape - if n_layers != len(self._codebook): + tuple_count = row_count * candidate_count + layer_count = flat.shape[0] // tuple_count + if layer_count < 1 or flat.shape[0] != tuple_count * layer_count: raise ValueError( - f"codes have {n_layers} layers but --codebook has " - f"{len(self._codebook)}." + "ragged candidate_codes: every candidate must share n_layers." ) - last = codes[:, n_layers - 1].astype(np.int64, copy=False) - - band_id = self._band_ids(codes) - order = _order_hash(item_ids) - - # Rank rows within their (band_id, last) bucket by order-hash: the RATE is - # invariant to which cap items are kept, but the choice is deterministic. - # One lexsort yields the rank, each bucket's representative row, and size. - rank, rep_rows, counts = self._within_bucket_rank(band_id, last, order) - overflow_order = np.flatnonzero(rank >= cap) - overflow_order = overflow_order[ - np.lexsort( - (order[overflow_order], last[overflow_order], band_id[overflow_order]) - ) - ] - - cand = self._candidates(item_ids, overflow_order) - last_size = self._codebook[-1] - - sid_key = band_id * last_size + last - slot_count: Dict[int, int] = dict( - zip(sid_key[rep_rows].tolist(), np.minimum(counts, cap).tolist()) - ) - - final_key = sid_key.copy() - final_index = rank + 1 - reassigned, unassigned = self._place_overflow( - overflow_order, - item_ids, - band_id, - last, - cand, - last_size, - cap, - slot_count, - final_key, - final_index, - ) - keep_mask = self._resolve_unassigned( - unassigned, sid_key, slot_count, final_index, n - ) - - final_codes = codes.copy() - final_codes[:, n_layers - 1] = final_key % last_size - - stats = self._stats(n, counts, slot_count, cap, reassigned, len(unassigned)) - return final_codes, final_index, keep_mask, stats - - def _band_ids(self, codes: np.ndarray) -> np.ndarray: - """Dense integer id per distinct ``(prefix)`` band (all layers but last). - - Packs the prefix layers into one int64 key using the declared per-layer - codebook sizes as radices, then factorizes -- far faster than - ``np.unique`` over 2-D void rows at scale, and independent of the values a - given batch happens to contain. - """ - n, n_layers = codes.shape - if n_layers == 1: - return np.zeros(n, dtype=np.int64) - key = codes[:, 0].astype(np.int64) # fresh copy: safe to pack in place below - for j in range(1, n_layers - 1): - key *= self._codebook[j] - key += codes[:, j] - inverse = np.unique(key, return_inverse=True)[1] - return inverse.astype(np.int64, copy=False).reshape(-1) + return flat.reshape(row_count, candidate_count, layer_count)[:, :, -1] @staticmethod - def _within_bucket_rank( - band_id: np.ndarray, last: np.ndarray, order: np.ndarray - ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Rank rows within their ``(band_id, last)`` bucket by order-hash. - - Returns the per-row 0-based rank, each bucket's representative original - row index, and each bucket's size -- all from a single column lexsort, so - the grouping is computed once (no packed key, no per-row size scatter). - """ - n = band_id.shape[0] - sort_idx = np.lexsort((order, last, band_id)) - sb = band_id[sort_idx] - sl = last[sort_idx] - is_first = np.empty(n, dtype=bool) - is_first[0] = True - is_first[1:] = (sb[1:] != sb[:-1]) | (sl[1:] != sl[:-1]) - first_sorted = np.flatnonzero(is_first) - counts = np.diff(np.append(first_sorted, n)) - rank = np.empty(n, dtype=np.int64) - rank[sort_idx] = np.arange(n) - np.repeat(first_sorted, counts) - return rank, sort_idx[first_sorted], counts - - def _candidates( - self, - item_ids: np.ndarray, - overflow_order: np.ndarray, - ) -> Dict[Any, List[int]]: - """Build per-overflow-item last-code candidate lists.""" - if self._strategy == "random": - size = self._codebook[-1] - if size < 2: - raise ValueError("strategy='random' requires a last codebook >= 2.") - num = min(self._random_num_candidates, size - 1) - draws = self._random_draws(overflow_order, item_ids, num, size) - ov_ids = item_ids[overflow_order].tolist() # Python-native dict keys - return {ov_ids[j]: draws[j].tolist() for j in range(len(ov_ids))} - if overflow_order.size == 0: - return {} - cand = self._load_candidate_last(item_ids[overflow_order]) - if not cand: + def _candidate_string_matrix(values: pa.Array) -> np.ndarray: + """Decode compact CSV candidate strings into an ``(N, K)`` matrix.""" + rows: List[List[int]] = [] + for value in values.to_pylist(): + last_codes: List[int] = [] + if value: + for candidate in str(value).split(_CAND_SEP): + parts = [ + int(part) + for part in candidate.strip().split(_CODE_SEP) + if part.strip() + ] + if parts: + last_codes.append(parts[-1]) + rows.append(last_codes) + + widths = {len(row) for row in rows} + if len(widths) > 1: + raise ValueError("ragged candidate_codes: all items must share topk.") + width = widths.pop() if widths else 0 + if width == 0: + return np.empty((len(rows), 0), dtype=np.int64) + return np.asarray(rows, dtype=np.int64).reshape(len(rows), width) + + def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarray: + """Load fixed-width candidates aligned to ``overflow_item_ids``.""" + item_count = overflow_item_ids.shape[0] + sorted_to_overflow = np.argsort(overflow_item_ids, kind="stable") + sorted_ids = overflow_item_ids[sorted_to_overflow] + if item_count > 1 and np.any(sorted_ids[1:] == sorted_ids[:-1]): + raise ValueError("overflow item IDs must be unique.") + + candidates: Optional[np.ndarray] = None + seen = np.zeros(item_count, dtype=bool) + field = self._config.candidate_codes_field + reader = self._make_reader([self._config.item_id_field, field]) + for batch in reader.to_batches(): + if field not in batch: + break + batch_ids = batch[self._config.item_id_field].to_numpy(zero_copy_only=False) + positions = np.searchsorted(sorted_ids, batch_ids) + in_bounds = positions < item_count + source_rows = np.flatnonzero(in_bounds) + if source_rows.size: + source_rows = source_rows[ + sorted_ids[positions[source_rows]] == batch_ids[source_rows] + ] + if source_rows.size == 0: + continue + + target_rows = sorted_to_overflow[positions[source_rows]] + 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_matrix(selected) + if candidates is None: + candidates = np.empty( + (item_count, batch_candidates.shape[1]), dtype=np.int64 + ) + elif candidates.shape[1] != batch_candidates.shape[1]: + raise ValueError( + "candidate topk changed between batches: " + f"{candidates.shape[1]} vs {batch_candidates.shape[1]}." + ) + candidates[target_rows] = batch_candidates + seen[target_rows] = True + + if candidates is None: raise ValueError( "map has overflow items but candidate_codes yielded no candidates." ) - return cand - - def _random_draws( - self, overflow_order: np.ndarray, item_ids: np.ndarray, num: int, size: int - ) -> np.ndarray: - """Order-independent random last-layer draws per overflow item, (M, num).""" - h = _order_hash(item_ids[overflow_order]) - k = np.arange(num, dtype=np.uint64) - with np.errstate(over="ignore"): - mixed = _splitmix64(h[:, None] + k[None, :] * np.uint64(0x9E3779B97F4A7C15)) - return (mixed % np.uint64(size)).astype(np.int64) - - def _place_overflow( - self, - overflow_order: np.ndarray, - item_ids: np.ndarray, - band_id: np.ndarray, - last: np.ndarray, - cand: Dict[Any, List[int]], - last_size: int, - cap: int, - slot_count: Dict[int, int], - final_key: np.ndarray, - final_index: np.ndarray, - ) -> Tuple[int, List[int]]: - """Greedy single pass: place each overflow item in the first free slot. - - Overflow-row inputs are pre-materialized to Python scalars so the loop - avoids per-iteration 0-d numpy indexing; placements are scattered back - into ``final_key`` / ``final_index`` vectorized at the end. - """ - rows = overflow_order.tolist() - bases = (band_id[overflow_order] * last_size).tolist() - origin_lasts = last[overflow_order].tolist() - ids = item_ids[overflow_order].tolist() - placed_rows: List[int] = [] - placed_key: List[int] = [] - placed_idx: List[int] = [] - unassigned: List[int] = [] - get = slot_count.get - for k, base in enumerate(bases): - origin_last = origin_lasts[k] - for code in cand.get(ids[k], ()): # ordered best-first - if code == origin_last: - continue - ck = base + code - count = get(ck, 0) - if count < cap: - slot_count[ck] = count + 1 - placed_rows.append(rows[k]) - placed_key.append(ck) - placed_idx.append(count + 1) - break - else: - unassigned.append(rows[k]) - if placed_rows: - final_key[placed_rows] = placed_key - final_index[placed_rows] = placed_idx - return len(placed_rows), unassigned - - def _resolve_unassigned( - self, - unassigned: List[int], - sid_key: np.ndarray, - slot_count: Dict[int, int], - final_index: np.ndarray, - n: int, - ) -> Optional[np.ndarray]: - """Apply --unassigned_policy; return the drop mask (None unless dropping).""" - if not unassigned: - return None - policy = self._unassigned_policy - if policy == "error": - preview = ",".join(str(i) for i in unassigned[:10]) - raise RuntimeError( - f"{len(unassigned)} items could not be assigned within capacity; " - f"first unassigned row indices: {preview}" + if not np.all(seen): + missing = np.flatnonzero(~seen) + preview = ",".join(str(value) for value in missing[:10]) + raise ValueError( + f"candidate_codes missing for {missing.size} overflow items; " + f"first overflow positions: {preview}." ) - if policy == "drop": - keep_mask = np.ones(n, dtype=bool) - keep_mask[unassigned] = False - return keep_mask - for i in unassigned: # keep_original - key = int(sid_key[i]) - slot_count[key] = slot_count.get(key, 0) + 1 - final_index[i] = slot_count[key] - return None + return candidates + + def _make_writer(self, output_path: str) -> BaseWriter: + """Create a repository writer for one independently resolved output.""" + if self._default_writer_type is None: + raise RuntimeError("writer type is unavailable before reading input.") + return create_writer( + output_path, + writer_type=self._default_writer_type, + quota_name=self._config.odps_data_quota_name, + world_size=1, + ) + + def _item_id_array(self, values: np.ndarray) -> pa.Array: + """Encode item IDs with the input Arrow type preserved.""" + if self._item_id_type is None: + raise RuntimeError("item ID type is unavailable before reading input.") + return pa.array(values, type=self._item_id_type) @staticmethod - def _stats( - n: int, - counts: np.ndarray, - slot_count: Dict[int, int], - cap: int, - reassigned: int, - unassigned_count: int, - ) -> AssignmentStats: - """Summarize raw vs final bucket occupancy. - - Final occupancy is read straight from ``slot_count`` (maintained during - placement) instead of re-grouping the final keys. - """ - if slot_count: - vals = slot_count.values() - final_collision = sum(1 for v in vals if v > cap) - max_final = max(vals) + def _grouped_item_ids_column( + values: pa.Array, offsets: np.ndarray, is_csv: bool + ) -> pa.Array: + """Encode flat item IDs and offsets as one grouped output column.""" + arrow_offsets = pa.array(offsets, type=pa.int32()) + if not is_csv: + return pa.ListArray.from_arrays(arrow_offsets, values) + + if pa.types.is_integer(values.type): + encoded_values = pc.cast(values, pa.string()) else: - final_collision = 0 - max_final = 0 - return AssignmentStats( - total_items=n, - raw_collision_buckets=int((counts > cap).sum()), - final_collision_buckets=final_collision, - reassigned_count=reassigned, - unassigned_count=unassigned_count, - max_final_bucket_size=max_final, - ) + try: + encoded_values = pa.array( + [ + json.dumps( + value, + ensure_ascii=False, + separators=(",", ":"), + ) + for value in values.to_pylist() + ], + type=pa.string(), + ) + except (TypeError, ValueError) as error: + raise ValueError( + "CSV SID group output cannot JSON-encode item ID type " + f"{values.type}." + ) from error + encoded_lists = pa.ListArray.from_arrays(arrow_offsets, encoded_values) + joined_values = pc.binary_join(encoded_lists, ",") + return pc.binary_join_element_wise("[", joined_values, "]", "") @staticmethod - def _item_id_array(values: np.ndarray) -> pa.Array: - """Encode item ids as int64 when integral, else as strings.""" - if np.issubdtype(values.dtype, np.integer): - return pa.array(values, type=pa.int64()) - return pa.array(values, type=pa.string()) - - def _codes_column(self, codes: np.ndarray) -> pa.Array: - """Encode an (M, n_layers) matrix as list (or a CSV string).""" - m, n_layers = codes.shape - if self._is_csv: - cols = [ - pc.cast(pa.array(codes[:, j]), pa.string()) for j in range(n_layers) + def _codes_column(codes: np.ndarray, is_csv: bool) -> pa.Array: + """Encode an SID matrix for the actual output writer.""" + row_count, layer_count = codes.shape + if layer_count < 1: + raise ValueError("SID output must contain at least one layer.") + if is_csv: + columns = [ + pc.cast(pa.array(codes[:, layer]), pa.string()) + for layer in range(layer_count) ] - return pc.binary_join_element_wise(*cols, _CODE_SEP) + return pc.binary_join_element_wise(*columns, _CODE_SEP) + if row_count > _ARROW_LIST_OFFSET_MAX // layer_count: + raise ValueError("SID output exceeds Arrow list offset capacity.") values = pa.array(codes.reshape(-1)) - offsets = pa.array(np.arange(0, (m + 1) * n_layers, n_layers, dtype=np.int32)) + offsets = pa.array( + np.arange( + 0, + (row_count + 1) * layer_count, + layer_count, + dtype=np.int32, + ) + ) return pa.ListArray.from_arrays(offsets, values) - def _write( + def _write_map( self, item_ids: np.ndarray, origin_codes: np.ndarray, - final_codes: np.ndarray, - index: np.ndarray, - keep_mask: Optional[np.ndarray], - stats: AssignmentStats, + result: CollisionResolutionResult, ) -> None: - """Chunked, vectorized write of the reassigned map (and diagnostics).""" - if keep_mask is not None: - item_ids = item_ids[keep_mask] - origin_codes = origin_codes[keep_mask] - final_codes = final_codes[keep_mask] - index = index[keep_mask] - writer = create_writer( - self._output_path, - writer_type=self._writer_type_str, - quota_name=self._odps_data_quota_name, - world_size=1, + """Write the resolved map in chunks without a full final-code copy.""" + writer = self._make_writer(self._config.output_path) + is_csv = writer.__class__.__name__ == "CsvWriter" + retained_rows = ( + np.flatnonzero(result.retained_mask) + if result.retained_mask is not None + else None + ) + output_count = ( + retained_rows.shape[0] if retained_rows is not None else item_ids.shape[0] + ) + try: + write_chunk = _WRITE_CHUNK + if not is_csv: + write_chunk = min( + write_chunk, + _ARROW_LIST_OFFSET_MAX // origin_codes.shape[1], + ) + if write_chunk < 1: + raise ValueError("one SID row exceeds Arrow list offset capacity.") + for start in range(0, output_count, write_chunk): + end = min(start + write_chunk, output_count) + selection: Union[slice, np.ndarray] + selection = ( + retained_rows[start:end] + if retained_rows is not None + else slice(start, end) + ) + origin_chunk = origin_codes[selection] + final_chunk = origin_chunk.copy() + final_chunk[:, -1] = result.resolved_last_codes[selection] + writer.write( + { + "item_id": self._item_id_array(item_ids[selection]), + "origin_codebook": self._codes_column(origin_chunk, is_csv), + "codebook": self._codes_column(final_chunk, is_csv), + "index": pa.array( + result.slot_indices[selection], type=pa.int64() + ), + } + ) + finally: + writer.close() + + def _write_group_outputs( + self, + item_ids: np.ndarray, + origin_codes: np.ndarray, + plan: CollisionPlan, + result: CollisionResolutionResult, + ) -> None: + """Build and write original and resolved SID item groups sequentially. + + Args: + item_ids: Input item IDs in original row order. + origin_codes: Original full SID matrix. + plan: Original SID aggregation and overflow plan. + result: Completed collision-resolution result. + """ + original_path = self._config.original_sid_groups_output_path + resolved_path = self._config.resolved_sid_groups_output_path + if original_path is None or resolved_path is None: + raise RuntimeError("SID group output paths were not validated.") + + original_grouping = build_original_item_grouping(plan) + self._write_sid_groups( + original_path, + item_ids, + origin_codes, + original_grouping, + resolved_last_codes=None, ) - n = len(item_ids) - for s in range(0, n, _WRITE_CHUNK): - e = min(s + _WRITE_CHUNK, n) + if plan.overflow_rows.size == 0: + self._write_sid_groups( + resolved_path, + item_ids, + origin_codes, + original_grouping, + resolved_last_codes=result.resolved_last_codes, + ) + del original_grouping + else: + del original_grouping + resolved_grouping = build_resolved_item_grouping(plan, result) + self._write_sid_groups( + resolved_path, + item_ids, + origin_codes, + resolved_grouping, + resolved_last_codes=result.resolved_last_codes, + ) + del resolved_grouping + + def _write_sid_groups( + self, + output_path: str, + item_ids: np.ndarray, + origin_codes: np.ndarray, + grouping: CodebookItemGrouping, + resolved_last_codes: Optional[np.ndarray], + ) -> None: + """Write codebook-to-item-ID groups in SID and slot order. + + Args: + output_path: Destination passed to the selected repository writer. + item_ids: Input item IDs in original row order. + origin_codes: Original full SID matrix. + grouping: Sorted SID groups and their flattened original row order. + resolved_last_codes: Final last-layer values, or ``None`` when + writing original SIDs. + + Raises: + RuntimeError: If grouping dimensions or counts are inconsistent. + ValueError: If one group exceeds Arrow list offset capacity. + """ + group_count = grouping.counts.shape[0] + if grouping.sid_keys.shape[0] != group_count: + raise RuntimeError("SID group keys and counts must have the same length.") + if np.any(grouping.counts <= 0): + raise RuntimeError("SID group counts must be positive.") + if np.any(grouping.counts > _ARROW_LIST_OFFSET_MAX): + raise ValueError("one SID group exceeds Arrow list offset capacity.") + + offsets = grouping.offsets + if int(offsets[-1]) != grouping.row_order.shape[0]: + raise RuntimeError("SID group counts do not match grouped row count.") + writer = self._make_writer(output_path) + is_csv = writer.__class__.__name__ == "CsvWriter" + child_chunk = ( + min(_WRITE_CHUNK, _CSV_GROUP_WRITE_CHUNK) if is_csv else _WRITE_CHUNK + ) + try: + max_codebook_rows = ( + group_count + if is_csv + else _ARROW_LIST_OFFSET_MAX // origin_codes.shape[1] + ) + if not is_csv and max_codebook_rows < 1: + raise ValueError("one SID row exceeds Arrow list offset capacity.") + group_start = 0 + while group_start < group_count: + child_limit = int(offsets[group_start]) + child_chunk + group_end = int(np.searchsorted(offsets, child_limit, side="right") - 1) + group_end = max(group_end, group_start + 1) + group_end = min(group_end, group_start + max_codebook_rows) + + child_start = int(offsets[group_start]) + child_end = int(offsets[group_end]) + rows = grouping.row_order[child_start:child_end] + local_offsets = ( + offsets[group_start : group_end + 1] - child_start + ).astype(np.int32, copy=False) + representative_rows = grouping.row_order[offsets[group_start:group_end]] + code_chunk = origin_codes[representative_rows] + if resolved_last_codes is not None: + code_chunk[:, -1] = resolved_last_codes[representative_rows] + flat_item_ids = self._item_id_array(item_ids[rows]) + writer.write( + { + "codebook": self._codes_column(code_chunk, is_csv), + "itemids": self._grouped_item_ids_column( + flat_item_ids, local_offsets, is_csv + ), + } + ) + group_start = group_end + finally: + writer.close() + + def _write_diagnostics(self, stats: CollisionResolutionStats) -> None: + """Write legacy-compatible diagnostic column names.""" + path = self._config.diagnostics_output_path + if path is None: + return + writer = self._make_writer(path) + try: writer.write( { - "item_id": self._item_id_array(item_ids[s:e]), - "origin_codebook": self._codes_column(origin_codes[s:e]), - "codebook": self._codes_column(final_codes[s:e]), - "index": pa.array(index[s:e], type=pa.int64()), + name: pa.array([value], type=pa.int64()) + for name, value in stats.to_output_dict().items() } ) - writer.close() - if self._diagnostics_output_path: - self._write_diagnostics(stats) - - def _write_diagnostics(self, stats: AssignmentStats) -> None: - """Write the one-row AssignmentStats table to the diagnostics path.""" - writer = create_writer( - self._diagnostics_output_path, - writer_type=self._writer_type_str, - quota_name=self._odps_data_quota_name, - world_size=1, - ) - writer.write( - {k: pa.array([v], type=pa.int64()) for k, v in asdict(stats).items()} - ) - writer.close() + finally: + writer.close() -if __name__ == "__main__": +def build_parser() -> argparse.ArgumentParser: + """Build the collision-prevention command-line parser.""" parser = argparse.ArgumentParser( description="Prevent SID codebook collisions (vectorized, within-band)." ) parser.add_argument("--input_path", required=True) parser.add_argument("--output_path", required=True) + parser.add_argument( + "--original_sid_groups_output_path", + default=None, + help=( + "Output grouped item IDs keyed by their original SID; required unless " + "--rate_only is set." + ), + ) + parser.add_argument( + "--resolved_sid_groups_output_path", + default=None, + help=( + "Output grouped retained item IDs keyed by their resolved SID; required " + "unless --rate_only is set." + ), + ) parser.add_argument("--diagnostics_output_path", default=None) parser.add_argument( "--reader_type", @@ -631,8 +740,7 @@ def _write_diagnostics(self, stats: AssignmentStats) -> None: "--writer_type", choices=["CsvWriter", "ParquetWriter", "OdpsWriter"], default=None, - help="Output writer; defaults to matching the input reader " - "(CsvReader -> CsvWriter, etc.).", + help="Output writer; defaults to matching the input reader.", ) parser.add_argument("--batch_size", type=int, default=100000) parser.add_argument("--item_id_field", default="item_id") @@ -640,16 +748,9 @@ def _write_diagnostics(self, stats: AssignmentStats) -> None: parser.add_argument( "--codebook", required=True, - help="Comma-separated per-layer codebook sizes, e.g. '8192,8192,8192'. " - "Declares n_layers and the last-layer code space used for reassignment.", + help="Comma-separated per-layer sizes, e.g. '8192,8192,8192'.", ) parser.add_argument("--candidate_codes_field", default="candidate_codes") - parser.add_argument( - "--candidate_depth", - type=int, - default=None, - help="Cap on candidate last-codes tried per overflow item (default: all).", - ) parser.add_argument("--max_items_per_codebook", type=int, required=True) parser.add_argument( "--unassigned_policy", @@ -660,20 +761,28 @@ def _write_diagnostics(self, stats: AssignmentStats) -> None: "--strategy", choices=["candidate", "random"], default="candidate", - help="Reassignment strategy: 'candidate' uses model candidate_codes; " - "'random' draws random within-band last-layer codes (no candidates).", + help="Use model candidates or deterministic legacy random draws.", ) parser.add_argument( "--random_num_candidates", type=int, default=64, - help="Random last-layer codes drawn per overflow item for --strategy random.", + help="Full-space random draws per overflow item for random strategy.", ) parser.add_argument( "--rate_only", action="store_true", - help="Compute + log stats only; skip writing the map (avoids the map-write " - "cost when only the collision rate is needed).", + help="Compute and log metrics without writing map or SID group outputs.", ) parser.add_argument("--odps_data_quota_name", default="pay-as-you-go") - CollisionRunner(parser.parse_args()).run() + return parser + + +def main() -> None: + """Run collision prevention from command-line arguments.""" + config = CollisionPreventionConfig.from_namespace(build_parser().parse_args()) + CollisionRunner(config).run() + + +if __name__ == "__main__": + main() diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 8653c01f..c1621cde 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -9,6 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os import random import shutil @@ -16,16 +17,21 @@ import unittest from argparse import Namespace from collections import Counter +from unittest import mock import pyarrow as pa from pyarrow import csv, parquet +from tzrec.datasets.odps_dataset import _type_pa_to_table +from tzrec.tools.sid import collision_prevention from tzrec.tools.sid.collision_prevention import CollisionRunner # Defaults mirror the __main__ argparse; tests override only what they exercise. _DEFAULTS = dict( input_path=None, output_path=None, + original_sid_groups_output_path=None, + resolved_sid_groups_output_path=None, diagnostics_output_path=None, reader_type=None, writer_type=None, @@ -33,7 +39,6 @@ item_id_field="item_id", code_field="codes", candidate_codes_field="candidate_codes", - candidate_depth=None, max_items_per_codebook=2, unassigned_policy="error", strategy="candidate", @@ -68,10 +73,20 @@ def tearDown(self) -> None: def _run(self, inp, out, **kw): args = dict(_DEFAULTS) - args.update(input_path=inp, output_path=out) + original_groups, resolved_groups = self._group_paths(out) + args.update( + input_path=inp, + output_path=out, + original_sid_groups_output_path=original_groups, + resolved_sid_groups_output_path=resolved_groups, + ) args.update(kw) return CollisionRunner(Namespace(**args)).run() + @staticmethod + def _group_paths(out): + return f"{out}_original_groups", f"{out}_resolved_groups" + def _read_parquet(self, out_dir): return parquet.read_table(os.path.join(out_dir, "part-0.parquet")).to_pydict() @@ -116,13 +131,29 @@ def test_keep_original_when_only_origin_candidate(self) -> None: self.assertEqual(stats.unassigned_count, 3) d = self._read_parquet(out) self.assertTrue(all(tuple(cb) == (0, 0) for cb in d["codebook"])) + _, resolved_path = self._group_paths(out) + resolved = self._read_parquet(resolved_path) + self.assertEqual(len(resolved["itemids"][0]), 5) + for item_id, index in zip(d["item_id"], d["index"]): + self.assertEqual(resolved["itemids"][0][index - 1], item_id) def test_error_policy_raises_on_overflow(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") + diagnostics = os.path.join(self.test_dir, "diagnostics") _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) with self.assertRaises(RuntimeError): - self._run(inp, out, max_items_per_codebook=2) + self._run( + inp, + out, + diagnostics_output_path=diagnostics, + max_items_per_codebook=2, + ) + original_path, resolved_path = self._group_paths(out) + self.assertFalse(os.path.exists(out)) + self.assertFalse(os.path.exists(original_path)) + self.assertFalse(os.path.exists(resolved_path)) + self.assertFalse(os.path.exists(diagnostics)) def test_drop_policy_excludes_unassigned(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") @@ -132,6 +163,13 @@ def test_drop_policy_excludes_unassigned(self) -> None: self.assertEqual(stats.unassigned_count, 3) d = self._read_parquet(out) self.assertEqual(len(d["item_id"]), 2) + original_path, resolved_path = self._group_paths(out) + original = self._read_parquet(original_path) + resolved = self._read_parquet(resolved_path) + self.assertEqual(len(original["itemids"][0]), 5) + self.assertEqual(len(resolved["itemids"][0]), 2) + for item_id, index in zip(d["item_id"], d["index"]): + self.assertEqual(resolved["itemids"][0][index - 1], item_id) def test_raises_when_overflow_but_no_candidates(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") @@ -140,6 +178,110 @@ def test_raises_when_overflow_but_no_candidates(self) -> None: with self.assertRaises(ValueError): self._run(inp, out, max_items_per_codebook=2) + def test_no_overflow_does_not_require_candidates(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0, 1, 2], [[0, 0], [0, 1], [0, 2]]) + stats = self._run(inp, out, max_items_per_codebook=1) + self.assertEqual(stats.reassigned_count, 0) + self.assertEqual(len(self._read_parquet(out)["item_id"]), 3) + + def test_candidates_align_across_batches(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + candidates = [[[0, code]] for code in [1, 2, 5, 3, 6]] + _parquet(inp, list(range(5)), [[0, 0]] * 5, candidates) + self._run(inp, out, batch_size=2, max_items_per_codebook=2) + result = self._read_parquet(out) + by_id = dict(zip(result["item_id"], result["codebook"])) + self.assertEqual(by_id[0], [0, 1]) + self.assertEqual(by_id[1], [0, 2]) + self.assertEqual(by_id[3], [0, 3]) + + def test_group_outputs_follow_sid_and_slot_order(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + item_ids = list(range(10)) + codes = [ + [0, 0], + [0, 0], + [0, 0], + [0, 0], + [0, 1], + [0, 2], + [0, 2], + [1, 0], + [1, 0], + [1, 0], + ] + candidates = [[[0, 0]] for _ in item_ids] + candidates[1] = [[0, 0], [0, 1], [0, 2]] + candidates[3] = [[0, 1], [0, 2], [0, 3]] + candidates[8] = [[1, 0], [1, 1], [1, 2]] + _parquet(inp, item_ids, codes, candidates) + + self._run(inp, out, max_items_per_codebook=2) + + original_path, resolved_path = self._group_paths(out) + original = parquet.read_table(os.path.join(original_path, "part-0.parquet")) + resolved = parquet.read_table(os.path.join(resolved_path, "part-0.parquet")) + self.assertEqual(original.column_names, ["codebook", "itemids"]) + self.assertEqual(resolved.column_names, ["codebook", "itemids"]) + self.assertEqual( + original.to_pydict(), + { + "codebook": [[0, 0], [0, 1], [0, 2], [1, 0]], + "itemids": [[2, 0, 1, 3], [4], [6, 5], [9, 7, 8]], + }, + ) + self.assertEqual( + resolved.to_pydict(), + { + "codebook": [[0, 0], [0, 1], [0, 2], [0, 3], [1, 0], [1, 1]], + "itemids": [[2, 0], [4, 1], [6, 5], [3], [9, 7], [8]], + }, + ) + + resolved_items = { + tuple(codebook): item_group + for codebook, item_group in zip( + resolved["codebook"].to_pylist(), + resolved["itemids"].to_pylist(), + ) + } + item_map = self._read_parquet(out) + for item_id, codebook, index in zip( + item_map["item_id"], item_map["codebook"], item_map["index"] + ): + self.assertEqual(resolved_items[tuple(codebook)][index - 1], item_id) + + def test_no_overflow_reuses_groups_with_gapped_prefixes(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet( + inp, + list(range(6)), + [[7, 10], [7, 2], [12, 1], [7, 2], [12, 1], [7, 10]], + ) + + with mock.patch.object( + collision_prevention, + "resolve_sid_collisions", + wraps=collision_prevention.resolve_sid_collisions, + ) as resolve: + self._run(inp, out, codebook="16,16", max_items_per_codebook=2) + self.assertFalse(resolve.call_args.kwargs["collect_grouping"]) + + original_path, resolved_path = self._group_paths(out) + original = parquet.read_table( + os.path.join(original_path, "part-0.parquet") + ).to_pydict() + resolved = parquet.read_table( + os.path.join(resolved_path, "part-0.parquet") + ).to_pydict() + self.assertEqual(original, resolved) + self.assertEqual(original["codebook"], [[7, 2], [7, 10], [12, 1]]) + # ---- random strategy ---- def test_random_reassigns_within_band(self) -> None: @@ -157,6 +299,14 @@ def test_random_reassigns_within_band(self) -> None: d = self._read_parquet(out) self.assertTrue(all(cb[0] == 0 for cb in d["codebook"])) self.assertLessEqual(max(Counter(map(tuple, d["codebook"])).values()), 2) + _, resolved_path = self._group_paths(out) + resolved = self._read_parquet(resolved_path) + resolved_items = { + tuple(codebook): item_group + for codebook, item_group in zip(resolved["codebook"], resolved["itemids"]) + } + for item_id, codebook, index in zip(d["item_id"], d["codebook"], d["index"]): + self.assertEqual(resolved_items[tuple(codebook)][index - 1], item_id) def test_random_requires_last_codebook_ge_2(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") @@ -208,10 +358,15 @@ def decisions(order): unassigned_policy="keep_original", ) d = self._read_parquet(out) - return { - d["item_id"][i]: tuple(d["codebook"][i]) - for i in range(len(d["item_id"])) - } + original_path, resolved_path = self._group_paths(out) + return ( + { + d["item_id"][i]: tuple(d["codebook"][i]) + for i in range(len(d["item_id"])) + }, + self._read_parquet(original_path), + self._read_parquet(resolved_path), + ) base = decisions(list(range(n))) self.assertEqual(base, decisions(list(range(n)))) # run-twice @@ -247,6 +402,118 @@ def test_csv_backend_within_band(self) -> None: self.assertEqual(result.schema.field("codebook").type, pa.string()) self.assertTrue(all(s.startswith("0|") for s in result["codebook"].to_pylist())) + original_path, resolved_path = self._group_paths(out) + for path in (original_path, resolved_path): + groups = csv.read_csv(os.path.join(path, "part-0.csv")) + self.assertEqual(groups.column_names, ["codebook", "itemids"]) + self.assertTrue( + all( + isinstance(json.loads(itemids), list) + for itemids in groups["itemids"].to_pylist() + ) + ) + + def test_csv_group_itemids_use_json_escaping(self) -> None: + inp = os.path.join(self.test_dir, "in.csv") + out = os.path.join(self.test_dir, "out") + item_ids = ["a;b", "x|y", "with,comma", 'with"quote', "back\\slash"] + csv.write_csv( + pa.table( + { + "item_id": item_ids, + "codes": ["0|0"] * len(item_ids), + } + ), + inp, + ) + + self._run( + inp, + out, + reader_type="CsvReader", + writer_type="CsvWriter", + max_items_per_codebook=len(item_ids), + ) + + original_path, _ = self._group_paths(out) + groups = csv.read_csv(os.path.join(original_path, "part-0.csv")) + decoded_ids = json.loads(groups["itemids"][0].as_py()) + self.assertCountEqual(decoded_ids, item_ids) + + def test_parquet_input_with_csv_output(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0, 1, 2], [[0, 0]] * 3, [[[0, 1]]] * 3) + self._run( + inp, + out, + writer_type="CsvWriter", + max_items_per_codebook=2, + ) + result = csv.read_csv(os.path.join(out, "part-0.csv")) + self.assertEqual(result.schema.field("codebook").type, pa.string()) + self.assertTrue(all("|" in sid for sid in result["codebook"].to_pylist())) + original_path, resolved_path = self._group_paths(out) + for path in (original_path, resolved_path): + schema = csv.read_csv(os.path.join(path, "part-0.csv")).schema + self.assertEqual(schema.names, ["codebook", "itemids"]) + self.assertEqual(schema.field("itemids").type, pa.string()) + + def test_csv_input_with_parquet_output(self) -> None: + inp = os.path.join(self.test_dir, "in.csv") + out = os.path.join(self.test_dir, "out") + csv.write_csv( + pa.table( + { + "item_id": ["a", "b", "c"], + "codes": ["0|0"] * 3, + "candidate_codes": ["0|1"] * 3, + } + ), + inp, + ) + self._run( + inp, + out, + reader_type="CsvReader", + writer_type="ParquetWriter", + max_items_per_codebook=2, + ) + schema = parquet.read_table(os.path.join(out, "part-0.parquet")).schema + self.assertEqual(schema.field("codebook").type, pa.list_(pa.int64())) + original_path, resolved_path = self._group_paths(out) + for path in (original_path, resolved_path): + group_schema = parquet.read_table( + os.path.join(path, "part-0.parquet") + ).schema + self.assertEqual(group_schema.names, ["codebook", "itemids"]) + self.assertEqual(group_schema.field("itemids").type, pa.list_(pa.string())) + + def test_preserves_integer_item_id_type(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + parquet.write_table( + pa.table( + { + "item_id": pa.array([0, 1, 2], type=pa.int32()), + "codes": pa.array( + [[0, 0], [0, 1], [0, 2]], + type=pa.list_(pa.int64()), + ), + } + ), + inp, + ) + self._run(inp, out, max_items_per_codebook=1) + schema = parquet.read_table(os.path.join(out, "part-0.parquet")).schema + self.assertEqual(schema.field("item_id").type, pa.int32()) + original_path, resolved_path = self._group_paths(out) + for path in (original_path, resolved_path): + group_schema = parquet.read_table( + os.path.join(path, "part-0.parquet") + ).schema + self.assertEqual(group_schema.field("itemids").type, pa.list_(pa.int32())) + def test_writer_defaults_to_reader(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") @@ -256,19 +523,209 @@ def test_writer_defaults_to_reader(self) -> None: # ---- misc ---- + def test_group_output_paths_are_required_for_normal_runs(self) -> None: + args = dict(_DEFAULTS) + args.update(input_path="input", output_path="map") + with self.assertRaisesRegex(ValueError, "both SID group output paths"): + CollisionRunner(Namespace(**args)) + + def test_group_output_paths_must_be_supplied_together(self) -> None: + args = dict(_DEFAULTS) + args.update( + input_path="input", + output_path="map", + original_sid_groups_output_path="original", + ) + with self.assertRaisesRegex(ValueError, "must be supplied together"): + CollisionRunner(Namespace(**args)) + + def test_output_paths_must_be_distinct(self) -> None: + args = dict(_DEFAULTS) + args.update( + input_path="input", + output_path="map", + original_sid_groups_output_path="map/../map", + resolved_sid_groups_output_path="resolved", + ) + with self.assertRaisesRegex(ValueError, "must differ from output_path"): + CollisionRunner(Namespace(**args)) + + def test_odps_output_paths_reject_trailing_slash_alias(self) -> None: + args = dict(_DEFAULTS) + args.update( + input_path="input", + output_path="odps://project/tables/map", + original_sid_groups_output_path="odps://project/tables/map/", + resolved_sid_groups_output_path="odps://project/tables/resolved", + ) + with self.assertRaisesRegex(ValueError, "must differ from output_path"): + CollisionRunner(Namespace(**args)) + + def test_odps_output_paths_reject_ignored_partition_alias(self) -> None: + args = dict(_DEFAULTS) + args.update( + input_path="input", + output_path="odps://project/tables/schema.map/dt=20260713", + original_sid_groups_output_path=( + "odps://project/tables/schema.map/dt=20260713&dt=20260712" + ), + resolved_sid_groups_output_path="odps://project/tables/schema.resolved", + ) + with self.assertRaisesRegex(ValueError, "must differ from output_path"): + CollisionRunner(Namespace(**args)) + + def test_odps_group_writes_use_native_array_columns(self) -> None: + class OdpsWriter: + def __init__(self) -> None: + self.writes = [] + self.closed = False + + def write(self, output_dict) -> None: + self.writes.append(output_dict) + + def close(self) -> None: + self.closed = True + + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "map") + diagnostics = os.path.join(self.test_dir, "diagnostics") + parquet.write_table( + pa.table( + { + "item_id": pa.array([1, 2], type=pa.int32()), + "codes": pa.array([[0, 0], [0, 1]], type=pa.list_(pa.int64())), + } + ), + inp, + ) + created_writers = [] + + def make_writer(output_path): + writer = OdpsWriter() + created_writers.append((output_path, writer)) + return writer + + with mock.patch.object( + CollisionRunner, "_make_writer", side_effect=make_writer + ): + self._run( + inp, + out, + writer_type="OdpsWriter", + diagnostics_output_path=diagnostics, + max_items_per_codebook=1, + ) + + original_path, resolved_path = self._group_paths(out) + self.assertEqual( + [path for path, _ in created_writers], + [original_path, resolved_path, out, diagnostics], + ) + self.assertTrue(all(writer.closed for _, writer in created_writers)) + for _, writer in created_writers[:2]: + columns = writer.writes[0] + self.assertEqual(list(columns), ["codebook", "itemids"]) + self.assertEqual(columns["codebook"].type, pa.list_(pa.int64())) + self.assertEqual(columns["itemids"].type, pa.list_(pa.int32())) + self.assertEqual( + _type_pa_to_table(columns["codebook"].type), "ARRAY" + ) + self.assertEqual(_type_pa_to_table(columns["itemids"].type), "ARRAY") + self.assertEqual(_type_pa_to_table(pa.list_(pa.string())), "ARRAY") + + def test_group_writer_chunks_by_item_count_without_splitting_groups(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(4)), [[0, 0], [0, 0], [0, 0], [0, 1]]) + + with mock.patch("tzrec.tools.sid.collision_prevention._WRITE_CHUNK", 2): + self._run(inp, out, max_items_per_codebook=3) + + original_path, _ = self._group_paths(out) + output_file = os.path.join(original_path, "part-0.parquet") + self.assertEqual(parquet.ParquetFile(output_file).metadata.num_row_groups, 2) + groups = parquet.read_table(output_file).to_pydict() + self.assertEqual([len(group) for group in groups["itemids"]], [3, 1]) + + def test_writers_bound_codebook_list_offsets(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0, 1, 2], [[0, 0], [0, 1], [0, 2]]) + + with ( + mock.patch( + "tzrec.tools.sid.collision_prevention._ARROW_LIST_OFFSET_MAX", 4 + ), + mock.patch("tzrec.tools.sid.collision_prevention._WRITE_CHUNK", 100), + ): + self._run(inp, out, max_items_per_codebook=1) + + original_path, resolved_path = self._group_paths(out) + for path in (out, original_path, resolved_path): + output_file = os.path.join(path, "part-0.parquet") + self.assertEqual( + parquet.ParquetFile(output_file).metadata.num_row_groups, 2 + ) + def test_rate_only_skips_write(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1], [0, 2], [0, 3]]] * 5) + with mock.patch.object( + collision_prevention, + "resolve_sid_collisions", + wraps=collision_prevention.resolve_sid_collisions, + ) as resolve: + stats = self._run( + inp, + out, + max_items_per_codebook=2, + unassigned_policy="keep_original", + rate_only=True, + ) + self.assertFalse(resolve.call_args.kwargs["collect_grouping"]) + self.assertEqual(stats.reassigned_count, 3) + self.assertFalse(os.path.exists(os.path.join(out, "part-0.parquet"))) + original_path, resolved_path = self._group_paths(out) + self.assertFalse(os.path.exists(original_path)) + self.assertFalse(os.path.exists(resolved_path)) + + def test_rate_only_allows_omitted_group_paths(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0, 1], [[0, 0], [0, 0]]) + stats = self._run( inp, out, + original_sid_groups_output_path=None, + resolved_sid_groups_output_path=None, + max_items_per_codebook=2, + rate_only=True, + ) + + self.assertEqual(stats.total_items, 2) + self.assertFalse(os.path.exists(out)) + + def test_rate_only_still_writes_diagnostics(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + diag = os.path.join(self.test_dir, "diag") + _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1]]] * 5) + self._run( + inp, + out, + diagnostics_output_path=diag, max_items_per_codebook=2, unassigned_policy="keep_original", rate_only=True, ) - self.assertEqual(stats.reassigned_count, 3) + result = parquet.read_table(os.path.join(diag, "part-0.parquet")).to_pydict() + self.assertEqual(result["total_items"], [5]) self.assertFalse(os.path.exists(os.path.join(out, "part-0.parquet"))) + original_path, resolved_path = self._group_paths(out) + self.assertFalse(os.path.exists(original_path)) + self.assertFalse(os.path.exists(resolved_path)) def test_diagnostics_output(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py new file mode 100644 index 00000000..cf79854b --- /dev/null +++ b/tzrec/tools/sid/collision_resolution.py @@ -0,0 +1,659 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NumPy collision-resolution core for semantic IDs.""" + +from dataclasses import dataclass +from typing import Dict, Iterable, Optional + +import numpy as np + +_MASK64 = (1 << 64) - 1 +_SEED = 2026 +_SPLITMIX_INCREMENT = 0x9E3779B97F4A7C15 +_FALLBACK_POLICIES = frozenset({"error", "drop", "keep_original"}) +_GROUPING_ROW_CHUNK = 1_000_000 + + +@dataclass(frozen=True) +class CollisionResolutionConfig: + """Configuration for within-band SID collision resolution. + + Args: + layer_sizes: Cardinality of each SID layer. + capacity: Maximum number of retained items in one SID bucket. + fallback_policy: Action when no candidate has a free slot. Supported + values are ``error``, ``drop``, and ``keep_original``. + """ + + layer_sizes: tuple[int, ...] + capacity: int + fallback_policy: str + + def __post_init__(self) -> None: + if not self.layer_sizes: + raise ValueError("layer_sizes must contain at least one layer.") + if any(size <= 0 for size in self.layer_sizes): + raise ValueError("all layer_sizes must be positive.") + if self.capacity < 1: + raise ValueError(f"capacity must be >= 1, got {self.capacity}.") + if self.fallback_policy not in _FALLBACK_POLICIES: + raise ValueError( + "fallback_policy must be one of error, drop, or keep_original, " + f"got {self.fallback_policy!r}." + ) + + +@dataclass(frozen=True) +class CollisionPlan: + """Compact grouping plan consumed by collision resolution. + + The plan stores full-length last-code, dense bucket-ID, and index arrays, + plus only the overflow-aligned identity and ordering data needed by + candidate loading. It deliberately does not retain the full input code + matrix. + """ + + item_count: int + original_last_codes: np.ndarray + origin_bucket_ids: np.ndarray + initial_slot_indices: np.ndarray + occupied_sid_keys: np.ndarray + raw_bucket_counts: np.ndarray + overflow_rows: np.ndarray + overflow_item_ids: np.ndarray + overflow_band_bases: np.ndarray + overflow_origin_last_codes: np.ndarray + config: CollisionResolutionConfig + + @property + def origin_sid_keys(self) -> np.ndarray: + """Return flattened original SID keys aligned with input rows.""" + return self.occupied_sid_keys[self.origin_bucket_ids] + + +@dataclass(frozen=True) +class CollisionResolutionStats: + """Summary statistics for a collision-resolution run.""" + + total_items: int + raw_collision_buckets: int + final_collision_buckets: int + relocated_count: int + unresolved_count: int + max_final_bucket_size: int + + @property + def reassigned_count(self) -> int: + """Return the legacy name for ``relocated_count``.""" + return self.relocated_count + + @property + def unassigned_count(self) -> int: + """Return the legacy name for ``unresolved_count``.""" + return self.unresolved_count + + def to_output_dict(self) -> Dict[str, int]: + """Return diagnostics using the existing output column names.""" + return { + "total_items": self.total_items, + "raw_collision_buckets": self.raw_collision_buckets, + "final_collision_buckets": self.final_collision_buckets, + "reassigned_count": self.relocated_count, + "unassigned_count": self.unresolved_count, + "max_final_bucket_size": self.max_final_bucket_size, + } + + +@dataclass(frozen=True) +class CollisionResolutionResult: + """Resolved last codes, indexes, row retention, and diagnostics.""" + + resolved_last_codes: np.ndarray + slot_indices: np.ndarray + retained_mask: Optional[np.ndarray] + unresolved_rows: np.ndarray + final_occupied_sid_keys: np.ndarray + final_bucket_counts: np.ndarray + slots_match_final_sid_keys: bool + grouping_collected: bool + stats: CollisionResolutionStats + + +@dataclass(frozen=True) +class CodebookItemGrouping: + """CSR-like row grouping for sorted flattened SID keys. + + ``row_order[offsets[i]:offsets[i + 1]]`` contains the original row indices + for ``sid_keys[i]``. Rows within a bucket follow their one-based slot order. + + Args: + sid_keys: Sorted flattened SID keys, one per occupied bucket. + counts: Item counts aligned with ``sid_keys``. + row_order: Original row indices flattened in SID and slot order. + """ + + sid_keys: np.ndarray + counts: np.ndarray + row_order: np.ndarray + + @property + def offsets(self) -> np.ndarray: + """Return CSR offsets into ``row_order`` for every SID bucket.""" + offsets = np.empty(self.counts.shape[0] + 1, dtype=np.int64) + offsets[0] = 0 + np.cumsum(self.counts, dtype=np.int64, out=offsets[1:]) + return offsets + + @property + def representative_rows(self) -> np.ndarray: + """Return the first original row index in every occupied bucket.""" + return self.row_order[self.offsets[:-1]] + + +def _splitmix64(values: np.ndarray) -> np.ndarray: + """Hash a uint64 array with the collision tool's fixed SplitMix64 seed.""" + with np.errstate(over="ignore"): + mixed = values.astype(np.uint64) + np.uint64( + (_SEED * _SPLITMIX_INCREMENT) & _MASK64 + ) + mixed = (mixed ^ (mixed >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) + mixed = (mixed ^ (mixed >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) + return mixed ^ (mixed >> np.uint64(31)) + + +def stable_order_hash(item_ids: np.ndarray) -> np.ndarray: + """Return the existing order-independent uint64 tie-break hashes. + + Integer IDs are converted directly to uint64. Other ID types retain the + existing ``pandas.util.hash_array`` folding step before SplitMix64. + + Args: + item_ids: One-dimensional item ID array. + + Returns: + A uint64 hash array aligned with ``item_ids``. + + Raises: + ValueError: If ``item_ids`` is not one-dimensional. + """ + item_ids = np.asarray(item_ids) + if item_ids.ndim != 1: + raise ValueError(f"item_ids must be 1-D, got shape {item_ids.shape}.") + if np.issubdtype(item_ids.dtype, np.integer): + base = item_ids.astype(np.uint64) + else: + import pandas as pd + + base = pd.util.hash_array(np.asarray(item_ids, dtype=object)) + return _splitmix64(base) + + +def _band_ids(codes: np.ndarray, layer_sizes: tuple[int, ...]) -> np.ndarray: + """Return a dense integer ID for each prefix band.""" + row_count, layer_count = codes.shape + if layer_count == 1: + return np.zeros(row_count, dtype=np.int64) + keys = codes[:, 0].astype(np.int64) + for layer in range(1, layer_count - 1): + keys *= layer_sizes[layer] + keys += codes[:, layer] + return np.unique(keys, return_inverse=True)[1].astype(np.int64, copy=False) + + +def _within_bucket_rank( + band_ids: np.ndarray, last_codes: np.ndarray, order_hashes: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return per-row rank, bucket IDs, sorted rows, representatives, and counts.""" + row_count = band_ids.shape[0] + if row_count == 0: + empty = np.empty(0, dtype=np.int64) + return empty, empty.copy(), empty.copy(), empty.copy(), empty.copy() + sorted_rows = np.lexsort((order_hashes, last_codes, band_ids)) + sorted_bands = band_ids[sorted_rows] + sorted_last_codes = last_codes[sorted_rows] + is_first = np.empty(row_count, dtype=bool) + is_first[0] = True + is_first[1:] = (sorted_bands[1:] != sorted_bands[:-1]) | ( + sorted_last_codes[1:] != sorted_last_codes[:-1] + ) + first_sorted_rows = np.flatnonzero(is_first) + bucket_counts = np.diff(np.append(first_sorted_rows, row_count)) + del sorted_bands, sorted_last_codes, is_first + bucket_ids = np.empty(row_count, dtype=np.int64) + bucket_ranks = np.empty(row_count, dtype=np.int64) + for start in range(0, row_count, _GROUPING_ROW_CHUNK): + end = min(start + _GROUPING_ROW_CHUNK, row_count) + first_bucket = int(np.searchsorted(first_sorted_rows, start, side="right") - 1) + last_bucket = int(np.searchsorted(first_sorted_rows, end - 1, side="right") - 1) + local_starts = first_sorted_rows[first_bucket : last_bucket + 1].copy() + local_starts[0] = start + local_counts = np.diff(np.append(local_starts, end)) + sorted_bucket_ids = np.repeat( + np.arange(first_bucket, last_bucket + 1, dtype=np.int64), + local_counts, + ) + rows = sorted_rows[start:end] + bucket_ids[rows] = sorted_bucket_ids + bucket_ranks[rows] = ( + np.arange(start, end, dtype=np.int64) - first_sorted_rows[sorted_bucket_ids] + ) + return ( + bucket_ranks, + bucket_ids, + sorted_rows, + sorted_rows[first_sorted_rows], + bucket_counts, + ) + + +def prepare_collision_plan( + item_ids: np.ndarray, + codes: np.ndarray, + config: CollisionResolutionConfig, +) -> CollisionPlan: + """Group SID buckets and identify overflow rows in stable processing order. + + Args: + item_ids: One-dimensional item IDs aligned with ``codes``. + codes: Integer SID matrix with shape ``(N, number_of_layers)``. + config: Collision capacity, shape, and fallback configuration. + + Returns: + A compact plan for candidate loading and collision resolution. + + Raises: + ValueError: If input dimensions, row counts, or layer counts disagree. + """ + item_ids = np.asarray(item_ids) + codes = np.asarray(codes) + if item_ids.ndim != 1: + raise ValueError(f"item_ids must be 1-D, got shape {item_ids.shape}.") + if codes.ndim != 2: + raise ValueError(f"codes must be 2-D, got shape {codes.shape}.") + if item_ids.shape[0] != codes.shape[0]: + raise ValueError( + "item_ids and codes must have the same row count, got " + f"{item_ids.shape[0]} and {codes.shape[0]}." + ) + layer_count = codes.shape[1] + if layer_count != len(config.layer_sizes): + raise ValueError( + f"codes have {layer_count} layers but layer_sizes has " + f"{len(config.layer_sizes)}." + ) + + item_count = codes.shape[0] + original_last_codes = codes[:, -1].astype(np.int64, copy=False) + band_ids = _band_ids(codes, config.layer_sizes) + order_hashes = stable_order_hash(item_ids) + ( + bucket_ranks, + origin_bucket_ids, + sorted_rows, + representative_rows, + raw_bucket_counts, + ) = _within_bucket_rank(band_ids, original_last_codes, order_hashes) + + overflow_rows = sorted_rows[bucket_ranks[sorted_rows] >= config.capacity] + last_size = config.layer_sizes[-1] + occupied_sid_keys = ( + band_ids[representative_rows] * last_size + + original_last_codes[representative_rows] + ) + overflow_band_bases = band_ids[overflow_rows] * last_size + + return CollisionPlan( + item_count=item_count, + original_last_codes=original_last_codes, + origin_bucket_ids=origin_bucket_ids, + initial_slot_indices=bucket_ranks + 1, + occupied_sid_keys=occupied_sid_keys, + raw_bucket_counts=raw_bucket_counts, + overflow_rows=overflow_rows, + overflow_item_ids=item_ids[overflow_rows].copy(), + overflow_band_bases=overflow_band_bases, + overflow_origin_last_codes=original_last_codes[overflow_rows].copy(), + config=config, + ) + + +def generate_random_candidate_last_codes( + item_ids: np.ndarray, num: int, last_size: int +) -> np.ndarray: + """Generate the existing deterministic full-space random candidate draws. + + Sampling is with replacement and includes each item's original code because + that code is not an input to this function. The placement step skips an + origin draw without replacing it, preserving current behavior. + + Args: + item_ids: One-dimensional IDs for the overflow rows. + num: Requested number of raw draws per row. As in the current strategy, + this is capped at ``last_size - 1``. + last_size: Cardinality of the last SID layer. + + Returns: + An ``(len(item_ids), min(num, last_size - 1))`` int64 matrix. + + Raises: + ValueError: If ``num`` is negative or ``last_size`` is smaller than two. + """ + if last_size < 2: + raise ValueError("random candidates require last_size >= 2.") + if num < 0: + raise ValueError(f"num must be >= 0, got {num}.") + effective_num = min(num, last_size - 1) + hashes = stable_order_hash(item_ids) + draw_indices = np.arange(effective_num, dtype=np.uint64) + with np.errstate(over="ignore"): + mixed = _splitmix64( + hashes[:, None] + draw_indices[None, :] * np.uint64(_SPLITMIX_INCREMENT) + ) + return (mixed % np.uint64(last_size)).astype(np.int64) + + +def resolve_sid_collisions( + plan: CollisionPlan, + candidate_last_codes: np.ndarray, + *, + collect_grouping: bool = True, +) -> CollisionResolutionResult: + """Greedily relocate overflow rows to their first free candidate slot. + + Candidate rows must be aligned with ``plan.overflow_rows``. Candidate values + intentionally receive no range validation so negative and out-of-range + values retain the existing key and modulo behavior. + + Args: + plan: Grouping and overflow plan from :func:`prepare_collision_plan`. + candidate_last_codes: Ordered candidate matrix with shape ``(M, K)``, + where ``M`` equals the number of overflow rows. + collect_grouping: Whether to collect final SID keys and bucket counts + for :func:`build_resolved_item_grouping`. Disable this in + diagnostics-only runs to avoid the grouping arrays and sorts. + + Returns: + Resolved last-layer codes, slot indices, retention mask, unresolved row + indices, and diagnostics. + + Raises: + RuntimeError: If a row is unresolved under the ``error`` policy. + TypeError: If candidate codes do not use an integer dtype. + ValueError: If candidates are not a row-aligned two-dimensional matrix. + """ + candidates = np.asarray(candidate_last_codes) + if candidates.ndim != 2: + raise ValueError( + f"candidate_last_codes must be 2-D, got shape {candidates.shape}." + ) + overflow_count = plan.overflow_rows.shape[0] + if candidates.shape[0] != overflow_count: + raise ValueError( + "candidate_last_codes must be row-aligned with overflow_rows, got " + f"{candidates.shape[0]} candidate rows for {overflow_count} " + "overflow rows." + ) + if not np.issubdtype(candidates.dtype, np.integer): + raise TypeError("candidate_last_codes must use an integer dtype.") + + capacity = plan.config.capacity + initial_counts = np.minimum(plan.raw_bucket_counts, capacity) + slot_counts = dict(zip(plan.occupied_sid_keys.tolist(), initial_counts.tolist())) + resolved_last_codes = plan.original_last_codes.copy() + slot_indices = plan.initial_slot_indices.copy() + relocated_rows = [] + relocated_last_codes = [] + relocated_indices = [] + unresolved_rows = [] + slots_match_final_sid_keys = True + get_slot_count = slot_counts.get + last_size = plan.config.layer_sizes[-1] + + overflow_rows = plan.overflow_rows.tolist() + band_bases = plan.overflow_band_bases.tolist() + origin_last_codes = plan.overflow_origin_last_codes.tolist() + for position, row in enumerate(overflow_rows): + band_base = band_bases[position] + origin_last_code = origin_last_codes[position] + for candidate in candidates[position].tolist(): + if candidate == origin_last_code: + continue + destination_key = band_base + candidate + destination_count = get_slot_count(destination_key, 0) + if destination_count < capacity: + slot_counts[destination_key] = destination_count + 1 + relocated_rows.append(row) + relocated_last_codes.append(destination_key % last_size) + relocated_indices.append(destination_count + 1) + if candidate < 0 or candidate >= last_size: + slots_match_final_sid_keys = False + break + else: + unresolved_rows.append(row) + + if relocated_rows: + resolved_last_codes[relocated_rows] = relocated_last_codes + slot_indices[relocated_rows] = relocated_indices + + retained_mask = None + fallback_policy = plan.config.fallback_policy + if unresolved_rows and fallback_policy == "error": + preview = ",".join(str(row) for row in unresolved_rows[:10]) + raise RuntimeError( + f"{len(unresolved_rows)} items could not be placed within capacity; " + f"first unresolved row indices: {preview}" + ) + if unresolved_rows and fallback_policy == "drop": + retained_mask = np.ones(plan.item_count, dtype=bool) + retained_mask[unresolved_rows] = False + elif unresolved_rows and fallback_policy == "keep_original": + for row in unresolved_rows: + origin_key = int(plan.occupied_sid_keys[plan.origin_bucket_ids[row]]) + origin_count = get_slot_count(origin_key, 0) + 1 + slot_counts[origin_key] = origin_count + slot_indices[row] = origin_count + + final_occupied_sid_keys = np.empty(0, dtype=np.int64) + final_bucket_counts = np.empty(0, dtype=np.int64) + if collect_grouping: + occupancy_counts = np.fromiter(slot_counts.values(), dtype=np.int64) + final_collision_buckets = int((occupancy_counts > capacity).sum()) + max_final_bucket_size = ( + int(occupancy_counts.max()) if occupancy_counts.size else 0 + ) + if slots_match_final_sid_keys: + final_occupied_sid_keys = np.fromiter(slot_counts, dtype=np.int64) + occupancy_order = np.argsort(final_occupied_sid_keys, kind="stable") + final_occupied_sid_keys = final_occupied_sid_keys[occupancy_order] + final_bucket_counts = occupancy_counts[occupancy_order] + else: + retained_rows = ( + np.flatnonzero(retained_mask) + if retained_mask is not None + else np.arange(plan.item_count, dtype=np.int64) + ) + origin_keys = plan.occupied_sid_keys[plan.origin_bucket_ids[retained_rows]] + final_sid_keys = ( + origin_keys + - plan.original_last_codes[retained_rows] + + resolved_last_codes[retained_rows] + ) + final_occupied_sid_keys, final_bucket_counts = np.unique( + final_sid_keys, return_counts=True + ) + else: + final_collision_buckets = 0 + max_final_bucket_size = 0 + for count in slot_counts.values(): + final_collision_buckets += count > capacity + max_final_bucket_size = max(max_final_bucket_size, count) + unresolved_array = np.asarray(unresolved_rows, dtype=np.int64) + stats = CollisionResolutionStats( + total_items=plan.item_count, + raw_collision_buckets=int((plan.raw_bucket_counts > capacity).sum()), + final_collision_buckets=final_collision_buckets, + relocated_count=len(relocated_rows), + unresolved_count=len(unresolved_rows), + max_final_bucket_size=max_final_bucket_size, + ) + return CollisionResolutionResult( + resolved_last_codes=resolved_last_codes, + slot_indices=slot_indices, + retained_mask=retained_mask, + unresolved_rows=unresolved_array, + final_occupied_sid_keys=final_occupied_sid_keys, + final_bucket_counts=final_bucket_counts, + slots_match_final_sid_keys=slots_match_final_sid_keys, + grouping_collected=collect_grouping, + stats=stats, + ) + + +def _scatter_item_grouping( + sid_keys: np.ndarray, + counts: np.ndarray, + row_count: int, + chunks: Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]], +) -> CodebookItemGrouping: + """Build a CSR row order from bounded row, bucket-ID, and slot chunks.""" + grouping = CodebookItemGrouping( + sid_keys=sid_keys, + counts=counts, + row_order=np.empty(row_count, dtype=np.int64), + ) + offsets = grouping.offsets + if int(offsets[-1]) != row_count: + raise RuntimeError( + "bucket counts do not match retained row count: " + f"{int(offsets[-1])} vs {row_count}." + ) + written_count = 0 + for rows, bucket_ids, slot_indices in chunks: + if rows.shape != bucket_ids.shape or rows.shape != slot_indices.shape: + raise RuntimeError("grouping row, bucket, and slot chunks must align.") + if rows.size == 0: + continue + if np.any(bucket_ids < 0) or np.any(bucket_ids >= counts.shape[0]): + raise RuntimeError("grouping bucket IDs are outside the SID key array.") + positions = offsets[bucket_ids] + slot_indices - 1 + if np.any(slot_indices < 1) or np.any(positions >= offsets[bucket_ids + 1]): + raise RuntimeError("slot indices are outside their SID bucket bounds.") + grouping.row_order[positions] = rows + written_count += rows.shape[0] + if written_count != row_count: + raise RuntimeError( + f"grouping chunks contain {written_count} rows, expected {row_count}." + ) + return grouping + + +def build_original_item_grouping(plan: CollisionPlan) -> CodebookItemGrouping: + """Group all original rows by SID and initial one-based slot index. + + Args: + plan: Grouping and overflow plan from :func:`prepare_collision_plan`. + + Returns: + Sorted original SID keys, bucket counts, and grouped original row order. + """ + + def row_chunks() -> Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]]: + for start in range(0, plan.item_count, _GROUPING_ROW_CHUNK): + end = min(start + _GROUPING_ROW_CHUNK, plan.item_count) + yield ( + np.arange(start, end, dtype=np.int64), + plan.origin_bucket_ids[start:end], + plan.initial_slot_indices[start:end], + ) + + return _scatter_item_grouping( + plan.occupied_sid_keys, + plan.raw_bucket_counts, + plan.item_count, + row_chunks(), + ) + + +def build_resolved_item_grouping( + plan: CollisionPlan, result: CollisionResolutionResult +) -> CodebookItemGrouping: + """Group retained rows by emitted SID and final one-based slot index. + + Canonical in-range candidates use a linear scatter from final slot indices. + Malformed candidates can make internal occupancy keys differ from emitted SID + keys because the legacy resolver applies modulo only to the emitted last code. + That compatibility path groups by emitted keys and reported slots instead. + + Args: + plan: Grouping and overflow plan from :func:`prepare_collision_plan`. + result: Collision output from :func:`resolve_sid_collisions`. + + Returns: + Sorted final SID keys, bucket counts, and grouped retained row order. + + Raises: + RuntimeError: If grouping metadata was not collected or result bucket + keys do not cover an emitted SID. + """ + if not result.grouping_collected: + raise RuntimeError( + "resolved item grouping metadata was not collected; call " + "resolve_sid_collisions with collect_grouping=True." + ) + if not result.slots_match_final_sid_keys: + rows = ( + np.flatnonzero(result.retained_mask) + if result.retained_mask is not None + else np.arange(plan.item_count, dtype=np.int64) + ) + origin_keys = plan.occupied_sid_keys[plan.origin_bucket_ids[rows]] + emitted_sid_keys = ( + origin_keys + - plan.original_last_codes[rows] + + result.resolved_last_codes[rows] + ) + row_order = np.lexsort((rows, result.slot_indices[rows], emitted_sid_keys)) + return CodebookItemGrouping( + sid_keys=result.final_occupied_sid_keys, + counts=result.final_bucket_counts, + row_order=rows[row_order], + ) + + def row_chunks() -> Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]]: + for start in range(0, plan.item_count, _GROUPING_ROW_CHUNK): + end = min(start + _GROUPING_ROW_CHUNK, plan.item_count) + if result.retained_mask is None: + rows = np.arange(start, end, dtype=np.int64) + else: + rows = np.flatnonzero(result.retained_mask[start:end]) + rows += start + if rows.size == 0: + continue + emitted_sid_keys = plan.occupied_sid_keys[plan.origin_bucket_ids[rows]] + emitted_sid_keys -= plan.original_last_codes[rows] + emitted_sid_keys += result.resolved_last_codes[rows] + bucket_ids = np.searchsorted( + result.final_occupied_sid_keys, emitted_sid_keys + ) + in_bounds = bucket_ids < result.final_occupied_sid_keys.shape[0] + if not np.all(in_bounds) or not np.all( + result.final_occupied_sid_keys[bucket_ids] == emitted_sid_keys + ): + raise RuntimeError("final bucket keys do not cover every emitted SID.") + yield rows, bucket_ids, result.slot_indices[rows] + + return _scatter_item_grouping( + result.final_occupied_sid_keys, + result.final_bucket_counts, + int(result.final_bucket_counts.sum()), + row_chunks(), + ) diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py new file mode 100644 index 00000000..0a0cb730 --- /dev/null +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -0,0 +1,300 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest import mock + +import numpy as np + +from tzrec.tools.sid import collision_resolution +from tzrec.tools.sid.collision_resolution import ( + CollisionResolutionConfig, + build_original_item_grouping, + build_resolved_item_grouping, + generate_random_candidate_last_codes, + prepare_collision_plan, + resolve_sid_collisions, +) + + +class CollisionResolutionTest(unittest.TestCase): + def test_golden_candidate_resolution(self) -> None: + item_ids = np.arange(10, dtype=np.int64) + codes = np.asarray( + [ + [0, 0], + [0, 0], + [0, 0], + [0, 0], + [0, 1], + [0, 2], + [0, 2], + [1, 0], + [1, 0], + [1, 0], + ], + dtype=np.int64, + ) + config = CollisionResolutionConfig((2, 4), 2, "error") + with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): + plan = prepare_collision_plan(item_ids, codes, config) + + np.testing.assert_array_equal(plan.overflow_rows, [1, 3, 8]) + np.testing.assert_array_equal(plan.overflow_item_ids, [1, 3, 8]) + np.testing.assert_array_equal( + plan.origin_bucket_ids, [0, 0, 0, 0, 1, 2, 2, 3, 3, 3] + ) + np.testing.assert_array_equal(plan.occupied_sid_keys, [0, 1, 2, 4]) + np.testing.assert_array_equal( + plan.origin_sid_keys, [0, 0, 0, 0, 1, 2, 2, 4, 4, 4] + ) + candidates = np.asarray([[0, 1, 2], [1, 2, 3], [0, 1, 2]], dtype=np.int64) + result = resolve_sid_collisions(plan, candidates) + + np.testing.assert_array_equal( + result.resolved_last_codes, [0, 1, 0, 3, 1, 2, 2, 0, 1, 0] + ) + np.testing.assert_array_equal( + result.slot_indices, [2, 2, 1, 1, 1, 2, 1, 2, 1, 1] + ) + np.testing.assert_array_equal(result.unresolved_rows, []) + np.testing.assert_array_equal( + result.final_occupied_sid_keys, [0, 1, 2, 3, 4, 5] + ) + np.testing.assert_array_equal(result.final_bucket_counts, [2, 2, 2, 1, 2, 1]) + self.assertTrue(result.slots_match_final_sid_keys) + self.assertTrue(result.grouping_collected) + self.assertIsNone(result.retained_mask) + self.assertEqual(result.stats.total_items, 10) + self.assertEqual(result.stats.raw_collision_buckets, 2) + self.assertEqual(result.stats.final_collision_buckets, 0) + self.assertEqual(result.stats.relocated_count, 3) + self.assertEqual(result.stats.unresolved_count, 0) + self.assertEqual(result.stats.max_final_bucket_size, 2) + self.assertEqual(result.stats.reassigned_count, 3) + self.assertEqual(result.stats.unassigned_count, 0) + self.assertEqual( + result.stats.to_output_dict(), + { + "total_items": 10, + "raw_collision_buckets": 2, + "final_collision_buckets": 0, + "reassigned_count": 3, + "unassigned_count": 0, + "max_final_bucket_size": 2, + }, + ) + + with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): + original_grouping = build_original_item_grouping(plan) + np.testing.assert_array_equal(original_grouping.sid_keys, [0, 1, 2, 4]) + np.testing.assert_array_equal(original_grouping.counts, [4, 1, 2, 3]) + np.testing.assert_array_equal(original_grouping.offsets, [0, 4, 5, 7, 10]) + np.testing.assert_array_equal( + original_grouping.row_order, [2, 0, 1, 3, 4, 6, 5, 9, 7, 8] + ) + np.testing.assert_array_equal( + original_grouping.representative_rows, [2, 4, 6, 9] + ) + + with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): + resolved_grouping = build_resolved_item_grouping(plan, result) + np.testing.assert_array_equal(resolved_grouping.sid_keys, [0, 1, 2, 3, 4, 5]) + np.testing.assert_array_equal(resolved_grouping.counts, [2, 2, 2, 1, 2, 1]) + np.testing.assert_array_equal( + resolved_grouping.row_order, [2, 0, 4, 1, 6, 5, 3, 9, 7, 8] + ) + + def test_error_fallback(self) -> None: + config = CollisionResolutionConfig((2,), 1, "error") + plan = prepare_collision_plan( + np.asarray([0, 1]), np.asarray([[0], [0]]), config + ) + + with self.assertRaisesRegex(RuntimeError, "first unresolved row indices: 1"): + resolve_sid_collisions(plan, np.asarray([[0]], dtype=np.int64)) + + def test_drop_fallback(self) -> None: + config = CollisionResolutionConfig((2,), 1, "drop") + plan = prepare_collision_plan( + np.asarray([0, 1]), np.asarray([[0], [0]]), config + ) + result = resolve_sid_collisions(plan, np.asarray([[0]], dtype=np.int64)) + + np.testing.assert_array_equal(result.resolved_last_codes, [0, 0]) + np.testing.assert_array_equal(result.slot_indices, [1, 2]) + np.testing.assert_array_equal(result.retained_mask, [True, False]) + np.testing.assert_array_equal(result.unresolved_rows, [1]) + self.assertEqual(result.stats.final_collision_buckets, 0) + self.assertEqual(result.stats.unresolved_count, 1) + with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 1): + original_grouping = build_original_item_grouping(plan) + resolved_grouping = build_resolved_item_grouping(plan, result) + np.testing.assert_array_equal(original_grouping.counts, [2]) + np.testing.assert_array_equal(original_grouping.row_order, [0, 1]) + np.testing.assert_array_equal(resolved_grouping.counts, [1]) + np.testing.assert_array_equal(resolved_grouping.row_order, [0]) + + def test_keep_original_fallback(self) -> None: + config = CollisionResolutionConfig((2,), 1, "keep_original") + plan = prepare_collision_plan( + np.asarray([0, 1]), np.asarray([[0], [0]]), config + ) + result = resolve_sid_collisions(plan, np.asarray([[0]], dtype=np.int64)) + + np.testing.assert_array_equal(result.resolved_last_codes, [0, 0]) + np.testing.assert_array_equal(result.slot_indices, [1, 2]) + self.assertIsNone(result.retained_mask) + np.testing.assert_array_equal(result.unresolved_rows, [1]) + self.assertEqual(result.stats.final_collision_buckets, 1) + self.assertEqual(result.stats.max_final_bucket_size, 2) + self.assertEqual(result.stats.unresolved_count, 1) + resolved_grouping = build_resolved_item_grouping(plan, result) + np.testing.assert_array_equal(resolved_grouping.sid_keys, [0]) + np.testing.assert_array_equal(resolved_grouping.counts, [2]) + np.testing.assert_array_equal(resolved_grouping.row_order, [0, 1]) + + def test_no_overflow(self) -> None: + config = CollisionResolutionConfig((2, 4), 2, "error") + codes = np.asarray([[0, 0], [0, 0], [1, 2]], dtype=np.int64) + plan = prepare_collision_plan(np.asarray([0, 1, 2]), codes, config) + result = resolve_sid_collisions(plan, np.empty((0, 3), dtype=np.int64)) + + np.testing.assert_array_equal(result.resolved_last_codes, [0, 0, 2]) + np.testing.assert_array_equal(result.slot_indices, [1, 2, 1]) + np.testing.assert_array_equal(result.unresolved_rows, []) + self.assertEqual(result.stats.raw_collision_buckets, 0) + self.assertEqual(result.stats.relocated_count, 0) + + def test_empty_groupings(self) -> None: + config = CollisionResolutionConfig((2, 4), 2, "error") + plan = prepare_collision_plan( + np.empty(0, dtype=np.int64), + np.empty((0, 2), dtype=np.int64), + config, + ) + result = resolve_sid_collisions(plan, np.empty((0, 0), dtype=np.int64)) + + for grouping in ( + build_original_item_grouping(plan), + build_resolved_item_grouping(plan, result), + ): + np.testing.assert_array_equal(grouping.sid_keys, []) + np.testing.assert_array_equal(grouping.counts, []) + np.testing.assert_array_equal(grouping.row_order, []) + np.testing.assert_array_equal(grouping.offsets, [0]) + np.testing.assert_array_equal(grouping.representative_rows, []) + + def test_random_candidate_golden_draws(self) -> None: + actual = generate_random_candidate_last_codes( + np.asarray([0, 1], dtype=np.int64), num=3, last_size=4 + ) + + np.testing.assert_array_equal(actual, [[1, 2, 0], [2, 0, 2]]) + + def test_candidate_matrix_must_be_two_dimensional(self) -> None: + config = CollisionResolutionConfig((2,), 1, "drop") + plan = prepare_collision_plan( + np.asarray([0, 1]), np.asarray([[0], [0]]), config + ) + + with self.assertRaisesRegex(ValueError, "must be 2-D"): + resolve_sid_collisions(plan, np.asarray([1], dtype=np.int64)) + + def test_candidate_matrix_must_align_with_overflow_rows(self) -> None: + config = CollisionResolutionConfig((2,), 1, "drop") + plan = prepare_collision_plan( + np.asarray([0, 1]), np.asarray([[0], [0]]), config + ) + + with self.assertRaisesRegex(ValueError, "row-aligned"): + resolve_sid_collisions(plan, np.empty((0, 1), dtype=np.int64)) + + def test_fixed_width_candidates_can_contend_for_free_slots(self) -> None: + config = CollisionResolutionConfig((6,), 1, "drop") + item_ids = np.asarray([0, 1, 2, 3, 4, 5], dtype=np.int64) + codes = np.asarray([[0], [0], [1], [3], [4], [4]], dtype=np.int64) + plan = prepare_collision_plan(item_ids, codes, config) + np.testing.assert_array_equal(plan.overflow_origin_last_codes, [0, 4]) + candidates = np.asarray( + [ + [0, 1, 2, 5], + [4, 2, 1, 3], + ], + dtype=np.int64, + ) + + result = resolve_sid_collisions(plan, candidates) + + np.testing.assert_array_equal(result.unresolved_rows, [5]) + self.assertEqual(result.resolved_last_codes[1], 2) + self.assertEqual(result.stats.relocated_count, 1) + self.assertEqual(result.stats.unresolved_count, 1) + + def test_resolved_grouping_uses_emitted_sid_for_out_of_range_candidate( + self, + ) -> None: + config = CollisionResolutionConfig((4,), 1, "drop") + item_ids = np.asarray([0, 1, 2], dtype=np.int64) + codes = np.asarray([[0], [0], [1]], dtype=np.int64) + plan = prepare_collision_plan(item_ids, codes, config) + + result = resolve_sid_collisions(plan, np.asarray([[5]], dtype=np.int64)) + + np.testing.assert_array_equal(result.resolved_last_codes, [0, 1, 1]) + np.testing.assert_array_equal(result.slot_indices, [1, 1, 1]) + self.assertFalse(result.slots_match_final_sid_keys) + np.testing.assert_array_equal(result.final_occupied_sid_keys, [0, 1]) + np.testing.assert_array_equal(result.final_bucket_counts, [1, 2]) + self.assertEqual(result.stats.final_collision_buckets, 0) + self.assertEqual(result.stats.max_final_bucket_size, 1) + + grouping = build_resolved_item_grouping(plan, result) + np.testing.assert_array_equal(grouping.sid_keys, [0, 1]) + np.testing.assert_array_equal(grouping.counts, [1, 2]) + np.testing.assert_array_equal(grouping.row_order, [0, 1, 2]) + + def test_resolution_can_skip_grouping_metadata(self) -> None: + config = CollisionResolutionConfig((4,), 1, "drop") + item_ids = np.asarray([0, 1, 2], dtype=np.int64) + codes = np.asarray([[0], [0], [1]], dtype=np.int64) + plan = prepare_collision_plan(item_ids, codes, config) + candidates = np.asarray([[5]], dtype=np.int64) + expected = resolve_sid_collisions(plan, candidates) + + with ( + mock.patch.object(collision_resolution.np, "fromiter") as fromiter, + mock.patch.object(collision_resolution.np, "argsort") as argsort, + mock.patch.object(collision_resolution.np, "unique") as unique, + ): + result = resolve_sid_collisions(plan, candidates, collect_grouping=False) + + fromiter.assert_not_called() + argsort.assert_not_called() + unique.assert_not_called() + np.testing.assert_array_equal( + result.resolved_last_codes, expected.resolved_last_codes + ) + np.testing.assert_array_equal(result.slot_indices, expected.slot_indices) + np.testing.assert_array_equal(result.retained_mask, expected.retained_mask) + np.testing.assert_array_equal(result.unresolved_rows, expected.unresolved_rows) + np.testing.assert_array_equal(result.final_occupied_sid_keys, []) + np.testing.assert_array_equal(result.final_bucket_counts, []) + self.assertFalse(result.slots_match_final_sid_keys) + self.assertFalse(result.grouping_collected) + self.assertEqual(result.stats, expected.stats) + + with self.assertRaisesRegex(RuntimeError, "metadata was not collected"): + build_resolved_item_grouping(plan, result) + + +if __name__ == "__main__": + unittest.main() From 31b2e6140c8517b00361dc7b634c29e25124fe73 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Mon, 13 Jul 2026 10:48:30 +0000 Subject: [PATCH 16/29] [refactor] simplify SID collision resolution --- tzrec/tools/sid/collision_prevention.py | 138 ++++++++-------- tzrec/tools/sid/collision_prevention_test.py | 48 +++++- tzrec/tools/sid/collision_resolution.py | 162 +++++++++---------- tzrec/tools/sid/collision_resolution_test.py | 34 ++-- 4 files changed, 212 insertions(+), 170 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index de0fb324..1fb1142a 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -27,7 +27,9 @@ import argparse import json import os +from contextlib import closing from dataclasses import dataclass +from functools import cached_property from typing import Dict, Iterable, List, Optional, Tuple, Union import numpy as np @@ -35,8 +37,8 @@ import pyarrow.compute as pc # Register the local reader/writer classes used through create_reader/writer. -import tzrec.datasets.csv_dataset # noqa: F401 import tzrec.datasets.parquet_dataset # noqa: F401 +from tzrec.datasets.csv_dataset import CsvWriter from tzrec.datasets.dataset import ( BaseReader, BaseWriter, @@ -60,8 +62,8 @@ _CODE_SEP = "|" _CAND_SEP = ";" -_WRITE_CHUNK = 20_000_000 -_CSV_GROUP_WRITE_CHUNK = 1_000_000 +_MAP_WRITE_ROWS = 1_000_000 +_GROUP_WRITE_ITEMS = 1_000_000 _ARROW_LIST_OFFSET_MAX = int(np.iinfo(np.int32).max) @@ -77,6 +79,28 @@ def _output_path_identity( return "odps", project, schema, table_name, partition_spec +class _ItemIdLookup: + """Map streamed item IDs to positions in a fixed requested-ID array.""" + + 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.") + + def match(self, item_ids: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Return matching source rows and requested-ID positions.""" + positions = np.searchsorted(self._sorted_ids, item_ids) + source_rows = np.flatnonzero(positions < self._sorted_ids.shape[0]) + if source_rows.size: + source_rows = source_rows[ + self._sorted_ids[positions[source_rows]] == item_ids[source_rows] + ] + return source_rows, self._sorted_to_requested[positions[source_rows]] + + @dataclass(frozen=True) class CollisionPreventionConfig: """Validated configuration for collision-prevention orchestration.""" @@ -171,7 +195,7 @@ def from_namespace(cls, args: argparse.Namespace) -> "CollisionPreventionConfig" odps_data_quota_name=args.odps_data_quota_name, ) - @property + @cached_property def resolution_config(self) -> CollisionResolutionConfig: """Return the pure-core configuration.""" return CollisionResolutionConfig( @@ -387,10 +411,7 @@ def _candidate_string_matrix(values: pa.Array) -> np.ndarray: def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarray: """Load fixed-width candidates aligned to ``overflow_item_ids``.""" item_count = overflow_item_ids.shape[0] - sorted_to_overflow = np.argsort(overflow_item_ids, kind="stable") - sorted_ids = overflow_item_ids[sorted_to_overflow] - if item_count > 1 and np.any(sorted_ids[1:] == sorted_ids[:-1]): - raise ValueError("overflow item IDs must be unique.") + item_id_lookup = _ItemIdLookup(overflow_item_ids) candidates: Optional[np.ndarray] = None seen = np.zeros(item_count, dtype=bool) @@ -400,17 +421,10 @@ def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarra if field not in batch: break batch_ids = batch[self._config.item_id_field].to_numpy(zero_copy_only=False) - positions = np.searchsorted(sorted_ids, batch_ids) - in_bounds = positions < item_count - source_rows = np.flatnonzero(in_bounds) - if source_rows.size: - source_rows = source_rows[ - sorted_ids[positions[source_rows]] == batch_ids[source_rows] - ] + source_rows, target_rows = item_id_lookup.match(batch_ids) if source_rows.size == 0: continue - target_rows = sorted_to_overflow[positions[source_rows]] 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())) @@ -497,11 +511,18 @@ def _codes_column(codes: np.ndarray, is_csv: bool) -> pa.Array: if layer_count < 1: raise ValueError("SID output must contain at least one layer.") if is_csv: - columns = [ - pc.cast(pa.array(codes[:, layer]), pa.string()) - for layer in range(layer_count) - ] - return pc.binary_join_element_wise(*columns, _CODE_SEP) + values = pc.cast(pa.array(codes.reshape(-1)), pa.string()) + offsets = pa.array( + np.arange( + 0, + (row_count + 1) * layer_count, + layer_count, + dtype=np.int64, + ) + ) + return pc.binary_join( + pa.LargeListArray.from_arrays(offsets, values), _CODE_SEP + ) if row_count > _ARROW_LIST_OFFSET_MAX // layer_count: raise ValueError("SID output exceeds Arrow list offset capacity.") values = pa.array(codes.reshape(-1)) @@ -522,18 +543,19 @@ def _write_map( result: CollisionResolutionResult, ) -> None: """Write the resolved map in chunks without a full final-code copy.""" - writer = self._make_writer(self._config.output_path) - is_csv = writer.__class__.__name__ == "CsvWriter" - retained_rows = ( - np.flatnonzero(result.retained_mask) - if result.retained_mask is not None - else None - ) - output_count = ( - retained_rows.shape[0] if retained_rows is not None else item_ids.shape[0] - ) - try: - write_chunk = _WRITE_CHUNK + with closing(self._make_writer(self._config.output_path)) as writer: + is_csv = isinstance(writer, CsvWriter) + retained_rows = ( + np.flatnonzero(result.retained_mask) + if result.retained_mask is not None + else None + ) + output_count = ( + retained_rows.shape[0] + if retained_rows is not None + else item_ids.shape[0] + ) + write_chunk = _MAP_WRITE_ROWS if not is_csv: write_chunk = min( write_chunk, @@ -562,8 +584,6 @@ def _write_map( ), } ) - finally: - writer.close() def _write_group_outputs( self, @@ -593,26 +613,21 @@ def _write_group_outputs( original_grouping, resolved_last_codes=None, ) - if plan.overflow_rows.size == 0: - self._write_sid_groups( - resolved_path, - item_ids, - origin_codes, - original_grouping, - resolved_last_codes=result.resolved_last_codes, - ) + if plan.overflow_rows.size: del original_grouping + grouping = build_resolved_item_grouping(plan, result) + resolved_last_codes = result.resolved_last_codes else: - del original_grouping - resolved_grouping = build_resolved_item_grouping(plan, result) - self._write_sid_groups( - resolved_path, - item_ids, - origin_codes, - resolved_grouping, - resolved_last_codes=result.resolved_last_codes, - ) - del resolved_grouping + grouping = original_grouping + resolved_last_codes = None + self._write_sid_groups( + resolved_path, + item_ids, + origin_codes, + grouping, + resolved_last_codes=resolved_last_codes, + ) + del grouping def _write_sid_groups( self, @@ -647,12 +662,8 @@ def _write_sid_groups( offsets = grouping.offsets if int(offsets[-1]) != grouping.row_order.shape[0]: raise RuntimeError("SID group counts do not match grouped row count.") - writer = self._make_writer(output_path) - is_csv = writer.__class__.__name__ == "CsvWriter" - child_chunk = ( - min(_WRITE_CHUNK, _CSV_GROUP_WRITE_CHUNK) if is_csv else _WRITE_CHUNK - ) - try: + with closing(self._make_writer(output_path)) as writer: + is_csv = isinstance(writer, CsvWriter) max_codebook_rows = ( group_count if is_csv @@ -662,7 +673,7 @@ def _write_sid_groups( raise ValueError("one SID row exceeds Arrow list offset capacity.") group_start = 0 while group_start < group_count: - child_limit = int(offsets[group_start]) + child_chunk + child_limit = int(offsets[group_start]) + _GROUP_WRITE_ITEMS group_end = int(np.searchsorted(offsets, child_limit, side="right") - 1) group_end = max(group_end, group_start + 1) group_end = min(group_end, group_start + max_codebook_rows) @@ -687,24 +698,19 @@ def _write_sid_groups( } ) group_start = group_end - finally: - writer.close() def _write_diagnostics(self, stats: CollisionResolutionStats) -> None: """Write legacy-compatible diagnostic column names.""" path = self._config.diagnostics_output_path if path is None: return - writer = self._make_writer(path) - try: + with closing(self._make_writer(path)) as writer: writer.write( { name: pa.array([value], type=pa.int64()) for name, value in stats.to_output_dict().items() } ) - finally: - writer.close() def build_parser() -> argparse.ArgumentParser: diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index c1621cde..7cc4b751 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -19,6 +19,7 @@ from collections import Counter from unittest import mock +import numpy as np import pyarrow as pa from pyarrow import csv, parquet @@ -376,6 +377,13 @@ def decisions(order): # ---- CSV backend ---- + def test_csv_code_encoder_joins_all_layers(self) -> None: + codes = np.asarray([[1, -2, 3], [4, 5, 6]], dtype=np.int64) + + encoded = CollisionRunner._codes_column(codes, is_csv=True) + + self.assertEqual(encoded.to_pylist(), ["1|-2|3", "4|5|6"]) + def test_csv_backend_within_band(self) -> None: inp = os.path.join(self.test_dir, "in.csv") out = os.path.join(self.test_dir, "out") @@ -633,12 +641,36 @@ def make_writer(output_path): self.assertEqual(_type_pa_to_table(columns["itemids"].type), "ARRAY") self.assertEqual(_type_pa_to_table(pa.list_(pa.string())), "ARRAY") + def test_writer_closes_when_write_fails(self) -> None: + class FailingWriter: + def __init__(self) -> None: + self.closed = False + + def write(self, output_dict) -> None: + raise RuntimeError("write failed") + + def close(self) -> None: + self.closed = True + + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0], [[0, 0]]) + writer = FailingWriter() + + with ( + mock.patch.object(CollisionRunner, "_make_writer", return_value=writer), + self.assertRaisesRegex(RuntimeError, "write failed"), + ): + self._run(inp, out, max_items_per_codebook=1) + + self.assertTrue(writer.closed) + def test_group_writer_chunks_by_item_count_without_splitting_groups(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") _parquet(inp, list(range(4)), [[0, 0], [0, 0], [0, 0], [0, 1]]) - with mock.patch("tzrec.tools.sid.collision_prevention._WRITE_CHUNK", 2): + with mock.patch("tzrec.tools.sid.collision_prevention._GROUP_WRITE_ITEMS", 2): self._run(inp, out, max_items_per_codebook=3) original_path, _ = self._group_paths(out) @@ -647,6 +679,17 @@ def test_group_writer_chunks_by_item_count_without_splitting_groups(self) -> Non groups = parquet.read_table(output_file).to_pydict() self.assertEqual([len(group) for group in groups["itemids"]], [3, 1]) + def test_map_writer_chunks_by_row_count(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(5)), [[0, code] for code in range(5)]) + + with mock.patch("tzrec.tools.sid.collision_prevention._MAP_WRITE_ROWS", 2): + self._run(inp, out, max_items_per_codebook=1) + + output_file = os.path.join(out, "part-0.parquet") + self.assertEqual(parquet.ParquetFile(output_file).metadata.num_row_groups, 3) + def test_writers_bound_codebook_list_offsets(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") @@ -656,7 +699,8 @@ def test_writers_bound_codebook_list_offsets(self) -> None: mock.patch( "tzrec.tools.sid.collision_prevention._ARROW_LIST_OFFSET_MAX", 4 ), - mock.patch("tzrec.tools.sid.collision_prevention._WRITE_CHUNK", 100), + mock.patch("tzrec.tools.sid.collision_prevention._MAP_WRITE_ROWS", 100), + mock.patch("tzrec.tools.sid.collision_prevention._GROUP_WRITE_ITEMS", 100), ): self._run(inp, out, max_items_per_codebook=1) diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py index cf79854b..30ca23da 100644 --- a/tzrec/tools/sid/collision_resolution.py +++ b/tzrec/tools/sid/collision_resolution.py @@ -12,6 +12,7 @@ """NumPy collision-resolution core for semantic IDs.""" from dataclasses import dataclass +from functools import cached_property from typing import Dict, Iterable, Optional import numpy as np @@ -123,7 +124,6 @@ class CollisionResolutionResult: unresolved_rows: np.ndarray final_occupied_sid_keys: np.ndarray final_bucket_counts: np.ndarray - slots_match_final_sid_keys: bool grouping_collected: bool stats: CollisionResolutionStats @@ -145,7 +145,7 @@ class CodebookItemGrouping: counts: np.ndarray row_order: np.ndarray - @property + @cached_property def offsets(self) -> np.ndarray: """Return CSR offsets into ``row_order`` for every SID bucket.""" offsets = np.empty(self.counts.shape[0] + 1, dtype=np.int64) @@ -153,7 +153,7 @@ def offsets(self) -> np.ndarray: np.cumsum(self.counts, dtype=np.int64, out=offsets[1:]) return offsets - @property + @cached_property def representative_rows(self) -> np.ndarray: """Return the first original row index in every occupied bucket.""" return self.row_order[self.offsets[:-1]] @@ -162,7 +162,7 @@ def representative_rows(self) -> np.ndarray: def _splitmix64(values: np.ndarray) -> np.ndarray: """Hash a uint64 array with the collision tool's fixed SplitMix64 seed.""" with np.errstate(over="ignore"): - mixed = values.astype(np.uint64) + np.uint64( + mixed = values.astype(np.uint64, copy=False) + np.uint64( (_SEED * _SPLITMIX_INCREMENT) & _MASK64 ) mixed = (mixed ^ (mixed >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) @@ -189,7 +189,7 @@ def stable_order_hash(item_ids: np.ndarray) -> np.ndarray: if item_ids.ndim != 1: raise ValueError(f"item_ids must be 1-D, got shape {item_ids.shape}.") if np.issubdtype(item_ids.dtype, np.integer): - base = item_ids.astype(np.uint64) + base = item_ids.astype(np.uint64, copy=False) else: import pandas as pd @@ -291,7 +291,6 @@ def prepare_collision_plan( f"{len(config.layer_sizes)}." ) - item_count = codes.shape[0] original_last_codes = codes[:, -1].astype(np.int64, copy=False) band_ids = _band_ids(codes, config.layer_sizes) order_hashes = stable_order_hash(item_ids) @@ -304,6 +303,7 @@ def prepare_collision_plan( ) = _within_bucket_rank(band_ids, original_last_codes, order_hashes) overflow_rows = sorted_rows[bucket_ranks[sorted_rows] >= config.capacity] + bucket_ranks += 1 last_size = config.layer_sizes[-1] occupied_sid_keys = ( band_ids[representative_rows] * last_size @@ -312,16 +312,16 @@ def prepare_collision_plan( overflow_band_bases = band_ids[overflow_rows] * last_size return CollisionPlan( - item_count=item_count, + item_count=codes.shape[0], original_last_codes=original_last_codes, origin_bucket_ids=origin_bucket_ids, - initial_slot_indices=bucket_ranks + 1, + initial_slot_indices=bucket_ranks, occupied_sid_keys=occupied_sid_keys, raw_bucket_counts=raw_bucket_counts, overflow_rows=overflow_rows, - overflow_item_ids=item_ids[overflow_rows].copy(), + overflow_item_ids=item_ids[overflow_rows], overflow_band_bases=overflow_band_bases, - overflow_origin_last_codes=original_last_codes[overflow_rows].copy(), + overflow_origin_last_codes=original_last_codes[overflow_rows], config=config, ) @@ -369,9 +369,8 @@ def resolve_sid_collisions( ) -> CollisionResolutionResult: """Greedily relocate overflow rows to their first free candidate slot. - Candidate rows must be aligned with ``plan.overflow_rows``. Candidate values - intentionally receive no range validation so negative and out-of-range - values retain the existing key and modulo behavior. + Candidate rows must be aligned with ``plan.overflow_rows`` and values must + be valid last-layer codebook indices. Args: plan: Grouping and overflow plan from :func:`prepare_collision_plan`. @@ -388,7 +387,8 @@ def resolve_sid_collisions( Raises: RuntimeError: If a row is unresolved under the ``error`` policy. TypeError: If candidate codes do not use an integer dtype. - ValueError: If candidates are not a row-aligned two-dimensional matrix. + ValueError: If candidates are not a row-aligned two-dimensional matrix + or contain out-of-range last-layer indices. """ candidates = np.asarray(candidate_last_codes) if candidates.ndim != 2: @@ -405,45 +405,79 @@ def resolve_sid_collisions( if not np.issubdtype(candidates.dtype, np.integer): raise TypeError("candidate_last_codes must use an integer dtype.") + last_size = plan.config.layer_sizes[-1] + if candidates.size and ( + int(candidates.min()) < 0 or int(candidates.max()) >= last_size + ): + raise ValueError( + f"candidate_last_codes must be in [0, {last_size}), got values " + "outside that range." + ) + + if overflow_count == 0: + final_occupied_sid_keys = ( + plan.occupied_sid_keys.copy() + if collect_grouping + else np.empty(0, dtype=np.int64) + ) + final_bucket_counts = ( + plan.raw_bucket_counts.copy() + if collect_grouping + else np.empty(0, dtype=np.int64) + ) + max_final_bucket_size = ( + int(plan.raw_bucket_counts.max()) if plan.raw_bucket_counts.size else 0 + ) + return CollisionResolutionResult( + resolved_last_codes=plan.original_last_codes, + slot_indices=plan.initial_slot_indices, + retained_mask=None, + unresolved_rows=np.empty(0, dtype=np.int64), + final_occupied_sid_keys=final_occupied_sid_keys, + final_bucket_counts=final_bucket_counts, + grouping_collected=collect_grouping, + stats=CollisionResolutionStats( + total_items=plan.item_count, + raw_collision_buckets=0, + final_collision_buckets=0, + relocated_count=0, + unresolved_count=0, + max_final_bucket_size=max_final_bucket_size, + ), + ) + capacity = plan.config.capacity initial_counts = np.minimum(plan.raw_bucket_counts, capacity) - slot_counts = dict(zip(plan.occupied_sid_keys.tolist(), initial_counts.tolist())) + slot_counts = dict(zip(map(int, plan.occupied_sid_keys), map(int, initial_counts))) resolved_last_codes = plan.original_last_codes.copy() slot_indices = plan.initial_slot_indices.copy() - relocated_rows = [] - relocated_last_codes = [] - relocated_indices = [] + relocated_count = 0 unresolved_rows = [] - slots_match_final_sid_keys = True get_slot_count = slot_counts.get - last_size = plan.config.layer_sizes[-1] - overflow_rows = plan.overflow_rows.tolist() - band_bases = plan.overflow_band_bases.tolist() - origin_last_codes = plan.overflow_origin_last_codes.tolist() - for position, row in enumerate(overflow_rows): - band_base = band_bases[position] - origin_last_code = origin_last_codes[position] - for candidate in candidates[position].tolist(): + for row_value, band_base_value, origin_value, candidate_row in zip( + plan.overflow_rows, + plan.overflow_band_bases, + plan.overflow_origin_last_codes, + candidates, + ): + row = int(row_value) + band_base = int(band_base_value) + origin_last_code = int(origin_value) + for candidate in candidate_row.tolist(): if candidate == origin_last_code: continue destination_key = band_base + candidate destination_count = get_slot_count(destination_key, 0) if destination_count < capacity: slot_counts[destination_key] = destination_count + 1 - relocated_rows.append(row) - relocated_last_codes.append(destination_key % last_size) - relocated_indices.append(destination_count + 1) - if candidate < 0 or candidate >= last_size: - slots_match_final_sid_keys = False + resolved_last_codes[row] = candidate + slot_indices[row] = destination_count + 1 + relocated_count += 1 break else: unresolved_rows.append(row) - if relocated_rows: - resolved_last_codes[relocated_rows] = relocated_last_codes - slot_indices[relocated_rows] = relocated_indices - retained_mask = None fallback_policy = plan.config.fallback_policy if unresolved_rows and fallback_policy == "error": @@ -465,31 +499,19 @@ def resolve_sid_collisions( final_occupied_sid_keys = np.empty(0, dtype=np.int64) final_bucket_counts = np.empty(0, dtype=np.int64) if collect_grouping: - occupancy_counts = np.fromiter(slot_counts.values(), dtype=np.int64) + occupancy_counts = np.fromiter( + slot_counts.values(), dtype=np.int64, count=len(slot_counts) + ) final_collision_buckets = int((occupancy_counts > capacity).sum()) max_final_bucket_size = ( int(occupancy_counts.max()) if occupancy_counts.size else 0 ) - if slots_match_final_sid_keys: - final_occupied_sid_keys = np.fromiter(slot_counts, dtype=np.int64) - occupancy_order = np.argsort(final_occupied_sid_keys, kind="stable") - final_occupied_sid_keys = final_occupied_sid_keys[occupancy_order] - final_bucket_counts = occupancy_counts[occupancy_order] - else: - retained_rows = ( - np.flatnonzero(retained_mask) - if retained_mask is not None - else np.arange(plan.item_count, dtype=np.int64) - ) - origin_keys = plan.occupied_sid_keys[plan.origin_bucket_ids[retained_rows]] - final_sid_keys = ( - origin_keys - - plan.original_last_codes[retained_rows] - + resolved_last_codes[retained_rows] - ) - final_occupied_sid_keys, final_bucket_counts = np.unique( - final_sid_keys, return_counts=True - ) + final_occupied_sid_keys = np.fromiter( + slot_counts, dtype=np.int64, count=len(slot_counts) + ) + occupancy_order = np.argsort(final_occupied_sid_keys, kind="stable") + final_occupied_sid_keys = final_occupied_sid_keys[occupancy_order] + final_bucket_counts = occupancy_counts[occupancy_order] else: final_collision_buckets = 0 max_final_bucket_size = 0 @@ -501,7 +523,7 @@ def resolve_sid_collisions( total_items=plan.item_count, raw_collision_buckets=int((plan.raw_bucket_counts > capacity).sum()), final_collision_buckets=final_collision_buckets, - relocated_count=len(relocated_rows), + relocated_count=relocated_count, unresolved_count=len(unresolved_rows), max_final_bucket_size=max_final_bucket_size, ) @@ -512,7 +534,6 @@ def resolve_sid_collisions( unresolved_rows=unresolved_array, final_occupied_sid_keys=final_occupied_sid_keys, final_bucket_counts=final_bucket_counts, - slots_match_final_sid_keys=slots_match_final_sid_keys, grouping_collected=collect_grouping, stats=stats, ) @@ -588,10 +609,7 @@ def build_resolved_item_grouping( ) -> CodebookItemGrouping: """Group retained rows by emitted SID and final one-based slot index. - Canonical in-range candidates use a linear scatter from final slot indices. - Malformed candidates can make internal occupancy keys differ from emitted SID - keys because the legacy resolver applies modulo only to the emitted last code. - That compatibility path groups by emitted keys and reported slots instead. + The grouping uses a linear scatter from final one-based slot indices. Args: plan: Grouping and overflow plan from :func:`prepare_collision_plan`. @@ -609,24 +627,6 @@ def build_resolved_item_grouping( "resolved item grouping metadata was not collected; call " "resolve_sid_collisions with collect_grouping=True." ) - if not result.slots_match_final_sid_keys: - rows = ( - np.flatnonzero(result.retained_mask) - if result.retained_mask is not None - else np.arange(plan.item_count, dtype=np.int64) - ) - origin_keys = plan.occupied_sid_keys[plan.origin_bucket_ids[rows]] - emitted_sid_keys = ( - origin_keys - - plan.original_last_codes[rows] - + result.resolved_last_codes[rows] - ) - row_order = np.lexsort((rows, result.slot_indices[rows], emitted_sid_keys)) - return CodebookItemGrouping( - sid_keys=result.final_occupied_sid_keys, - counts=result.final_bucket_counts, - row_order=rows[row_order], - ) def row_chunks() -> Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]]: for start in range(0, plan.item_count, _GROUPING_ROW_CHUNK): diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index 0a0cb730..91a93ec5 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -70,7 +70,6 @@ def test_golden_candidate_resolution(self) -> None: result.final_occupied_sid_keys, [0, 1, 2, 3, 4, 5] ) np.testing.assert_array_equal(result.final_bucket_counts, [2, 2, 2, 1, 2, 1]) - self.assertTrue(result.slots_match_final_sid_keys) self.assertTrue(result.grouping_collected) self.assertIsNone(result.retained_mask) self.assertEqual(result.stats.total_items, 10) @@ -166,11 +165,19 @@ def test_no_overflow(self) -> None: config = CollisionResolutionConfig((2, 4), 2, "error") codes = np.asarray([[0, 0], [0, 0], [1, 2]], dtype=np.int64) plan = prepare_collision_plan(np.asarray([0, 1, 2]), codes, config) - result = resolve_sid_collisions(plan, np.empty((0, 3), dtype=np.int64)) + with ( + mock.patch.object(collision_resolution.np, "fromiter") as fromiter, + mock.patch.object(collision_resolution.np, "argsort") as argsort, + ): + result = resolve_sid_collisions(plan, np.empty((0, 3), dtype=np.int64)) + fromiter.assert_not_called() + argsort.assert_not_called() np.testing.assert_array_equal(result.resolved_last_codes, [0, 0, 2]) np.testing.assert_array_equal(result.slot_indices, [1, 2, 1]) np.testing.assert_array_equal(result.unresolved_rows, []) + np.testing.assert_array_equal(result.final_occupied_sid_keys, [0, 6]) + np.testing.assert_array_equal(result.final_bucket_counts, [2, 1]) self.assertEqual(result.stats.raw_collision_buckets, 0) self.assertEqual(result.stats.relocated_count, 0) @@ -239,35 +246,21 @@ def test_fixed_width_candidates_can_contend_for_free_slots(self) -> None: self.assertEqual(result.stats.relocated_count, 1) self.assertEqual(result.stats.unresolved_count, 1) - def test_resolved_grouping_uses_emitted_sid_for_out_of_range_candidate( - self, - ) -> None: + def test_rejects_out_of_range_candidate(self) -> None: config = CollisionResolutionConfig((4,), 1, "drop") item_ids = np.asarray([0, 1, 2], dtype=np.int64) codes = np.asarray([[0], [0], [1]], dtype=np.int64) plan = prepare_collision_plan(item_ids, codes, config) - result = resolve_sid_collisions(plan, np.asarray([[5]], dtype=np.int64)) - - np.testing.assert_array_equal(result.resolved_last_codes, [0, 1, 1]) - np.testing.assert_array_equal(result.slot_indices, [1, 1, 1]) - self.assertFalse(result.slots_match_final_sid_keys) - np.testing.assert_array_equal(result.final_occupied_sid_keys, [0, 1]) - np.testing.assert_array_equal(result.final_bucket_counts, [1, 2]) - self.assertEqual(result.stats.final_collision_buckets, 0) - self.assertEqual(result.stats.max_final_bucket_size, 1) - - grouping = build_resolved_item_grouping(plan, result) - np.testing.assert_array_equal(grouping.sid_keys, [0, 1]) - np.testing.assert_array_equal(grouping.counts, [1, 2]) - np.testing.assert_array_equal(grouping.row_order, [0, 1, 2]) + with self.assertRaisesRegex(ValueError, r"must be in \[0, 4\)"): + resolve_sid_collisions(plan, np.asarray([[5]], dtype=np.int64)) def test_resolution_can_skip_grouping_metadata(self) -> None: config = CollisionResolutionConfig((4,), 1, "drop") item_ids = np.asarray([0, 1, 2], dtype=np.int64) codes = np.asarray([[0], [0], [1]], dtype=np.int64) plan = prepare_collision_plan(item_ids, codes, config) - candidates = np.asarray([[5]], dtype=np.int64) + candidates = np.asarray([[2]], dtype=np.int64) expected = resolve_sid_collisions(plan, candidates) with ( @@ -288,7 +281,6 @@ def test_resolution_can_skip_grouping_metadata(self) -> None: np.testing.assert_array_equal(result.unresolved_rows, expected.unresolved_rows) np.testing.assert_array_equal(result.final_occupied_sid_keys, []) np.testing.assert_array_equal(result.final_bucket_counts, []) - self.assertFalse(result.slots_match_final_sid_keys) self.assertFalse(result.grouping_collected) self.assertEqual(result.stats, expected.stats) From 749a1399cd4640cb24f3ff96b58f84b9af4b6019 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Tue, 14 Jul 2026 09:52:26 +0000 Subject: [PATCH 17/29] [doc] add SID append refactoring plan --- docs/sid_collision_append_refactoring_plan.md | 156 ++++++++++++++++++ ...id_collision_append_refactoring_plan_zh.md | 156 ++++++++++++++++++ 2 files changed, 312 insertions(+) create mode 100644 docs/sid_collision_append_refactoring_plan.md create mode 100644 docs/sid_collision_append_refactoring_plan_zh.md diff --git a/docs/sid_collision_append_refactoring_plan.md b/docs/sid_collision_append_refactoring_plan.md new file mode 100644 index 00000000..00e9ea03 --- /dev/null +++ b/docs/sid_collision_append_refactoring_plan.md @@ -0,0 +1,156 @@ +# SID Collision Resolution: Fresh and Append Refactoring Plan + +> Status: design proposal; no implementation changes are included. + +## 1. Scope and terminology + +The refactoring should expose two explicit modes: + +- **Fresh** (the requested “Pure Addition” scenario): read one prediction file and run exactly the current collision-resolution algorithm. With no prior state, assignments, indexes, fallback behavior, statistics, and output ordering must remain unchanged. +- **Append**: read a new prediction file plus an existing `original_groups` snapshot and `resolved_groups` snapshot. Existing assignments are immutable; only new items are resolved and appended. + +`original_groups` records the historical mapping from each original SID to its item IDs. `resolved_groups` records the final SID buckets. The position of an item in a resolved bucket is semantic: position `k` corresponds to the current 1-based map index `k`. + +The inspected workspace artifacts are `experiments/sid_collision/sample.parquet` and `sid_collision/{original_groups,resolved_groups}`. The literal `/mnt/data/...` path was not present in this environment. The sample contains 500 unique items and is exactly the source of the existing snapshots, so using it as an append batch must fail duplicate validation rather than append another 500 copies. + +## 2. Behavioral contract + +| Property | Fresh | Append | +| ------------------------------- | --------------------- | -------------------------------- | +| Prior group inputs | Forbidden | Both snapshots required | +| Existing final SIDs and indexes | Not applicable | Never changed | +| Capacity seed | Empty | Counts from `resolved_groups` | +| Original membership | Current input | Historical groups plus new input | +| Candidate resolution | All current overflows | New overflows only | +| Map output | Current input | New-input delta only | +| Group outputs | Complete snapshots | New complete snapshots | + +Append is intentionally **not** equivalent to rerunning Fresh over the union. Old candidate lists are unavailable, historical assignments must remain stable, and existing items receive priority over new items. Consequently, append results are history-dependent. + +Default duplicate policy is `error`: item IDs must be unique in the new batch and disjoint from historical `original_groups`. Neither silent skipping nor upsert is safe because the grouped state cannot prove that a repeated item has identical SID and candidate data. + +All append outputs must use paths distinct from their input snapshots. `rate_only` should perform the same validation and planning and report new and combined statistics, but should not publish group or map data. + +## 3. Stable persisted identity + +The current dense band IDs are produced from prefixes present in one run. They are suitable for in-memory grouping but cannot identify a persisted bucket: for example, prefix `B` may be band 1 in one batch and band 0 in another. + +Introduce a `SidKeyCodec` that validates every code against `layer_sizes` and converts a full SID to a stable, checked mixed-radix `int64` key. Because the last code has unit stride, the corresponding band key can be derived independently of the current batch. Reject configurations whose key space cannot fit in `int64`; a structured multi-column key is the future fallback if such configurations must be supported. + +Persist a versioned state manifest. The group files alone cannot safely recover capacity or layer sizes—the inspected sample generator uses `(64, 64, 64)`, while the local invocation script specifies `(256, 256, 256)`. The manifest should contain: + +- state format, SID-key codec, ordering/hash, and assignment-algorithm versions; +- `layer_sizes`, capacity, strategy, fallback policy, and random-count settings; +- Arrow item-ID type/schema fingerprint; +- original and resolved item counts, parent state version, and map scope (`delta`); +- artifact locations/formats and a completion marker. + +Legacy snapshots should be bootstrapped once using explicit configuration, fully validated, and then accompanied by a manifest without changing any assignment. + +## 4. Append planning algorithm + +1. **Validate state.** Load the manifest, verify configuration and item-ID type compatibility, reject mixed CSV/Parquet directories, and ensure input and output paths do not overlap. +1. **Load compact occupancy.** Read `resolved_groups`, validate unique sorted SID keys and bucket contents, and retain only canonical SID keys and counts for planning. These counts—not `original_groups`—represent occupied capacity. +1. **Validate new IDs.** Check uniqueness in the new input. Sort its IDs once, then stream historical `original_groups` item lists through lookup against that sorted array. This avoids a Python set containing every historical item. +1. **Rank new origins.** Preserve the current deterministic hash ranking within each new original-SID bucket. +1. **Admit at origins.** For origin bucket `s`, compute `free(s) = max(capacity - existing_resolved_count(s), 0)`. Keep the first `free(s)` new items there; their absolute indexes start at `existing_resolved_count(s) + 1`. Mark the remainder as overflow. +1. **Reserve all origins first.** Add every newly retained origin item to occupancy before processing any candidate. This preserves the current Fresh rule that direct origin ownership has priority over candidate relocation. +1. **Resolve new overflow.** In deterministic order, apply the current first-fit or random strategy only to new overflow rows. Candidate occupancy starts from historical resolved counts plus all newly retained origin items. Candidates remain confined to the same canonical prefix band. +1. **Apply fallback.** `error` aborts before publication; `drop` omits unresolved new items from resolved groups and the map but retains their original membership; `keep_original` appends them to their origin even when this exceeds capacity. +1. **Build the delta.** Record each new item’s final SID and absolute 1-based index. Existing map rows are neither reconstructed nor rewritten. + +With empty base occupancy, the generalized planner must produce byte-for-byte-equivalent logical results to the current Fresh implementation. + +## 5. Snapshot merge and output rules + +The first implementation should publish complete replacement snapshots: + +- **Original snapshot:** one row per original SID; preserve the old item list as a prefix and append new IDs in the current deterministic within-bucket order. +- **Resolved snapshot:** one row per final SID; preserve the old list exactly as a prefix and append new IDs in assigned absolute-index order. This invariant makes every emitted map index verifiable by direct lookup. +- **Map:** emit rows for the new input only, matching the existing meaning of `output_path` as the mapping for the input being processed. + +Use TorchEasyRec’s existing readers and writers, including the CSV writer. A small state-specific adapter is justified for manifest validation, normalized CSV/Parquet group decoding, streaming merge, and publication; a second generic I/O framework is not. + +CSV represents an entire `itemids` group as one JSON-like field, so a hot bucket can create a line larger than the reader block size. Parquet or ODPS should be preferred for large state. Write chunks bound table/CSV batch memory, but they cannot split one oversized group without changing the schema. + +Writers currently replace outputs independently, so do not update state in place. Write all artifacts under a new version, close and validate them, and publish the manifest/current-version pointer last. Include the parent version and use a lock or compare-and-swap check to prevent concurrent append jobs from losing updates. + +If repeated full snapshots become too expensive at 100 million items, add an explicit base-plus-ordered-deltas format with periodic compaction as a later phase. Do not emit duplicate codebook rows under the current one-row-per-bucket snapshot contract. + +## 6. ODPS append support check + +**Conclusion (verified 2026-07-14): MaxCompute supports physical row append, but the current TorchEasyRec writer does not expose it, and direct physical append is incorrect for the grouped-state schema.** + +| Layer | Append capability | Consequence for this design | +| ------------------------- | --------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | +| MaxCompute SQL | `INSERT INTO` appends to a standard table or static partition; `INSERT OVERWRITE` replaces it | The platform has append capability, subject to table restrictions | +| PyODPS / Storage API | Non-overwrite writes can append; `overwrite=True` requests replacement | The underlying client can represent both modes | +| TorchEasyRec `OdpsWriter` | Hardcodes `TableBatchWriteRequest(..., overwrite=True)` | Every writer run replaces the target table or partition | +| SID group snapshot | Requires one row per `codebook` with the complete `itemids` list | Appending a touched bucket would create a duplicate row, not update its list | + +The platform behavior is documented by [MaxCompute `INSERT INTO` / `INSERT OVERWRITE`](https://www.alibabacloud.com/help/en/maxcompute/user-guide/insert-or-update-data-into-a-table-or-a-static-partition), [TableTunnel overwrite semantics](https://www.alibabacloud.com/help/en/maxcompute/user-guide/tabletunnel), [PyODPS table writes](https://pyodps.readthedocs.io/en/latest/base-tables.html#write-data-to-tables), and [PyODPS Storage API writes](https://pyodps.readthedocs.io/en/stable/base-storage-api-v2.html#writing-data). MaxCompute also states that `INSERT INTO` is unsupported for clustered tables and that concurrent inserts to the same target are not protected by table locking. Its [ACID documentation](https://www.alibabacloud.com/help/en/maxcompute/product-overview/acid-semantics) does not provide a transaction spanning the separate SID output tables. + +Repository inspection confirms that `tzrec/datasets/odps_dataset.py:762-764` always supplies `overwrite=True`; neither `OdpsWriter.__init__` nor the collision workflow passes an append option. Repeated `write()` calls accumulate batches only inside that one overwrite session. Existing ODPS tests validate a single writer session, not retention across two writer runs. The repository pins PyODPS 0.12.5.1, whose [`TableBatchWriteRequest`](https://raw.githubusercontent.com/aliyun/aliyun-odps-python-sdk/v0.12.5.1/odps/apis/storage_api/storage_api.py) accepts an overwrite flag, but the explicit `True` in TorchEasyRec prevents append. + +For example, physical ODPS append changes `(A, [1, 2])` plus `(A, [3])` into two rows for codebook `A`; it does not produce `(A, [1, 2, 3])`. This violates the current snapshot invariant even though the database append itself succeeds. + +Therefore, the initial ODPS implementation should keep **logical Append** separate from **physical row append**: + +1. Read the parent `original_groups` and `resolved_groups` partitions. +1. Merge the old lists with the new delta in the application. +1. Write complete group snapshots to new immutable `state_version=` partitions using the existing overwrite writer. +1. Validate all artifacts and publish the completed-version manifest last. +1. Enforce a single writer or parent-version check because the four outputs do not share one transaction and MaxCompute provides no insert lock. + +Do not add a generic `append=True` switch to `OdpsWriter` merely for this workflow. Physical append becomes appropriate only if a future, explicitly different schema stores ordered bucket fragments such as `(state_version, sequence, codebook, itemids_delta)` and readers fold those deltas before use. + +No ODPS credentials or CI project are configured in this workspace, so this verification combines official documentation, the pinned and installed PyODPS request models, and repository source inspection rather than a live remote write. Before release, add an ODPS integration gate using a temporary non-clustered table and static partition with `ARRAY` item IDs: write `(A, [1, 2])`, append `(B, [3])` and `(A, [4])` with `overwrite=False`, confirm that old rows survive and `A` is duplicated rather than merged, verify `overwrite=True` replacement, and inject a failed version publication to prove the parent snapshot remains active. + +## 7. Proposed modularization + +- `collision_resolution.py`: pure array-based Fresh/Append planning; no storage concerns. +- `SidKeyCodec`: stable SID encoding, decoding, range checks, and band identity. +- `BucketOccupancy`: compact historical counts plus a mutable overlay for touched buckets. +- `CollisionPlan` / `CollisionResult`: distinguish base counts, new admissions, new final codes/indexes, unresolved rows, and combined counts. +- A state module such as `collision_state.py`: manifest, group readers, integrity checks, snapshot merge, and atomic publication using existing repository I/O. +- `collision_prevention.py`: CLI parsing, mode validation, orchestration, statistics, and writer selection. + +Keep Arrow arrays at the I/O boundary and NumPy arrays in the CPU planner. Converting everything to tensors would add copies and device concerns without benefiting this sort/group/merge workload. Avoid a global Python dictionary or item-ID set at 100-million scale; use sorted arrays, vectorized search, streaming historical IDs, and a small overlay for buckets touched by the new batch. + +The existing grouping helper must not scatter new rows directly into arrays sized by combined counts. Append indexes are absolute, while a delta buffer is relative; use `delta_position = absolute_index - base_bucket_count` when constructing new group fragments. + +## 8. Result-affecting decisions + +These rules must be documented and treated as compatibility contracts: + +- Existing items always win capacity and never move. +- Existing resolved item order—and therefore indexes—is immutable. +- All new origin admissions are reserved before any new candidate assignment. +- New-item ordering remains deterministic and input-order invariant. +- The capacity, `layer_sizes`, item-ID type, hash version, and assignment version must match the parent state. +- `drop` permits resolved state to contain fewer items than original state; `keep_original` permits over-capacity buckets. +- The append map is a delta. A full historical map would require the old map as an additional state input or a costly reconstruction. + +Temporary full-array/debug printing currently present in the working tree must be removed before scale tests because it would dominate runtime and memory, but that cleanup is outside this design-only change. + +## 9. Verification plan + +1. Golden/differential test: Fresh with empty state exactly matches the current implementation for first-fit, random, all fallbacks, and `rate_only`. +1. Stable-key test: historical prefixes `A,B`, followed by a new batch containing only `B`, catches batch-local band IDs. +1. Append cases: empty, partially filled, full, and already-overfull origins; occupied candidate targets; new-only buckets; and competing candidates. +1. Immutability checks: every old final SID, item-list prefix, and index remains unchanged; every new map index points to the merged resolved list. +1. Duplicate checks: within-new and new-versus-history; appending the supplied sample to the supplied snapshots must reject all 500 overlaps. +1. Compatibility checks: capacity, layer sizes, item-ID type, manifest version, schema, format, and path mismatch. +1. Fallback and format matrix: `error`, `drop`, `keep_original`; CSV and Parquet round trips; mixed-format directories rejected. +1. Operational tests: two sequential appends, input-order permutation, failed publication, concurrent-parent mismatch, and recovery of the prior active version. +1. Scale benchmarks: 100-million historical items, many buckets, one hot bucket, bounded planner memory, snapshot merge throughput, and CSV line-size limits. +1. ODPS integration: verify platform append in a temporary standard table/static partition, current `OdpsWriter` replacement across separate runs, version-partition isolation, and rejection of duplicate-codebook snapshot rows. + +## 10. Implementation sequence + +1. Freeze Fresh behavior with golden tests and define the manifest/key compatibility contract. +1. Add stable SID keys, state data classes, and legacy-state validation/bootstrap. +1. Generalize the pure planner to accept base occupancy; prove empty-base equivalence. +1. Add delta grouping and streaming full-snapshot merge through existing I/O. +1. Add explicit `fresh|append` CLI validation, versioned publication, statistics, and failure handling. +1. Complete the correctness/format/scale matrix, then consider delta files plus compaction only if snapshot rewrite cost warrants it. diff --git a/docs/sid_collision_append_refactoring_plan_zh.md b/docs/sid_collision_append_refactoring_plan_zh.md new file mode 100644 index 00000000..e47cdfc4 --- /dev/null +++ b/docs/sid_collision_append_refactoring_plan_zh.md @@ -0,0 +1,156 @@ +# SID 碰撞处理：Fresh 与 Append 重构方案 + +> 状态：设计方案；不包含代码实现变更。 + +## 1. 范围与术语 + +重构后应显式支持两种模式： + +- **Fresh**（需求中的“纯新增”）：只读取一个预测文件，完整执行当前碰撞处理算法。没有历史状态时，分配结果、索引、fallback 行为、统计和输出顺序都必须与当前实现保持一致。 +- **Append**：读取一个新预测文件，以及已有的 `original_groups` 和 `resolved_groups` 快照。历史分配不可变，只处理并追加新 item。 + +`original_groups` 保存“原始 SID -> 历史 item IDs”；`resolved_groups` 保存“最终 SID -> item IDs”。item 在 resolved bucket 中的位置具有业务含义：第 `k` 个位置对应当前 map 中从 1 开始的索引 `k`。 + +本次实际检查的是工作区中的 `experiments/sid_collision/sample.parquet` 和 `sid_collision/{original_groups,resolved_groups}`；当前环境中不存在字面量 `/mnt/data/...` 路径。sample 包含 500 个唯一 item，而且正是现有快照的数据来源。因此，若把该 sample 作为 append 新批次，正确行为应是重复 ID 校验失败，而不是再次追加 500 份数据。 + +## 2. 行为契约 + +| 属性 | Fresh | Append | +| ------------------- | --------------------- | ----------------------------- | +| 历史 group 输入 | 禁止 | 必须同时提供两个快照 | +| 历史最终 SID 与索引 | 不适用 | 永不改变 | +| 容量初始值 | 空 | 来自 `resolved_groups` 的计数 | +| 原始归属 | 当前输入 | 历史 groups 加新输入 | +| Candidate 处理范围 | 当前批次全部 overflow | 只处理新 item 的 overflow | +| Map 输出 | 当前输入 | 只输出新批次增量 | +| Group 输出 | 完整快照 | 新的完整快照 | + +Append 有意地**不等价于**对“历史 + 新数据”重新执行 Fresh。原因是历史 candidate 列表不可用、历史分配必须稳定，而且旧 item 的优先级高于新 item。因此，append 结果必然与批次历史有关。 + +重复 ID 的默认策略应为 `error`：新批次内部必须唯一，并且不得与历史 `original_groups` 相交。不能默认静默跳过或 upsert，因为 group 状态无法证明重复 item 的 SID 和 candidate 数据完全一致。 + +Append 输出路径必须与两个输入快照不同。`rate_only` 仍应完成相同的校验和规划，并报告新增与合并后的统计，但不发布 group 或 map 数据。 + +## 3. 稳定的持久化标识 + +当前 dense band ID 由“本批次实际出现的 prefix”生成，适合单次内存分组，却不能标识持久化 bucket。例如，prefix `B` 在一个批次中可能是 band 1，在另一个只有 `B` 的批次中却变成 band 0。 + +应引入 `SidKeyCodec`：先按 `layer_sizes` 校验每一层 code，再把完整 SID 编码为稳定且有溢出检查的 mixed-radix `int64` key。最后一层的步长为 1，因此可以得到不依赖当前批次的 band key。如果整个 key 空间无法放入 `int64`，应拒绝配置；若未来确有需要，再使用结构化多列 key 作为替代方案。 + +还应持久化带版本的 state manifest。仅从 group 文件无法安全恢复 capacity 或 layer sizes：本次检查到 sample 生成器使用 `(64, 64, 64)`，而本地调用脚本配置为 `(256, 256, 256)`。manifest 至少应包含： + +- state 格式、SID key codec、排序/hash 和分配算法版本； +- `layer_sizes`、capacity、strategy、fallback policy 和 random-count 配置； +- Arrow item-ID 类型或 schema fingerprint； +- original/resolved item 数、父 state 版本，以及 map 范围（`delta`）； +- 各 artifact 的位置、格式和完成标记。 + +旧快照需要进行一次 bootstrap：显式提供配置，完整校验后生成 manifest，但不改变任何历史分配。 + +## 4. Append 规划算法 + +1. **校验 state。** 读取 manifest，检查配置与 item-ID 类型兼容性；拒绝 CSV/Parquet 混合目录；确认输入与输出路径不重叠。 +1. **读取紧凑 occupancy。** 从 `resolved_groups` 读取并校验唯一、已排序的 SID key 和 bucket 内容；规划阶段只保留 canonical SID key 与计数。占用容量必须来自 `resolved_groups`，而不是 `original_groups`。 +1. **校验新 ID。** 检查新批次内部唯一性；对新 ID 排序一次，再流式扫描历史 `original_groups.itemids` 并查找交集，避免为全部历史 item 创建 Python set。 +1. **对新 origin 排序。** 在每个新 original-SID bucket 内继续使用当前确定性的 hash rank。 +1. **在 origin 接纳新 item。** 对 origin bucket `s`，计算 `free(s) = max(capacity - existing_resolved_count(s), 0)`。前 `free(s)` 个新 item 留在 origin，其绝对索引从 `existing_resolved_count(s) + 1` 开始；其余标记为 overflow。 +1. **先预占全部 origin。** 处理任何 candidate 之前，先把所有被 origin 接纳的新 item 加入 occupancy。这样才能保留当前 Fresh 中“直接归属优先于 candidate 迁入”的规则。 +1. **处理新 overflow。** 按确定性顺序，只对新 overflow 执行当前 first-fit 或 random 策略。candidate occupancy 的初值为历史 resolved 计数加上全部新 origin 接纳项；candidate 仍须限制在同一个 canonical prefix band 中。 +1. **应用 fallback。** `error` 在发布前终止；`drop` 不把 unresolved 新 item 写入 resolved groups 和 map，但仍保留其 original membership；`keep_original` 把它追加回 origin，即使 bucket 因此超过 capacity。 +1. **构建增量结果。** 保存每个新 item 的最终 SID 和绝对的、从 1 开始的索引；不重建或重写旧 map 行。 + +当 base occupancy 为空时，通用 planner 的逻辑结果必须与当前 Fresh 实现完全一致。 + +## 5. 快照合并与输出规则 + +第一阶段应输出完整的新快照： + +- **Original 快照：** 每个 original SID 一行；旧 item 列表必须保持为前缀，新 ID 按当前 bucket 内确定性顺序追加。 +- **Resolved 快照：** 每个 final SID 一行；旧列表必须原样保持为前缀，新 ID 按已分配的绝对索引顺序追加。这样每个新 map index 都可以直接在合并后的列表中验证。 +- **Map：** 只输出当前新输入的记录，保持 `output_path` 表示“本次输入的映射”这一现有语义。 + +继续复用 TorchEasyRec 已有 reader/writer，包括 CSV writer。可以增加一个小型、面向 state 的适配模块，负责 manifest 校验、CSV/Parquet group 统一解码、流式合并和发布；没有必要再建立一套通用 I/O 框架。 + +CSV 会把整个 `itemids` group 放进一个类似 JSON 的字段，因此热点 bucket 可能产生超过 reader block size 的单行。大规模 state 应优先选择 Parquet 或 ODPS。写入 chunk 可以控制 table/CSV 批次内存，但在不修改 schema 的前提下，无法拆分单个超大 group。 + +当前 writer 会彼此独立地覆盖输出，因此不能原地更新 state。应先把所有 artifact 写入一个新版本，关闭并校验后，最后发布 manifest 或 current-version 指针。manifest 中记录父版本，并使用锁或 compare-and-swap 校验，避免并发 append 互相覆盖。 + +如果 1 亿 item 场景下频繁重写完整快照成本过高，可以在后续阶段引入显式的“base + 有序 delta + 定期 compaction”格式。当前“一 bucket 一行”的快照契约下，不能直接输出重复 codebook 行。 + +## 6. ODPS Append 能力检查 + +**结论（验证日期：2026-07-14）：MaxCompute 支持物理行追加，但 TorchEasyRec 当前 writer 没有暴露该能力；而且对于现有 grouped-state schema，直接物理追加是错误的。** + +| 层次 | Append 能力 | 对本设计的影响 | +| ------------------------- | --------------------------------------------------------------------- | ---------------------------------------------------- | +| MaxCompute SQL | `INSERT INTO` 可追加到普通表或静态分区；`INSERT OVERWRITE` 会替换目标 | 平台具备追加能力，但受表类型等限制 | +| PyODPS / Storage API | non-overwrite write 可以追加；`overwrite=True` 表示替换 | 底层 client 能表达两种模式 | +| TorchEasyRec `OdpsWriter` | 固定使用 `TableBatchWriteRequest(..., overwrite=True)` | 每次 writer 运行都会替换目标表或分区 | +| SID group 快照 | 每个 `codebook` 必须只有一行，且 `itemids` 是完整列表 | 对已存在 bucket 追加行只会产生重复行，不会更新原列表 | + +平台能力可参见 [MaxCompute `INSERT INTO` / `INSERT OVERWRITE`](https://www.alibabacloud.com/help/en/maxcompute/user-guide/insert-or-update-data-into-a-table-or-a-static-partition)、[TableTunnel overwrite 语义](https://www.alibabacloud.com/help/en/maxcompute/user-guide/tabletunnel)、[PyODPS table write](https://pyodps.readthedocs.io/en/latest/base-tables.html#write-data-to-tables) 和 [PyODPS Storage API write](https://pyodps.readthedocs.io/en/stable/base-storage-api-v2.html#writing-data)。MaxCompute 还明确说明：`INSERT INTO` 不支持 clustered table，并且对同一目标的并发 INSERT 没有 table lock 保护；其 [ACID 文档](https://www.alibabacloud.com/help/en/maxcompute/product-overview/acid-semantics) 也没有提供跨多个 SID 输出表的统一事务。 + +仓库检查显示，`tzrec/datasets/odps_dataset.py:762-764` 始终传入 `overwrite=True`；`OdpsWriter.__init__` 和 collision 流程都没有 append 参数。同一个 writer 内多次调用 `write()` 只是在同一个 overwrite session 内累积 batch。现有 ODPS 测试只验证一次 writer session，没有验证两次独立运行后是否保留旧数据。仓库固定使用 PyODPS 0.12.5.1；该版本的 [`TableBatchWriteRequest`](https://raw.githubusercontent.com/aliyun/aliyun-odps-python-sdk/v0.12.5.1/odps/apis/storage_api/storage_api.py) 接受 overwrite flag，但 TorchEasyRec 显式传入的 `True` 阻止了 append。 + +例如，ODPS 物理 append 会把 `(A, [1, 2])` 加 `(A, [3])` 变成两行 codebook `A`，而不是 `(A, [1, 2, 3])`。因此，即使数据库追加成功，结果仍违反当前 snapshot invariant。 + +因此，第一阶段 ODPS 实现必须区分**逻辑 Append**和**物理行 append**： + +1. 读取父版本的 `original_groups` 与 `resolved_groups` 分区。 +1. 在应用层合并旧列表和新 delta。 +1. 使用现有 overwrite writer，把完整 group 快照写到新的不可变 `state_version=` 分区。 +1. 校验全部 artifact 后，最后发布 completed-version manifest。 +1. 使用 single-writer 或父版本校验，因为四类输出不共享一个事务，而且 MaxCompute 没有 insert lock。 + +不要仅为本流程给通用 `OdpsWriter` 增加简单的 `append=True` 开关。只有未来显式引入不同的 delta schema，例如 `(state_version, sequence, codebook, itemids_delta)`，并要求 reader 在使用前折叠这些 delta 时，物理 append 才适合。 + +当前工作区没有配置 ODPS credential 或 CI project，因此本次验证基于官方文档、固定版本和已安装版本的 PyODPS request model，以及仓库源码，而不是远程 live write。发布前应增加 ODPS integration gate：使用带 `ARRAY` item IDs 的临时非 clustered 普通表和静态分区，先写 `(A, [1, 2])`，再用 `overwrite=False` 追加 `(B, [3])` 与 `(A, [4])`；确认旧行仍存在且 `A` 产生重复行而不是合并，验证 `overwrite=True` 会替换目标，并注入版本发布失败以确认父快照仍为 active。 + +## 7. 建议的模块划分 + +- `collision_resolution.py`：纯数组的 Fresh/Append 规划逻辑，不涉及存储。 +- `SidKeyCodec`：稳定 SID 编解码、范围检查和 band 标识。 +- `BucketOccupancy`：紧凑的历史计数，以及仅包含本次触达 bucket 的可变 overlay。 +- `CollisionPlan` / `CollisionResult`：明确区分 base counts、新接纳项、新 item 的最终 code/index、unresolved 行和合并后计数。 +- 新的 state 模块（如 `collision_state.py`）：manifest、group reader、完整性检查、快照合并，以及基于现有仓库 I/O 的原子发布。 +- `collision_prevention.py`：CLI 参数、mode 校验、流程编排、统计和 writer 选择。 + +I/O 边界继续使用 Arrow，CPU planner 使用 NumPy。把全部数据统一成 tensor 会增加复制和 device 管理，却不会改善这种排序、分组和合并工作负载。1 亿数据规模下，应避免全量 Python dict 或 item-ID set，改用排序数组、向量化查找、流式历史 ID，以及只保存新批次触达 bucket 的小型 overlay。 + +现有 grouping helper 不能直接把新行 scatter 到以“合并后 counts”为大小的数组中。Append index 是绝对位置，而 delta buffer 使用相对位置；构造新 group fragment 时应使用 `delta_position = absolute_index - base_bucket_count`。 + +## 8. 会影响结果的关键决策 + +以下规则需要写入兼容性契约： + +- 历史 item 永远优先占用容量，且绝不迁移。 +- 历史 resolved item 顺序以及对应 index 永远不变。 +- 必须先预占所有新 origin 接纳项，再执行任何新 candidate 分配。 +- 新 item 排序保持确定性，并且不受输入行顺序影响。 +- capacity、`layer_sizes`、item-ID 类型、hash 版本和分配算法版本必须与父 state 一致。 +- `drop` 允许 resolved state 的 item 数少于 original state；`keep_original` 允许 bucket 超过 capacity。 +- Append map 是增量。若需要完整历史 map，就必须把旧 map 作为额外 state 输入，或承担昂贵的重建成本。 + +当前工作树中存在临时的全数组/debug 输出；在大规模测试前必须删除，否则输出本身会主导耗时和内存。但该清理不属于本次仅设计文档的变更范围。 + +## 9. 验证计划 + +1. Golden/differential 测试：空 state 下的 Fresh 在 first-fit、random、全部 fallback 和 `rate_only` 模式中都与当前实现一致。 +1. 稳定 key 测试：历史 prefix 为 `A,B`，新批次只有 `B`，用于捕获 batch-local band ID 问题。 +1. Append 场景：空、部分占用、满和已超容量 origin；已占用 candidate target；仅新数据 bucket；candidate 竞争。 +1. 不变性检查：所有旧 final SID、item-list 前缀和 index 保持不变；每个新 map index 都指向合并后的 resolved list。 +1. 重复检查：新批次内部、新批次与历史之间；把提供的 sample append 到提供的快照时，应检测并拒绝 500 个重复 ID。 +1. 兼容性检查：capacity、layer sizes、item-ID 类型、manifest 版本、schema、格式和路径不匹配。 +1. Fallback 与格式矩阵：`error`、`drop`、`keep_original`；CSV/Parquet round trip；拒绝混合格式目录。 +1. 运维测试：连续两次 append、输入顺序置换、发布失败、并发父版本冲突，以及旧 active version 恢复。 +1. 规模测试：1 亿历史 item、大量 bucket、单热点 bucket、planner 内存上界、快照合并吞吐和 CSV 单行大小限制。 +1. ODPS integration：在临时普通表和静态分区验证平台 append、验证当前 `OdpsWriter` 两次独立运行的替换行为、版本分区隔离，以及拒绝 snapshot 中的重复 codebook 行。 + +## 10. 实施顺序 + +1. 用 golden 测试冻结 Fresh 行为，并定义 manifest/key 兼容性契约。 +1. 增加稳定 SID key、state 数据类，以及旧 state 校验和 bootstrap。 +1. 让纯 planner 接受 base occupancy，并证明空 base 与当前实现等价。 +1. 基于现有 I/O 增加 delta grouping 和完整快照的流式合并。 +1. 增加显式 `fresh|append` CLI 校验、版本化发布、统计和失败处理。 +1. 完成正确性、格式与规模测试矩阵；只有完整快照重写成本确实不可接受时，再引入 delta 文件和 compaction。 From 11989c5cac4b0962865a23827c7818e1cb7098d5 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 02:34:09 +0000 Subject: [PATCH 18/29] [refactor] clarify SID collision plan field names and trim dead code Rename the CollisionPlan/CollisionResolutionResult arrays for clarity (origin_bucket_indices, bucket_keys/bucket_counts, final_bucket_keys, overflow_bucket_key_prefixes), document CollisionPlan's fields, and remove test-only properties, single-use helpers, and redundant guards. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 75 ++++++------ tzrec/tools/sid/collision_resolution.py | 121 +++++++++++-------- tzrec/tools/sid/collision_resolution_test.py | 19 +-- 3 files changed, 111 insertions(+), 104 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 1fb1142a..ecf4be08 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -9,7 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Offline SID collision prevention with TorchEasyRec-native I/O. +r"""Offline SID collision prevention with TorchEasyRec-native I/O. The runner caps each SID bucket, loads fixed-width last-layer candidates only for overflow items, delegates collision resolution to the pure NumPy core, and @@ -20,6 +20,22 @@ The random strategy intentionally preserves the legacy deterministic baseline: it draws with replacement from the full last-layer space. Placement skips an item's origin, so an origin draw or a duplicate draw is not replaced. + +It is a single-process tool -- launch it with ``python -m`` (no torchrun / +process group). The input is a Semantic-ID table from ``tzrec.predict`` (an +``item_id`` column, a ``codes`` ``list`` column, and -- for the default +``--strategy candidate`` -- a ``candidate_codes`` ``list>`` column). + +Example:: + + python -m tzrec.tools.sid.collision_prevention \ + --input_path 'sid_predict_output/*.parquet' \ + --codebook 256,256,256 --max_items_per_codebook 5 \ + --strategy candidate --unassigned_policy keep_original \ + --output_path sid_collision/map \ + --original_sid_groups_output_path sid_collision/original_groups \ + --resolved_sid_groups_output_path sid_collision/resolved_groups \ + --diagnostics_output_path sid_collision/diagnostics """ from __future__ import annotations @@ -30,7 +46,7 @@ from contextlib import closing from dataclasses import dataclass from functools import cached_property -from typing import Dict, Iterable, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import numpy as np import pyarrow as pa @@ -66,10 +82,10 @@ _GROUP_WRITE_ITEMS = 1_000_000 _ARROW_LIST_OFFSET_MAX = int(np.iinfo(np.int32).max) +_PathIdentity = Union[str, Tuple[str, str, Optional[str], str, Optional[str]]] + -def _output_path_identity( - path: str, -) -> Union[str, Tuple[str, str, Optional[str], str, Optional[str]]]: +def _output_path_identity(path: str) -> _PathIdentity: """Return the destination identity used by the repository writer.""" if not path.startswith("odps://"): return os.path.realpath(path) @@ -151,9 +167,7 @@ def __post_init__(self) -> None: "resolved_sid_groups_output_path": self.resolved_sid_groups_output_path, "diagnostics_output_path": self.diagnostics_output_path, } - seen_paths: Dict[ - Union[str, Tuple[str, str, Optional[str], str, Optional[str]]], str - ] = {} + seen_paths: Dict[_PathIdentity, str] = {} for name, path in named_paths.items(): if not path: continue @@ -261,17 +275,6 @@ def _make_reader(self, selected_cols: List[str]) -> BaseReader: quota_name=self._config.odps_data_quota_name, ) - def _read_codes(self) -> Iterable[Dict[str, pa.Array]]: - """Stream item IDs and SIDs while resolving the default writer type.""" - reader = self._make_reader( - [self._config.item_id_field, self._config.code_field] - ) - reader_name = reader.__class__.__name__ - self._default_writer_type = self._config.writer_type or reader_name.replace( - "Reader", "Writer" - ) - yield from reader.to_batches() - @staticmethod def _is_list_type(data_type: pa.DataType) -> bool: """Return whether ``data_type`` is an Arrow list representation.""" @@ -326,10 +329,20 @@ def _validate_item_ids(self, item_ids: pa.Array) -> None: ) def _load_codes(self) -> Tuple[np.ndarray, np.ndarray]: - """Load item IDs and SIDs into arrays used by the NumPy core.""" + """Load item IDs and SIDs into arrays used by the NumPy core. + + Also resolves the default writer type from the reader class name before + streaming, so downstream writes can default to the input format. + """ + reader = self._make_reader( + [self._config.item_id_field, self._config.code_field] + ) + self._default_writer_type = self._config.writer_type or ( + reader.__class__.__name__.replace("Reader", "Writer") + ) id_chunks: List[np.ndarray] = [] code_chunks: List[np.ndarray] = [] - for batch in self._read_codes(): + for batch in reader.to_batches(): item_ids = batch[self._config.item_id_field] self._validate_item_ids(item_ids) id_chunks.append(item_ids.to_numpy(zero_copy_only=False)) @@ -355,12 +368,6 @@ def _candidate_last_codes(self, plan: CollisionPlan) -> np.ndarray: ) return self._load_candidate_last_codes(plan.overflow_item_ids) - def _candidate_matrix(self, values: pa.Array) -> np.ndarray: - """Decode candidate SIDs and retain only their last-layer codes.""" - if self._is_list_type(values.type): - return self._candidate_list_matrix(values) - return self._candidate_string_matrix(values) - def _candidate_list_matrix(self, values: pa.Array) -> np.ndarray: """Decode a nested Arrow candidate column into an ``(N, K)`` matrix.""" row_count = len(values) @@ -428,7 +435,10 @@ def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarra 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_matrix(selected) + if self._is_list_type(selected.type): + batch_candidates = self._candidate_list_matrix(selected) + else: + batch_candidates = self._candidate_string_matrix(selected) if candidates is None: candidates = np.empty( (item_count, batch_candidates.shape[1]), dtype=np.int64 @@ -508,8 +518,6 @@ def _grouped_item_ids_column( def _codes_column(codes: np.ndarray, is_csv: bool) -> pa.Array: """Encode an SID matrix for the actual output writer.""" row_count, layer_count = codes.shape - if layer_count < 1: - raise ValueError("SID output must contain at least one layer.") if is_csv: values = pc.cast(pa.array(codes.reshape(-1)), pa.string()) offsets = pa.array( @@ -648,20 +656,13 @@ def _write_sid_groups( writing original SIDs. Raises: - RuntimeError: If grouping dimensions or counts are inconsistent. ValueError: If one group exceeds Arrow list offset capacity. """ group_count = grouping.counts.shape[0] - if grouping.sid_keys.shape[0] != group_count: - raise RuntimeError("SID group keys and counts must have the same length.") - if np.any(grouping.counts <= 0): - raise RuntimeError("SID group counts must be positive.") if np.any(grouping.counts > _ARROW_LIST_OFFSET_MAX): raise ValueError("one SID group exceeds Arrow list offset capacity.") offsets = grouping.offsets - if int(offsets[-1]) != grouping.row_order.shape[0]: - raise RuntimeError("SID group counts do not match grouped row count.") with closing(self._make_writer(output_path)) as writer: is_csv = isinstance(writer, CsvWriter) max_codebook_rows = ( diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py index 30ca23da..0c66151f 100644 --- a/tzrec/tools/sid/collision_resolution.py +++ b/tzrec/tools/sid/collision_resolution.py @@ -61,25 +61,49 @@ class CollisionPlan: plus only the overflow-aligned identity and ordering data needed by candidate loading. It deliberately does not retain the full input code matrix. + + Args: + item_count: Total number of input rows. + original_last_codes: Last-layer code of every input row, int64 shape + ``(item_count,)`` in original row order. + origin_bucket_indices: Per-row position of each input row's origin bucket + *within* ``bucket_keys`` / ``bucket_counts`` (a row->bucket + index in ``[0, num_buckets)``, not a key value), int64 shape + ``(item_count,)`` in original row order. + initial_slot_indices: One-based rank of each row within its origin + bucket in deterministic (hash) order, int64 shape ``(item_count,)`` + in original row order. + bucket_keys: Flattened SID key ``band_id * last_size + last_code`` + for every occupied bucket, int64 shape ``(num_buckets,)`` indexed by + bucket index (the values in ``origin_bucket_indices``) and ascending. + bucket_counts: Item count of every bucket before capacity capping, + int64 shape ``(num_buckets,)`` aligned with ``bucket_keys``. + overflow_rows: Original row indices ranked at or beyond ``capacity`` + within their bucket (the rows to relocate), int64 in deterministic + processing order. + overflow_item_ids: Item IDs of ``overflow_rows``, aligned with it. + overflow_bucket_key_prefixes: The prefix part of each overflow row's + bucket key (``band_id * last_size``), so + ``prefix + candidate_last_code`` is the destination bucket key within + the same band; int64 aligned with ``overflow_rows``. + overflow_origin_last_codes: Origin last-layer code of each overflow row + (used to skip a candidate equal to the origin), int64 aligned with + ``overflow_rows``. + config: Collision capacity, SID shape, and fallback configuration. """ item_count: int original_last_codes: np.ndarray - origin_bucket_ids: np.ndarray + origin_bucket_indices: np.ndarray initial_slot_indices: np.ndarray - occupied_sid_keys: np.ndarray - raw_bucket_counts: np.ndarray + bucket_keys: np.ndarray + bucket_counts: np.ndarray overflow_rows: np.ndarray overflow_item_ids: np.ndarray - overflow_band_bases: np.ndarray + overflow_bucket_key_prefixes: np.ndarray overflow_origin_last_codes: np.ndarray config: CollisionResolutionConfig - @property - def origin_sid_keys(self) -> np.ndarray: - """Return flattened original SID keys aligned with input rows.""" - return self.occupied_sid_keys[self.origin_bucket_ids] - @dataclass(frozen=True) class CollisionResolutionStats: @@ -122,7 +146,7 @@ class CollisionResolutionResult: slot_indices: np.ndarray retained_mask: Optional[np.ndarray] unresolved_rows: np.ndarray - final_occupied_sid_keys: np.ndarray + final_bucket_keys: np.ndarray final_bucket_counts: np.ndarray grouping_collected: bool stats: CollisionResolutionStats @@ -153,11 +177,6 @@ def offsets(self) -> np.ndarray: np.cumsum(self.counts, dtype=np.int64, out=offsets[1:]) return offsets - @cached_property - def representative_rows(self) -> np.ndarray: - """Return the first original row index in every occupied bucket.""" - return self.row_order[self.offsets[:-1]] - def _splitmix64(values: np.ndarray) -> np.ndarray: """Hash a uint64 array with the collision tool's fixed SplitMix64 seed.""" @@ -296,31 +315,31 @@ def prepare_collision_plan( order_hashes = stable_order_hash(item_ids) ( bucket_ranks, - origin_bucket_ids, + origin_bucket_indices, sorted_rows, representative_rows, - raw_bucket_counts, + bucket_counts, ) = _within_bucket_rank(band_ids, original_last_codes, order_hashes) overflow_rows = sorted_rows[bucket_ranks[sorted_rows] >= config.capacity] bucket_ranks += 1 last_size = config.layer_sizes[-1] - occupied_sid_keys = ( + bucket_keys = ( band_ids[representative_rows] * last_size + original_last_codes[representative_rows] ) - overflow_band_bases = band_ids[overflow_rows] * last_size + overflow_bucket_key_prefixes = band_ids[overflow_rows] * last_size return CollisionPlan( item_count=codes.shape[0], original_last_codes=original_last_codes, - origin_bucket_ids=origin_bucket_ids, + origin_bucket_indices=origin_bucket_indices, initial_slot_indices=bucket_ranks, - occupied_sid_keys=occupied_sid_keys, - raw_bucket_counts=raw_bucket_counts, + bucket_keys=bucket_keys, + bucket_counts=bucket_counts, overflow_rows=overflow_rows, overflow_item_ids=item_ids[overflow_rows], - overflow_band_bases=overflow_band_bases, + overflow_bucket_key_prefixes=overflow_bucket_key_prefixes, overflow_origin_last_codes=original_last_codes[overflow_rows], config=config, ) @@ -415,25 +434,23 @@ def resolve_sid_collisions( ) if overflow_count == 0: - final_occupied_sid_keys = ( - plan.occupied_sid_keys.copy() - if collect_grouping - else np.empty(0, dtype=np.int64) + final_bucket_keys = ( + plan.bucket_keys.copy() if collect_grouping else np.empty(0, dtype=np.int64) ) final_bucket_counts = ( - plan.raw_bucket_counts.copy() + plan.bucket_counts.copy() if collect_grouping else np.empty(0, dtype=np.int64) ) max_final_bucket_size = ( - int(plan.raw_bucket_counts.max()) if plan.raw_bucket_counts.size else 0 + int(plan.bucket_counts.max()) if plan.bucket_counts.size else 0 ) return CollisionResolutionResult( resolved_last_codes=plan.original_last_codes, slot_indices=plan.initial_slot_indices, retained_mask=None, unresolved_rows=np.empty(0, dtype=np.int64), - final_occupied_sid_keys=final_occupied_sid_keys, + final_bucket_keys=final_bucket_keys, final_bucket_counts=final_bucket_counts, grouping_collected=collect_grouping, stats=CollisionResolutionStats( @@ -447,27 +464,27 @@ def resolve_sid_collisions( ) capacity = plan.config.capacity - initial_counts = np.minimum(plan.raw_bucket_counts, capacity) - slot_counts = dict(zip(map(int, plan.occupied_sid_keys), map(int, initial_counts))) + initial_counts = np.minimum(plan.bucket_counts, capacity) + slot_counts = dict(zip(map(int, plan.bucket_keys), map(int, initial_counts))) resolved_last_codes = plan.original_last_codes.copy() slot_indices = plan.initial_slot_indices.copy() relocated_count = 0 unresolved_rows = [] get_slot_count = slot_counts.get - for row_value, band_base_value, origin_value, candidate_row in zip( + for row_value, key_prefix_value, origin_value, candidate_row in zip( plan.overflow_rows, - plan.overflow_band_bases, + plan.overflow_bucket_key_prefixes, plan.overflow_origin_last_codes, candidates, ): row = int(row_value) - band_base = int(band_base_value) + key_prefix = int(key_prefix_value) origin_last_code = int(origin_value) for candidate in candidate_row.tolist(): if candidate == origin_last_code: continue - destination_key = band_base + candidate + destination_key = key_prefix + candidate destination_count = get_slot_count(destination_key, 0) if destination_count < capacity: slot_counts[destination_key] = destination_count + 1 @@ -491,12 +508,12 @@ def resolve_sid_collisions( retained_mask[unresolved_rows] = False elif unresolved_rows and fallback_policy == "keep_original": for row in unresolved_rows: - origin_key = int(plan.occupied_sid_keys[plan.origin_bucket_ids[row]]) + origin_key = int(plan.bucket_keys[plan.origin_bucket_indices[row]]) origin_count = get_slot_count(origin_key, 0) + 1 slot_counts[origin_key] = origin_count slot_indices[row] = origin_count - final_occupied_sid_keys = np.empty(0, dtype=np.int64) + final_bucket_keys = np.empty(0, dtype=np.int64) final_bucket_counts = np.empty(0, dtype=np.int64) if collect_grouping: occupancy_counts = np.fromiter( @@ -506,11 +523,11 @@ def resolve_sid_collisions( max_final_bucket_size = ( int(occupancy_counts.max()) if occupancy_counts.size else 0 ) - final_occupied_sid_keys = np.fromiter( + final_bucket_keys = np.fromiter( slot_counts, dtype=np.int64, count=len(slot_counts) ) - occupancy_order = np.argsort(final_occupied_sid_keys, kind="stable") - final_occupied_sid_keys = final_occupied_sid_keys[occupancy_order] + occupancy_order = np.argsort(final_bucket_keys, kind="stable") + final_bucket_keys = final_bucket_keys[occupancy_order] final_bucket_counts = occupancy_counts[occupancy_order] else: final_collision_buckets = 0 @@ -521,7 +538,7 @@ def resolve_sid_collisions( unresolved_array = np.asarray(unresolved_rows, dtype=np.int64) stats = CollisionResolutionStats( total_items=plan.item_count, - raw_collision_buckets=int((plan.raw_bucket_counts > capacity).sum()), + raw_collision_buckets=int((plan.bucket_counts > capacity).sum()), final_collision_buckets=final_collision_buckets, relocated_count=relocated_count, unresolved_count=len(unresolved_rows), @@ -532,7 +549,7 @@ def resolve_sid_collisions( slot_indices=slot_indices, retained_mask=retained_mask, unresolved_rows=unresolved_array, - final_occupied_sid_keys=final_occupied_sid_keys, + final_bucket_keys=final_bucket_keys, final_bucket_counts=final_bucket_counts, grouping_collected=collect_grouping, stats=stats, @@ -592,13 +609,13 @@ def row_chunks() -> Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]]: end = min(start + _GROUPING_ROW_CHUNK, plan.item_count) yield ( np.arange(start, end, dtype=np.int64), - plan.origin_bucket_ids[start:end], + plan.origin_bucket_indices[start:end], plan.initial_slot_indices[start:end], ) return _scatter_item_grouping( - plan.occupied_sid_keys, - plan.raw_bucket_counts, + plan.bucket_keys, + plan.bucket_counts, plan.item_count, row_chunks(), ) @@ -638,21 +655,19 @@ def row_chunks() -> Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]]: rows += start if rows.size == 0: continue - emitted_sid_keys = plan.occupied_sid_keys[plan.origin_bucket_ids[rows]] + emitted_sid_keys = plan.bucket_keys[plan.origin_bucket_indices[rows]] emitted_sid_keys -= plan.original_last_codes[rows] emitted_sid_keys += result.resolved_last_codes[rows] - bucket_ids = np.searchsorted( - result.final_occupied_sid_keys, emitted_sid_keys - ) - in_bounds = bucket_ids < result.final_occupied_sid_keys.shape[0] + bucket_ids = np.searchsorted(result.final_bucket_keys, emitted_sid_keys) + in_bounds = bucket_ids < result.final_bucket_keys.shape[0] if not np.all(in_bounds) or not np.all( - result.final_occupied_sid_keys[bucket_ids] == emitted_sid_keys + result.final_bucket_keys[bucket_ids] == emitted_sid_keys ): raise RuntimeError("final bucket keys do not cover every emitted SID.") yield rows, bucket_ids, result.slot_indices[rows] return _scatter_item_grouping( - result.final_occupied_sid_keys, + result.final_bucket_keys, result.final_bucket_counts, int(result.final_bucket_counts.sum()), row_chunks(), diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index 91a93ec5..0dc1acd5 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -50,12 +50,9 @@ def test_golden_candidate_resolution(self) -> None: np.testing.assert_array_equal(plan.overflow_rows, [1, 3, 8]) np.testing.assert_array_equal(plan.overflow_item_ids, [1, 3, 8]) np.testing.assert_array_equal( - plan.origin_bucket_ids, [0, 0, 0, 0, 1, 2, 2, 3, 3, 3] - ) - np.testing.assert_array_equal(plan.occupied_sid_keys, [0, 1, 2, 4]) - np.testing.assert_array_equal( - plan.origin_sid_keys, [0, 0, 0, 0, 1, 2, 2, 4, 4, 4] + plan.origin_bucket_indices, [0, 0, 0, 0, 1, 2, 2, 3, 3, 3] ) + np.testing.assert_array_equal(plan.bucket_keys, [0, 1, 2, 4]) candidates = np.asarray([[0, 1, 2], [1, 2, 3], [0, 1, 2]], dtype=np.int64) result = resolve_sid_collisions(plan, candidates) @@ -66,9 +63,7 @@ def test_golden_candidate_resolution(self) -> None: result.slot_indices, [2, 2, 1, 1, 1, 2, 1, 2, 1, 1] ) np.testing.assert_array_equal(result.unresolved_rows, []) - np.testing.assert_array_equal( - result.final_occupied_sid_keys, [0, 1, 2, 3, 4, 5] - ) + np.testing.assert_array_equal(result.final_bucket_keys, [0, 1, 2, 3, 4, 5]) np.testing.assert_array_equal(result.final_bucket_counts, [2, 2, 2, 1, 2, 1]) self.assertTrue(result.grouping_collected) self.assertIsNone(result.retained_mask) @@ -100,9 +95,6 @@ def test_golden_candidate_resolution(self) -> None: np.testing.assert_array_equal( original_grouping.row_order, [2, 0, 1, 3, 4, 6, 5, 9, 7, 8] ) - np.testing.assert_array_equal( - original_grouping.representative_rows, [2, 4, 6, 9] - ) with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): resolved_grouping = build_resolved_item_grouping(plan, result) @@ -176,7 +168,7 @@ def test_no_overflow(self) -> None: np.testing.assert_array_equal(result.resolved_last_codes, [0, 0, 2]) np.testing.assert_array_equal(result.slot_indices, [1, 2, 1]) np.testing.assert_array_equal(result.unresolved_rows, []) - np.testing.assert_array_equal(result.final_occupied_sid_keys, [0, 6]) + np.testing.assert_array_equal(result.final_bucket_keys, [0, 6]) np.testing.assert_array_equal(result.final_bucket_counts, [2, 1]) self.assertEqual(result.stats.raw_collision_buckets, 0) self.assertEqual(result.stats.relocated_count, 0) @@ -198,7 +190,6 @@ def test_empty_groupings(self) -> None: np.testing.assert_array_equal(grouping.counts, []) np.testing.assert_array_equal(grouping.row_order, []) np.testing.assert_array_equal(grouping.offsets, [0]) - np.testing.assert_array_equal(grouping.representative_rows, []) def test_random_candidate_golden_draws(self) -> None: actual = generate_random_candidate_last_codes( @@ -279,7 +270,7 @@ def test_resolution_can_skip_grouping_metadata(self) -> None: np.testing.assert_array_equal(result.slot_indices, expected.slot_indices) np.testing.assert_array_equal(result.retained_mask, expected.retained_mask) np.testing.assert_array_equal(result.unresolved_rows, expected.unresolved_rows) - np.testing.assert_array_equal(result.final_occupied_sid_keys, []) + np.testing.assert_array_equal(result.final_bucket_keys, []) np.testing.assert_array_equal(result.final_bucket_counts, []) self.assertFalse(result.grouping_collected) self.assertEqual(result.stats, expected.stats) From a1106a0da35abce24ddc58b2f98d85d532d5b062 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 03:04:07 +0000 Subject: [PATCH 19/29] [refactor] drop unused diagnostics output from SID collision tool Remove the diagnostics-output feature (--diagnostics_output_path, the _write_diagnostics writer, and the CollisionResolutionStats.to_output_dict serializer), and with it the now-redundant reassigned_count/unassigned_count stat aliases whose only purpose was the ALGR-named diagnostics columns; the stats object now speaks one vocabulary (relocated_count/unresolved_count). Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 22 +------ tzrec/tools/sid/collision_prevention_test.py | 64 ++++---------------- tzrec/tools/sid/collision_resolution.py | 25 +------- tzrec/tools/sid/collision_resolution_test.py | 13 ---- 4 files changed, 14 insertions(+), 110 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index ecf4be08..23725ee7 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -34,8 +34,7 @@ --strategy candidate --unassigned_policy keep_original \ --output_path sid_collision/map \ --original_sid_groups_output_path sid_collision/original_groups \ - --resolved_sid_groups_output_path sid_collision/resolved_groups \ - --diagnostics_output_path sid_collision/diagnostics + --resolved_sid_groups_output_path sid_collision/resolved_groups """ from __future__ import annotations @@ -125,7 +124,6 @@ class CollisionPreventionConfig: output_path: str original_sid_groups_output_path: Optional[str] resolved_sid_groups_output_path: Optional[str] - diagnostics_output_path: Optional[str] reader_type: Optional[str] writer_type: Optional[str] batch_size: int @@ -165,7 +163,6 @@ def __post_init__(self) -> None: "output_path": self.output_path, "original_sid_groups_output_path": self.original_sid_groups_output_path, "resolved_sid_groups_output_path": self.resolved_sid_groups_output_path, - "diagnostics_output_path": self.diagnostics_output_path, } seen_paths: Dict[_PathIdentity, str] = {} for name, path in named_paths.items(): @@ -193,7 +190,6 @@ def from_namespace(cls, args: argparse.Namespace) -> "CollisionPreventionConfig" resolved_sid_groups_output_path=getattr( args, "resolved_sid_groups_output_path", None ), - diagnostics_output_path=args.diagnostics_output_path, reader_type=args.reader_type, writer_type=args.writer_type, batch_size=args.batch_size, @@ -259,8 +255,6 @@ def run(self) -> CollisionResolutionStats: self._write_group_outputs(item_ids, codes, plan, result) del plan self._write_map(item_ids, codes, result) - if self._config.diagnostics_output_path: - self._write_diagnostics(result.stats) logger.info("SID collision prevention finished: %s", result.stats) return result.stats @@ -700,19 +694,6 @@ def _write_sid_groups( ) group_start = group_end - def _write_diagnostics(self, stats: CollisionResolutionStats) -> None: - """Write legacy-compatible diagnostic column names.""" - path = self._config.diagnostics_output_path - if path is None: - return - with closing(self._make_writer(path)) as writer: - writer.write( - { - name: pa.array([value], type=pa.int64()) - for name, value in stats.to_output_dict().items() - } - ) - def build_parser() -> argparse.ArgumentParser: """Build the collision-prevention command-line parser.""" @@ -737,7 +718,6 @@ def build_parser() -> argparse.ArgumentParser: "unless --rate_only is set." ), ) - parser.add_argument("--diagnostics_output_path", default=None) parser.add_argument( "--reader_type", choices=["CsvReader", "ParquetReader", "OdpsReader"], diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 7cc4b751..c80b9f83 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -33,7 +33,6 @@ output_path=None, original_sid_groups_output_path=None, resolved_sid_groups_output_path=None, - diagnostics_output_path=None, reader_type=None, writer_type=None, batch_size=100000, @@ -102,8 +101,8 @@ def test_candidate_reassigns_within_band(self) -> None: inp, out, max_items_per_codebook=2, unassigned_policy="keep_original" ) self.assertEqual(stats.total_items, 5) - self.assertEqual(stats.reassigned_count, 3) - self.assertEqual(stats.unassigned_count, 0) + self.assertEqual(stats.relocated_count, 3) + self.assertEqual(stats.unresolved_count, 0) self.assertEqual(stats.max_final_bucket_size, 2) d = self._read_parquet(out) # every final SID keeps prefix 0 (stays in band) ... @@ -128,8 +127,8 @@ def test_keep_original_when_only_origin_candidate(self) -> None: stats = self._run( inp, out, max_items_per_codebook=2, unassigned_policy="keep_original" ) - self.assertEqual(stats.reassigned_count, 0) - self.assertEqual(stats.unassigned_count, 3) + self.assertEqual(stats.relocated_count, 0) + self.assertEqual(stats.unresolved_count, 3) d = self._read_parquet(out) self.assertTrue(all(tuple(cb) == (0, 0) for cb in d["codebook"])) _, resolved_path = self._group_paths(out) @@ -141,27 +140,24 @@ def test_keep_original_when_only_origin_candidate(self) -> None: def test_error_policy_raises_on_overflow(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") - diagnostics = os.path.join(self.test_dir, "diagnostics") _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) with self.assertRaises(RuntimeError): self._run( inp, out, - diagnostics_output_path=diagnostics, max_items_per_codebook=2, ) original_path, resolved_path = self._group_paths(out) self.assertFalse(os.path.exists(out)) self.assertFalse(os.path.exists(original_path)) self.assertFalse(os.path.exists(resolved_path)) - self.assertFalse(os.path.exists(diagnostics)) def test_drop_policy_excludes_unassigned(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) stats = self._run(inp, out, max_items_per_codebook=2, unassigned_policy="drop") - self.assertEqual(stats.unassigned_count, 3) + self.assertEqual(stats.unresolved_count, 3) d = self._read_parquet(out) self.assertEqual(len(d["item_id"]), 2) original_path, resolved_path = self._group_paths(out) @@ -184,7 +180,7 @@ def test_no_overflow_does_not_require_candidates(self) -> None: out = os.path.join(self.test_dir, "out") _parquet(inp, [0, 1, 2], [[0, 0], [0, 1], [0, 2]]) stats = self._run(inp, out, max_items_per_codebook=1) - self.assertEqual(stats.reassigned_count, 0) + self.assertEqual(stats.relocated_count, 0) self.assertEqual(len(self._read_parquet(out)["item_id"]), 3) def test_candidates_align_across_batches(self) -> None: @@ -296,7 +292,7 @@ def test_random_reassigns_within_band(self) -> None: strategy="random", unassigned_policy="keep_original", ) - self.assertEqual(stats.reassigned_count, 3) + self.assertEqual(stats.relocated_count, 3) d = self._read_parquet(out) self.assertTrue(all(cb[0] == 0 for cb in d["codebook"])) self.assertLessEqual(max(Counter(map(tuple, d["codebook"])).values()), 2) @@ -330,7 +326,7 @@ def test_single_layer_sid(self) -> None: codebook="16", unassigned_policy="keep_original", ) - self.assertEqual(stats.reassigned_count, 3) + self.assertEqual(stats.relocated_count, 3) d = self._read_parquet(out) self.assertEqual(len({tuple(cb) for cb in d["codebook"]}), 4) @@ -405,7 +401,7 @@ def test_csv_backend_within_band(self) -> None: max_items_per_codebook=2, unassigned_policy="keep_original", ) - self.assertEqual(stats.reassigned_count, 3) + self.assertEqual(stats.relocated_count, 3) result = csv.read_csv(os.path.join(out, "part-0.csv")) self.assertEqual(result.schema.field("codebook").type, pa.string()) self.assertTrue(all(s.startswith("0|") for s in result["codebook"].to_pylist())) @@ -596,7 +592,6 @@ def close(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "map") - diagnostics = os.path.join(self.test_dir, "diagnostics") parquet.write_table( pa.table( { @@ -620,14 +615,13 @@ def make_writer(output_path): inp, out, writer_type="OdpsWriter", - diagnostics_output_path=diagnostics, max_items_per_codebook=1, ) original_path, resolved_path = self._group_paths(out) self.assertEqual( [path for path, _ in created_writers], - [original_path, resolved_path, out, diagnostics], + [original_path, resolved_path, out], ) self.assertTrue(all(writer.closed for _, writer in created_writers)) for _, writer in created_writers[:2]: @@ -728,7 +722,7 @@ def test_rate_only_skips_write(self) -> None: rate_only=True, ) self.assertFalse(resolve.call_args.kwargs["collect_grouping"]) - self.assertEqual(stats.reassigned_count, 3) + self.assertEqual(stats.relocated_count, 3) self.assertFalse(os.path.exists(os.path.join(out, "part-0.parquet"))) original_path, resolved_path = self._group_paths(out) self.assertFalse(os.path.exists(original_path)) @@ -751,42 +745,6 @@ def test_rate_only_allows_omitted_group_paths(self) -> None: self.assertEqual(stats.total_items, 2) self.assertFalse(os.path.exists(out)) - def test_rate_only_still_writes_diagnostics(self) -> None: - inp = os.path.join(self.test_dir, "in.parquet") - out = os.path.join(self.test_dir, "out") - diag = os.path.join(self.test_dir, "diag") - _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1]]] * 5) - self._run( - inp, - out, - diagnostics_output_path=diag, - max_items_per_codebook=2, - unassigned_policy="keep_original", - rate_only=True, - ) - result = parquet.read_table(os.path.join(diag, "part-0.parquet")).to_pydict() - self.assertEqual(result["total_items"], [5]) - self.assertFalse(os.path.exists(os.path.join(out, "part-0.parquet"))) - original_path, resolved_path = self._group_paths(out) - self.assertFalse(os.path.exists(original_path)) - self.assertFalse(os.path.exists(resolved_path)) - - def test_diagnostics_output(self) -> None: - inp = os.path.join(self.test_dir, "in.parquet") - out = os.path.join(self.test_dir, "out") - diag = os.path.join(self.test_dir, "diag") - _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1], [0, 2], [0, 3]]] * 5) - self._run( - inp, - out, - max_items_per_codebook=2, - unassigned_policy="keep_original", - diagnostics_output_path=diag, - ) - d = parquet.read_table(os.path.join(diag, "part-0.parquet")).to_pydict() - self.assertEqual(d["total_items"][0], 5) - self.assertEqual(d["reassigned_count"][0], 3) - def test_empty_input_raises(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py index 0c66151f..88bbf655 100644 --- a/tzrec/tools/sid/collision_resolution.py +++ b/tzrec/tools/sid/collision_resolution.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from functools import cached_property -from typing import Dict, Iterable, Optional +from typing import Iterable, Optional import numpy as np @@ -116,27 +116,6 @@ class CollisionResolutionStats: unresolved_count: int max_final_bucket_size: int - @property - def reassigned_count(self) -> int: - """Return the legacy name for ``relocated_count``.""" - return self.relocated_count - - @property - def unassigned_count(self) -> int: - """Return the legacy name for ``unresolved_count``.""" - return self.unresolved_count - - def to_output_dict(self) -> Dict[str, int]: - """Return diagnostics using the existing output column names.""" - return { - "total_items": self.total_items, - "raw_collision_buckets": self.raw_collision_buckets, - "final_collision_buckets": self.final_collision_buckets, - "reassigned_count": self.relocated_count, - "unassigned_count": self.unresolved_count, - "max_final_bucket_size": self.max_final_bucket_size, - } - @dataclass(frozen=True) class CollisionResolutionResult: @@ -397,7 +376,7 @@ def resolve_sid_collisions( where ``M`` equals the number of overflow rows. collect_grouping: Whether to collect final SID keys and bucket counts for :func:`build_resolved_item_grouping`. Disable this in - diagnostics-only runs to avoid the grouping arrays and sorts. + rate-only runs to avoid the grouping arrays and sorts. Returns: Resolved last-layer codes, slot indices, retention mask, unresolved row diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index 0dc1acd5..516db782 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -73,19 +73,6 @@ def test_golden_candidate_resolution(self) -> None: self.assertEqual(result.stats.relocated_count, 3) self.assertEqual(result.stats.unresolved_count, 0) self.assertEqual(result.stats.max_final_bucket_size, 2) - self.assertEqual(result.stats.reassigned_count, 3) - self.assertEqual(result.stats.unassigned_count, 0) - self.assertEqual( - result.stats.to_output_dict(), - { - "total_items": 10, - "raw_collision_buckets": 2, - "final_collision_buckets": 0, - "reassigned_count": 3, - "unassigned_count": 0, - "max_final_bucket_size": 2, - }, - ) with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): original_grouping = build_original_item_grouping(plan) From c6074594c9f5073c1fee435e521c2cf300cc1867 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 03:19:34 +0000 Subject: [PATCH 20/29] [chore] drop SID append design-plan docs These were a design proposal for an unimplemented append mode and referenced the local workspace; not part of the collision-prevention tool contribution. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/sid_collision_append_refactoring_plan.md | 156 ------------------ ...id_collision_append_refactoring_plan_zh.md | 156 ------------------ 2 files changed, 312 deletions(-) delete mode 100644 docs/sid_collision_append_refactoring_plan.md delete mode 100644 docs/sid_collision_append_refactoring_plan_zh.md diff --git a/docs/sid_collision_append_refactoring_plan.md b/docs/sid_collision_append_refactoring_plan.md deleted file mode 100644 index 00e9ea03..00000000 --- a/docs/sid_collision_append_refactoring_plan.md +++ /dev/null @@ -1,156 +0,0 @@ -# SID Collision Resolution: Fresh and Append Refactoring Plan - -> Status: design proposal; no implementation changes are included. - -## 1. Scope and terminology - -The refactoring should expose two explicit modes: - -- **Fresh** (the requested “Pure Addition” scenario): read one prediction file and run exactly the current collision-resolution algorithm. With no prior state, assignments, indexes, fallback behavior, statistics, and output ordering must remain unchanged. -- **Append**: read a new prediction file plus an existing `original_groups` snapshot and `resolved_groups` snapshot. Existing assignments are immutable; only new items are resolved and appended. - -`original_groups` records the historical mapping from each original SID to its item IDs. `resolved_groups` records the final SID buckets. The position of an item in a resolved bucket is semantic: position `k` corresponds to the current 1-based map index `k`. - -The inspected workspace artifacts are `experiments/sid_collision/sample.parquet` and `sid_collision/{original_groups,resolved_groups}`. The literal `/mnt/data/...` path was not present in this environment. The sample contains 500 unique items and is exactly the source of the existing snapshots, so using it as an append batch must fail duplicate validation rather than append another 500 copies. - -## 2. Behavioral contract - -| Property | Fresh | Append | -| ------------------------------- | --------------------- | -------------------------------- | -| Prior group inputs | Forbidden | Both snapshots required | -| Existing final SIDs and indexes | Not applicable | Never changed | -| Capacity seed | Empty | Counts from `resolved_groups` | -| Original membership | Current input | Historical groups plus new input | -| Candidate resolution | All current overflows | New overflows only | -| Map output | Current input | New-input delta only | -| Group outputs | Complete snapshots | New complete snapshots | - -Append is intentionally **not** equivalent to rerunning Fresh over the union. Old candidate lists are unavailable, historical assignments must remain stable, and existing items receive priority over new items. Consequently, append results are history-dependent. - -Default duplicate policy is `error`: item IDs must be unique in the new batch and disjoint from historical `original_groups`. Neither silent skipping nor upsert is safe because the grouped state cannot prove that a repeated item has identical SID and candidate data. - -All append outputs must use paths distinct from their input snapshots. `rate_only` should perform the same validation and planning and report new and combined statistics, but should not publish group or map data. - -## 3. Stable persisted identity - -The current dense band IDs are produced from prefixes present in one run. They are suitable for in-memory grouping but cannot identify a persisted bucket: for example, prefix `B` may be band 1 in one batch and band 0 in another. - -Introduce a `SidKeyCodec` that validates every code against `layer_sizes` and converts a full SID to a stable, checked mixed-radix `int64` key. Because the last code has unit stride, the corresponding band key can be derived independently of the current batch. Reject configurations whose key space cannot fit in `int64`; a structured multi-column key is the future fallback if such configurations must be supported. - -Persist a versioned state manifest. The group files alone cannot safely recover capacity or layer sizes—the inspected sample generator uses `(64, 64, 64)`, while the local invocation script specifies `(256, 256, 256)`. The manifest should contain: - -- state format, SID-key codec, ordering/hash, and assignment-algorithm versions; -- `layer_sizes`, capacity, strategy, fallback policy, and random-count settings; -- Arrow item-ID type/schema fingerprint; -- original and resolved item counts, parent state version, and map scope (`delta`); -- artifact locations/formats and a completion marker. - -Legacy snapshots should be bootstrapped once using explicit configuration, fully validated, and then accompanied by a manifest without changing any assignment. - -## 4. Append planning algorithm - -1. **Validate state.** Load the manifest, verify configuration and item-ID type compatibility, reject mixed CSV/Parquet directories, and ensure input and output paths do not overlap. -1. **Load compact occupancy.** Read `resolved_groups`, validate unique sorted SID keys and bucket contents, and retain only canonical SID keys and counts for planning. These counts—not `original_groups`—represent occupied capacity. -1. **Validate new IDs.** Check uniqueness in the new input. Sort its IDs once, then stream historical `original_groups` item lists through lookup against that sorted array. This avoids a Python set containing every historical item. -1. **Rank new origins.** Preserve the current deterministic hash ranking within each new original-SID bucket. -1. **Admit at origins.** For origin bucket `s`, compute `free(s) = max(capacity - existing_resolved_count(s), 0)`. Keep the first `free(s)` new items there; their absolute indexes start at `existing_resolved_count(s) + 1`. Mark the remainder as overflow. -1. **Reserve all origins first.** Add every newly retained origin item to occupancy before processing any candidate. This preserves the current Fresh rule that direct origin ownership has priority over candidate relocation. -1. **Resolve new overflow.** In deterministic order, apply the current first-fit or random strategy only to new overflow rows. Candidate occupancy starts from historical resolved counts plus all newly retained origin items. Candidates remain confined to the same canonical prefix band. -1. **Apply fallback.** `error` aborts before publication; `drop` omits unresolved new items from resolved groups and the map but retains their original membership; `keep_original` appends them to their origin even when this exceeds capacity. -1. **Build the delta.** Record each new item’s final SID and absolute 1-based index. Existing map rows are neither reconstructed nor rewritten. - -With empty base occupancy, the generalized planner must produce byte-for-byte-equivalent logical results to the current Fresh implementation. - -## 5. Snapshot merge and output rules - -The first implementation should publish complete replacement snapshots: - -- **Original snapshot:** one row per original SID; preserve the old item list as a prefix and append new IDs in the current deterministic within-bucket order. -- **Resolved snapshot:** one row per final SID; preserve the old list exactly as a prefix and append new IDs in assigned absolute-index order. This invariant makes every emitted map index verifiable by direct lookup. -- **Map:** emit rows for the new input only, matching the existing meaning of `output_path` as the mapping for the input being processed. - -Use TorchEasyRec’s existing readers and writers, including the CSV writer. A small state-specific adapter is justified for manifest validation, normalized CSV/Parquet group decoding, streaming merge, and publication; a second generic I/O framework is not. - -CSV represents an entire `itemids` group as one JSON-like field, so a hot bucket can create a line larger than the reader block size. Parquet or ODPS should be preferred for large state. Write chunks bound table/CSV batch memory, but they cannot split one oversized group without changing the schema. - -Writers currently replace outputs independently, so do not update state in place. Write all artifacts under a new version, close and validate them, and publish the manifest/current-version pointer last. Include the parent version and use a lock or compare-and-swap check to prevent concurrent append jobs from losing updates. - -If repeated full snapshots become too expensive at 100 million items, add an explicit base-plus-ordered-deltas format with periodic compaction as a later phase. Do not emit duplicate codebook rows under the current one-row-per-bucket snapshot contract. - -## 6. ODPS append support check - -**Conclusion (verified 2026-07-14): MaxCompute supports physical row append, but the current TorchEasyRec writer does not expose it, and direct physical append is incorrect for the grouped-state schema.** - -| Layer | Append capability | Consequence for this design | -| ------------------------- | --------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | -| MaxCompute SQL | `INSERT INTO` appends to a standard table or static partition; `INSERT OVERWRITE` replaces it | The platform has append capability, subject to table restrictions | -| PyODPS / Storage API | Non-overwrite writes can append; `overwrite=True` requests replacement | The underlying client can represent both modes | -| TorchEasyRec `OdpsWriter` | Hardcodes `TableBatchWriteRequest(..., overwrite=True)` | Every writer run replaces the target table or partition | -| SID group snapshot | Requires one row per `codebook` with the complete `itemids` list | Appending a touched bucket would create a duplicate row, not update its list | - -The platform behavior is documented by [MaxCompute `INSERT INTO` / `INSERT OVERWRITE`](https://www.alibabacloud.com/help/en/maxcompute/user-guide/insert-or-update-data-into-a-table-or-a-static-partition), [TableTunnel overwrite semantics](https://www.alibabacloud.com/help/en/maxcompute/user-guide/tabletunnel), [PyODPS table writes](https://pyodps.readthedocs.io/en/latest/base-tables.html#write-data-to-tables), and [PyODPS Storage API writes](https://pyodps.readthedocs.io/en/stable/base-storage-api-v2.html#writing-data). MaxCompute also states that `INSERT INTO` is unsupported for clustered tables and that concurrent inserts to the same target are not protected by table locking. Its [ACID documentation](https://www.alibabacloud.com/help/en/maxcompute/product-overview/acid-semantics) does not provide a transaction spanning the separate SID output tables. - -Repository inspection confirms that `tzrec/datasets/odps_dataset.py:762-764` always supplies `overwrite=True`; neither `OdpsWriter.__init__` nor the collision workflow passes an append option. Repeated `write()` calls accumulate batches only inside that one overwrite session. Existing ODPS tests validate a single writer session, not retention across two writer runs. The repository pins PyODPS 0.12.5.1, whose [`TableBatchWriteRequest`](https://raw.githubusercontent.com/aliyun/aliyun-odps-python-sdk/v0.12.5.1/odps/apis/storage_api/storage_api.py) accepts an overwrite flag, but the explicit `True` in TorchEasyRec prevents append. - -For example, physical ODPS append changes `(A, [1, 2])` plus `(A, [3])` into two rows for codebook `A`; it does not produce `(A, [1, 2, 3])`. This violates the current snapshot invariant even though the database append itself succeeds. - -Therefore, the initial ODPS implementation should keep **logical Append** separate from **physical row append**: - -1. Read the parent `original_groups` and `resolved_groups` partitions. -1. Merge the old lists with the new delta in the application. -1. Write complete group snapshots to new immutable `state_version=` partitions using the existing overwrite writer. -1. Validate all artifacts and publish the completed-version manifest last. -1. Enforce a single writer or parent-version check because the four outputs do not share one transaction and MaxCompute provides no insert lock. - -Do not add a generic `append=True` switch to `OdpsWriter` merely for this workflow. Physical append becomes appropriate only if a future, explicitly different schema stores ordered bucket fragments such as `(state_version, sequence, codebook, itemids_delta)` and readers fold those deltas before use. - -No ODPS credentials or CI project are configured in this workspace, so this verification combines official documentation, the pinned and installed PyODPS request models, and repository source inspection rather than a live remote write. Before release, add an ODPS integration gate using a temporary non-clustered table and static partition with `ARRAY` item IDs: write `(A, [1, 2])`, append `(B, [3])` and `(A, [4])` with `overwrite=False`, confirm that old rows survive and `A` is duplicated rather than merged, verify `overwrite=True` replacement, and inject a failed version publication to prove the parent snapshot remains active. - -## 7. Proposed modularization - -- `collision_resolution.py`: pure array-based Fresh/Append planning; no storage concerns. -- `SidKeyCodec`: stable SID encoding, decoding, range checks, and band identity. -- `BucketOccupancy`: compact historical counts plus a mutable overlay for touched buckets. -- `CollisionPlan` / `CollisionResult`: distinguish base counts, new admissions, new final codes/indexes, unresolved rows, and combined counts. -- A state module such as `collision_state.py`: manifest, group readers, integrity checks, snapshot merge, and atomic publication using existing repository I/O. -- `collision_prevention.py`: CLI parsing, mode validation, orchestration, statistics, and writer selection. - -Keep Arrow arrays at the I/O boundary and NumPy arrays in the CPU planner. Converting everything to tensors would add copies and device concerns without benefiting this sort/group/merge workload. Avoid a global Python dictionary or item-ID set at 100-million scale; use sorted arrays, vectorized search, streaming historical IDs, and a small overlay for buckets touched by the new batch. - -The existing grouping helper must not scatter new rows directly into arrays sized by combined counts. Append indexes are absolute, while a delta buffer is relative; use `delta_position = absolute_index - base_bucket_count` when constructing new group fragments. - -## 8. Result-affecting decisions - -These rules must be documented and treated as compatibility contracts: - -- Existing items always win capacity and never move. -- Existing resolved item order—and therefore indexes—is immutable. -- All new origin admissions are reserved before any new candidate assignment. -- New-item ordering remains deterministic and input-order invariant. -- The capacity, `layer_sizes`, item-ID type, hash version, and assignment version must match the parent state. -- `drop` permits resolved state to contain fewer items than original state; `keep_original` permits over-capacity buckets. -- The append map is a delta. A full historical map would require the old map as an additional state input or a costly reconstruction. - -Temporary full-array/debug printing currently present in the working tree must be removed before scale tests because it would dominate runtime and memory, but that cleanup is outside this design-only change. - -## 9. Verification plan - -1. Golden/differential test: Fresh with empty state exactly matches the current implementation for first-fit, random, all fallbacks, and `rate_only`. -1. Stable-key test: historical prefixes `A,B`, followed by a new batch containing only `B`, catches batch-local band IDs. -1. Append cases: empty, partially filled, full, and already-overfull origins; occupied candidate targets; new-only buckets; and competing candidates. -1. Immutability checks: every old final SID, item-list prefix, and index remains unchanged; every new map index points to the merged resolved list. -1. Duplicate checks: within-new and new-versus-history; appending the supplied sample to the supplied snapshots must reject all 500 overlaps. -1. Compatibility checks: capacity, layer sizes, item-ID type, manifest version, schema, format, and path mismatch. -1. Fallback and format matrix: `error`, `drop`, `keep_original`; CSV and Parquet round trips; mixed-format directories rejected. -1. Operational tests: two sequential appends, input-order permutation, failed publication, concurrent-parent mismatch, and recovery of the prior active version. -1. Scale benchmarks: 100-million historical items, many buckets, one hot bucket, bounded planner memory, snapshot merge throughput, and CSV line-size limits. -1. ODPS integration: verify platform append in a temporary standard table/static partition, current `OdpsWriter` replacement across separate runs, version-partition isolation, and rejection of duplicate-codebook snapshot rows. - -## 10. Implementation sequence - -1. Freeze Fresh behavior with golden tests and define the manifest/key compatibility contract. -1. Add stable SID keys, state data classes, and legacy-state validation/bootstrap. -1. Generalize the pure planner to accept base occupancy; prove empty-base equivalence. -1. Add delta grouping and streaming full-snapshot merge through existing I/O. -1. Add explicit `fresh|append` CLI validation, versioned publication, statistics, and failure handling. -1. Complete the correctness/format/scale matrix, then consider delta files plus compaction only if snapshot rewrite cost warrants it. diff --git a/docs/sid_collision_append_refactoring_plan_zh.md b/docs/sid_collision_append_refactoring_plan_zh.md deleted file mode 100644 index e47cdfc4..00000000 --- a/docs/sid_collision_append_refactoring_plan_zh.md +++ /dev/null @@ -1,156 +0,0 @@ -# SID 碰撞处理：Fresh 与 Append 重构方案 - -> 状态：设计方案；不包含代码实现变更。 - -## 1. 范围与术语 - -重构后应显式支持两种模式： - -- **Fresh**（需求中的“纯新增”）：只读取一个预测文件，完整执行当前碰撞处理算法。没有历史状态时，分配结果、索引、fallback 行为、统计和输出顺序都必须与当前实现保持一致。 -- **Append**：读取一个新预测文件，以及已有的 `original_groups` 和 `resolved_groups` 快照。历史分配不可变，只处理并追加新 item。 - -`original_groups` 保存“原始 SID -> 历史 item IDs”；`resolved_groups` 保存“最终 SID -> item IDs”。item 在 resolved bucket 中的位置具有业务含义：第 `k` 个位置对应当前 map 中从 1 开始的索引 `k`。 - -本次实际检查的是工作区中的 `experiments/sid_collision/sample.parquet` 和 `sid_collision/{original_groups,resolved_groups}`；当前环境中不存在字面量 `/mnt/data/...` 路径。sample 包含 500 个唯一 item，而且正是现有快照的数据来源。因此，若把该 sample 作为 append 新批次，正确行为应是重复 ID 校验失败，而不是再次追加 500 份数据。 - -## 2. 行为契约 - -| 属性 | Fresh | Append | -| ------------------- | --------------------- | ----------------------------- | -| 历史 group 输入 | 禁止 | 必须同时提供两个快照 | -| 历史最终 SID 与索引 | 不适用 | 永不改变 | -| 容量初始值 | 空 | 来自 `resolved_groups` 的计数 | -| 原始归属 | 当前输入 | 历史 groups 加新输入 | -| Candidate 处理范围 | 当前批次全部 overflow | 只处理新 item 的 overflow | -| Map 输出 | 当前输入 | 只输出新批次增量 | -| Group 输出 | 完整快照 | 新的完整快照 | - -Append 有意地**不等价于**对“历史 + 新数据”重新执行 Fresh。原因是历史 candidate 列表不可用、历史分配必须稳定，而且旧 item 的优先级高于新 item。因此，append 结果必然与批次历史有关。 - -重复 ID 的默认策略应为 `error`：新批次内部必须唯一，并且不得与历史 `original_groups` 相交。不能默认静默跳过或 upsert，因为 group 状态无法证明重复 item 的 SID 和 candidate 数据完全一致。 - -Append 输出路径必须与两个输入快照不同。`rate_only` 仍应完成相同的校验和规划，并报告新增与合并后的统计，但不发布 group 或 map 数据。 - -## 3. 稳定的持久化标识 - -当前 dense band ID 由“本批次实际出现的 prefix”生成，适合单次内存分组，却不能标识持久化 bucket。例如，prefix `B` 在一个批次中可能是 band 1，在另一个只有 `B` 的批次中却变成 band 0。 - -应引入 `SidKeyCodec`：先按 `layer_sizes` 校验每一层 code，再把完整 SID 编码为稳定且有溢出检查的 mixed-radix `int64` key。最后一层的步长为 1，因此可以得到不依赖当前批次的 band key。如果整个 key 空间无法放入 `int64`，应拒绝配置；若未来确有需要，再使用结构化多列 key 作为替代方案。 - -还应持久化带版本的 state manifest。仅从 group 文件无法安全恢复 capacity 或 layer sizes：本次检查到 sample 生成器使用 `(64, 64, 64)`，而本地调用脚本配置为 `(256, 256, 256)`。manifest 至少应包含： - -- state 格式、SID key codec、排序/hash 和分配算法版本； -- `layer_sizes`、capacity、strategy、fallback policy 和 random-count 配置； -- Arrow item-ID 类型或 schema fingerprint； -- original/resolved item 数、父 state 版本，以及 map 范围（`delta`）； -- 各 artifact 的位置、格式和完成标记。 - -旧快照需要进行一次 bootstrap：显式提供配置，完整校验后生成 manifest，但不改变任何历史分配。 - -## 4. Append 规划算法 - -1. **校验 state。** 读取 manifest，检查配置与 item-ID 类型兼容性；拒绝 CSV/Parquet 混合目录；确认输入与输出路径不重叠。 -1. **读取紧凑 occupancy。** 从 `resolved_groups` 读取并校验唯一、已排序的 SID key 和 bucket 内容；规划阶段只保留 canonical SID key 与计数。占用容量必须来自 `resolved_groups`，而不是 `original_groups`。 -1. **校验新 ID。** 检查新批次内部唯一性；对新 ID 排序一次，再流式扫描历史 `original_groups.itemids` 并查找交集，避免为全部历史 item 创建 Python set。 -1. **对新 origin 排序。** 在每个新 original-SID bucket 内继续使用当前确定性的 hash rank。 -1. **在 origin 接纳新 item。** 对 origin bucket `s`，计算 `free(s) = max(capacity - existing_resolved_count(s), 0)`。前 `free(s)` 个新 item 留在 origin，其绝对索引从 `existing_resolved_count(s) + 1` 开始；其余标记为 overflow。 -1. **先预占全部 origin。** 处理任何 candidate 之前，先把所有被 origin 接纳的新 item 加入 occupancy。这样才能保留当前 Fresh 中“直接归属优先于 candidate 迁入”的规则。 -1. **处理新 overflow。** 按确定性顺序，只对新 overflow 执行当前 first-fit 或 random 策略。candidate occupancy 的初值为历史 resolved 计数加上全部新 origin 接纳项；candidate 仍须限制在同一个 canonical prefix band 中。 -1. **应用 fallback。** `error` 在发布前终止；`drop` 不把 unresolved 新 item 写入 resolved groups 和 map，但仍保留其 original membership；`keep_original` 把它追加回 origin，即使 bucket 因此超过 capacity。 -1. **构建增量结果。** 保存每个新 item 的最终 SID 和绝对的、从 1 开始的索引；不重建或重写旧 map 行。 - -当 base occupancy 为空时，通用 planner 的逻辑结果必须与当前 Fresh 实现完全一致。 - -## 5. 快照合并与输出规则 - -第一阶段应输出完整的新快照： - -- **Original 快照：** 每个 original SID 一行；旧 item 列表必须保持为前缀，新 ID 按当前 bucket 内确定性顺序追加。 -- **Resolved 快照：** 每个 final SID 一行；旧列表必须原样保持为前缀，新 ID 按已分配的绝对索引顺序追加。这样每个新 map index 都可以直接在合并后的列表中验证。 -- **Map：** 只输出当前新输入的记录，保持 `output_path` 表示“本次输入的映射”这一现有语义。 - -继续复用 TorchEasyRec 已有 reader/writer，包括 CSV writer。可以增加一个小型、面向 state 的适配模块，负责 manifest 校验、CSV/Parquet group 统一解码、流式合并和发布；没有必要再建立一套通用 I/O 框架。 - -CSV 会把整个 `itemids` group 放进一个类似 JSON 的字段，因此热点 bucket 可能产生超过 reader block size 的单行。大规模 state 应优先选择 Parquet 或 ODPS。写入 chunk 可以控制 table/CSV 批次内存，但在不修改 schema 的前提下，无法拆分单个超大 group。 - -当前 writer 会彼此独立地覆盖输出，因此不能原地更新 state。应先把所有 artifact 写入一个新版本，关闭并校验后，最后发布 manifest 或 current-version 指针。manifest 中记录父版本，并使用锁或 compare-and-swap 校验，避免并发 append 互相覆盖。 - -如果 1 亿 item 场景下频繁重写完整快照成本过高，可以在后续阶段引入显式的“base + 有序 delta + 定期 compaction”格式。当前“一 bucket 一行”的快照契约下，不能直接输出重复 codebook 行。 - -## 6. ODPS Append 能力检查 - -**结论（验证日期：2026-07-14）：MaxCompute 支持物理行追加，但 TorchEasyRec 当前 writer 没有暴露该能力；而且对于现有 grouped-state schema，直接物理追加是错误的。** - -| 层次 | Append 能力 | 对本设计的影响 | -| ------------------------- | --------------------------------------------------------------------- | ---------------------------------------------------- | -| MaxCompute SQL | `INSERT INTO` 可追加到普通表或静态分区；`INSERT OVERWRITE` 会替换目标 | 平台具备追加能力，但受表类型等限制 | -| PyODPS / Storage API | non-overwrite write 可以追加；`overwrite=True` 表示替换 | 底层 client 能表达两种模式 | -| TorchEasyRec `OdpsWriter` | 固定使用 `TableBatchWriteRequest(..., overwrite=True)` | 每次 writer 运行都会替换目标表或分区 | -| SID group 快照 | 每个 `codebook` 必须只有一行，且 `itemids` 是完整列表 | 对已存在 bucket 追加行只会产生重复行，不会更新原列表 | - -平台能力可参见 [MaxCompute `INSERT INTO` / `INSERT OVERWRITE`](https://www.alibabacloud.com/help/en/maxcompute/user-guide/insert-or-update-data-into-a-table-or-a-static-partition)、[TableTunnel overwrite 语义](https://www.alibabacloud.com/help/en/maxcompute/user-guide/tabletunnel)、[PyODPS table write](https://pyodps.readthedocs.io/en/latest/base-tables.html#write-data-to-tables) 和 [PyODPS Storage API write](https://pyodps.readthedocs.io/en/stable/base-storage-api-v2.html#writing-data)。MaxCompute 还明确说明：`INSERT INTO` 不支持 clustered table，并且对同一目标的并发 INSERT 没有 table lock 保护；其 [ACID 文档](https://www.alibabacloud.com/help/en/maxcompute/product-overview/acid-semantics) 也没有提供跨多个 SID 输出表的统一事务。 - -仓库检查显示，`tzrec/datasets/odps_dataset.py:762-764` 始终传入 `overwrite=True`；`OdpsWriter.__init__` 和 collision 流程都没有 append 参数。同一个 writer 内多次调用 `write()` 只是在同一个 overwrite session 内累积 batch。现有 ODPS 测试只验证一次 writer session，没有验证两次独立运行后是否保留旧数据。仓库固定使用 PyODPS 0.12.5.1；该版本的 [`TableBatchWriteRequest`](https://raw.githubusercontent.com/aliyun/aliyun-odps-python-sdk/v0.12.5.1/odps/apis/storage_api/storage_api.py) 接受 overwrite flag，但 TorchEasyRec 显式传入的 `True` 阻止了 append。 - -例如，ODPS 物理 append 会把 `(A, [1, 2])` 加 `(A, [3])` 变成两行 codebook `A`，而不是 `(A, [1, 2, 3])`。因此，即使数据库追加成功，结果仍违反当前 snapshot invariant。 - -因此，第一阶段 ODPS 实现必须区分**逻辑 Append**和**物理行 append**： - -1. 读取父版本的 `original_groups` 与 `resolved_groups` 分区。 -1. 在应用层合并旧列表和新 delta。 -1. 使用现有 overwrite writer，把完整 group 快照写到新的不可变 `state_version=` 分区。 -1. 校验全部 artifact 后，最后发布 completed-version manifest。 -1. 使用 single-writer 或父版本校验，因为四类输出不共享一个事务，而且 MaxCompute 没有 insert lock。 - -不要仅为本流程给通用 `OdpsWriter` 增加简单的 `append=True` 开关。只有未来显式引入不同的 delta schema，例如 `(state_version, sequence, codebook, itemids_delta)`，并要求 reader 在使用前折叠这些 delta 时，物理 append 才适合。 - -当前工作区没有配置 ODPS credential 或 CI project，因此本次验证基于官方文档、固定版本和已安装版本的 PyODPS request model，以及仓库源码，而不是远程 live write。发布前应增加 ODPS integration gate：使用带 `ARRAY` item IDs 的临时非 clustered 普通表和静态分区，先写 `(A, [1, 2])`，再用 `overwrite=False` 追加 `(B, [3])` 与 `(A, [4])`；确认旧行仍存在且 `A` 产生重复行而不是合并，验证 `overwrite=True` 会替换目标，并注入版本发布失败以确认父快照仍为 active。 - -## 7. 建议的模块划分 - -- `collision_resolution.py`：纯数组的 Fresh/Append 规划逻辑，不涉及存储。 -- `SidKeyCodec`：稳定 SID 编解码、范围检查和 band 标识。 -- `BucketOccupancy`：紧凑的历史计数，以及仅包含本次触达 bucket 的可变 overlay。 -- `CollisionPlan` / `CollisionResult`：明确区分 base counts、新接纳项、新 item 的最终 code/index、unresolved 行和合并后计数。 -- 新的 state 模块（如 `collision_state.py`）：manifest、group reader、完整性检查、快照合并，以及基于现有仓库 I/O 的原子发布。 -- `collision_prevention.py`：CLI 参数、mode 校验、流程编排、统计和 writer 选择。 - -I/O 边界继续使用 Arrow，CPU planner 使用 NumPy。把全部数据统一成 tensor 会增加复制和 device 管理，却不会改善这种排序、分组和合并工作负载。1 亿数据规模下，应避免全量 Python dict 或 item-ID set，改用排序数组、向量化查找、流式历史 ID，以及只保存新批次触达 bucket 的小型 overlay。 - -现有 grouping helper 不能直接把新行 scatter 到以“合并后 counts”为大小的数组中。Append index 是绝对位置，而 delta buffer 使用相对位置；构造新 group fragment 时应使用 `delta_position = absolute_index - base_bucket_count`。 - -## 8. 会影响结果的关键决策 - -以下规则需要写入兼容性契约： - -- 历史 item 永远优先占用容量，且绝不迁移。 -- 历史 resolved item 顺序以及对应 index 永远不变。 -- 必须先预占所有新 origin 接纳项，再执行任何新 candidate 分配。 -- 新 item 排序保持确定性，并且不受输入行顺序影响。 -- capacity、`layer_sizes`、item-ID 类型、hash 版本和分配算法版本必须与父 state 一致。 -- `drop` 允许 resolved state 的 item 数少于 original state；`keep_original` 允许 bucket 超过 capacity。 -- Append map 是增量。若需要完整历史 map，就必须把旧 map 作为额外 state 输入，或承担昂贵的重建成本。 - -当前工作树中存在临时的全数组/debug 输出；在大规模测试前必须删除，否则输出本身会主导耗时和内存。但该清理不属于本次仅设计文档的变更范围。 - -## 9. 验证计划 - -1. Golden/differential 测试：空 state 下的 Fresh 在 first-fit、random、全部 fallback 和 `rate_only` 模式中都与当前实现一致。 -1. 稳定 key 测试：历史 prefix 为 `A,B`，新批次只有 `B`，用于捕获 batch-local band ID 问题。 -1. Append 场景：空、部分占用、满和已超容量 origin；已占用 candidate target；仅新数据 bucket；candidate 竞争。 -1. 不变性检查：所有旧 final SID、item-list 前缀和 index 保持不变；每个新 map index 都指向合并后的 resolved list。 -1. 重复检查：新批次内部、新批次与历史之间；把提供的 sample append 到提供的快照时，应检测并拒绝 500 个重复 ID。 -1. 兼容性检查：capacity、layer sizes、item-ID 类型、manifest 版本、schema、格式和路径不匹配。 -1. Fallback 与格式矩阵：`error`、`drop`、`keep_original`；CSV/Parquet round trip；拒绝混合格式目录。 -1. 运维测试：连续两次 append、输入顺序置换、发布失败、并发父版本冲突，以及旧 active version 恢复。 -1. 规模测试：1 亿历史 item、大量 bucket、单热点 bucket、planner 内存上界、快照合并吞吐和 CSV 单行大小限制。 -1. ODPS integration：在临时普通表和静态分区验证平台 append、验证当前 `OdpsWriter` 两次独立运行的替换行为、版本分区隔离，以及拒绝 snapshot 中的重复 codebook 行。 - -## 10. 实施顺序 - -1. 用 golden 测试冻结 Fresh 行为，并定义 manifest/key 兼容性契约。 -1. 增加稳定 SID key、state 数据类，以及旧 state 校验和 bootstrap。 -1. 让纯 planner 接受 base occupancy，并证明空 base 与当前实现等价。 -1. 基于现有 I/O 增加 delta grouping 和完整快照的流式合并。 -1. 增加显式 `fresh|append` CLI 校验、版本化发布、统计和失败处理。 -1. 完成正确性、格式与规模测试矩阵；只有完整快照重写成本确实不可接受时，再引入 delta 文件和 compaction。 From 379eb3d969132dffd759a2c3ddaad9fb3a8fee0a Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 03:21:17 +0000 Subject: [PATCH 21/29] [chore] keep tools/sid/__init__.py identical to upstream Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tzrec/tools/sid/__init__.py b/tzrec/tools/sid/__init__.py index c4aae99f..eedc773b 100644 --- a/tzrec/tools/sid/__init__.py +++ b/tzrec/tools/sid/__init__.py @@ -8,5 +8,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -"""SID offline tools.""" From 02bb1a7e4ad560997fb2dfd8cc43f88a0511e101 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 03:54:23 +0000 Subject: [PATCH 22/29] [refactor] SID collision: keep_original as the only unassigned policy Drop the error and drop policies (and the --unassigned_policy config): an overflow item that cannot be relocated now always keeps its original SID over capacity, so every input item is preserved. Removes the retained_mask machinery that only existed for the drop policy. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 28 +------- tzrec/tools/sid/collision_prevention_test.py | 50 ++------------ tzrec/tools/sid/collision_resolution.py | 70 +++++++------------- tzrec/tools/sid/collision_resolution_test.py | 51 +++----------- 4 files changed, 40 insertions(+), 159 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 23725ee7..50f0ac13 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -31,7 +31,7 @@ python -m tzrec.tools.sid.collision_prevention \ --input_path 'sid_predict_output/*.parquet' \ --codebook 256,256,256 --max_items_per_codebook 5 \ - --strategy candidate --unassigned_policy keep_original \ + --strategy candidate \ --output_path sid_collision/map \ --original_sid_groups_output_path sid_collision/original_groups \ --resolved_sid_groups_output_path sid_collision/resolved_groups @@ -132,7 +132,6 @@ class CollisionPreventionConfig: candidate_codes_field: str layer_sizes: Tuple[int, ...] max_items_per_codebook: int - unassigned_policy: str strategy: str random_num_candidates: int rate_only: bool @@ -198,7 +197,6 @@ def from_namespace(cls, args: argparse.Namespace) -> "CollisionPreventionConfig" candidate_codes_field=args.candidate_codes_field, layer_sizes=layer_sizes, max_items_per_codebook=args.max_items_per_codebook, - unassigned_policy=args.unassigned_policy, strategy=args.strategy, random_num_candidates=args.random_num_candidates, rate_only=args.rate_only, @@ -211,7 +209,6 @@ def resolution_config(self) -> CollisionResolutionConfig: return CollisionResolutionConfig( layer_sizes=self.layer_sizes, capacity=self.max_items_per_codebook, - fallback_policy=self.unassigned_policy, ) @@ -547,16 +544,7 @@ def _write_map( """Write the resolved map in chunks without a full final-code copy.""" with closing(self._make_writer(self._config.output_path)) as writer: is_csv = isinstance(writer, CsvWriter) - retained_rows = ( - np.flatnonzero(result.retained_mask) - if result.retained_mask is not None - else None - ) - output_count = ( - retained_rows.shape[0] - if retained_rows is not None - else item_ids.shape[0] - ) + output_count = item_ids.shape[0] write_chunk = _MAP_WRITE_ROWS if not is_csv: write_chunk = min( @@ -567,12 +555,7 @@ def _write_map( raise ValueError("one SID row exceeds Arrow list offset capacity.") for start in range(0, output_count, write_chunk): end = min(start + write_chunk, output_count) - selection: Union[slice, np.ndarray] - selection = ( - retained_rows[start:end] - if retained_rows is not None - else slice(start, end) - ) + selection = slice(start, end) origin_chunk = origin_codes[selection] final_chunk = origin_chunk.copy() final_chunk[:, -1] = result.resolved_last_codes[selection] @@ -739,11 +722,6 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--candidate_codes_field", default="candidate_codes") parser.add_argument("--max_items_per_codebook", type=int, required=True) - parser.add_argument( - "--unassigned_policy", - choices=["error", "drop", "keep_original"], - default="error", - ) parser.add_argument( "--strategy", choices=["candidate", "random"], diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index c80b9f83..f87023d4 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -40,7 +40,6 @@ code_field="codes", candidate_codes_field="candidate_codes", max_items_per_codebook=2, - unassigned_policy="error", strategy="candidate", codebook="8,8", random_num_candidates=64, @@ -97,9 +96,7 @@ def test_candidate_reassigns_within_band(self) -> None: out = os.path.join(self.test_dir, "out") # 5 items in bucket (0,0); candidates vary only the last layer within band 0. _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 1], [0, 2], [0, 3]]] * 5) - stats = self._run( - inp, out, max_items_per_codebook=2, unassigned_policy="keep_original" - ) + stats = self._run(inp, out, max_items_per_codebook=2) self.assertEqual(stats.total_items, 5) self.assertEqual(stats.relocated_count, 3) self.assertEqual(stats.unresolved_count, 0) @@ -114,7 +111,7 @@ def test_output_is_list_int64(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") _parquet(inp, [0, 1, 2], [[0, 0]] * 3, [[[0, 1]]] * 3) - self._run(inp, out, max_items_per_codebook=2, unassigned_policy="keep_original") + self._run(inp, out, max_items_per_codebook=2) schema = parquet.read_table(os.path.join(out, "part-0.parquet")).schema self.assertEqual(schema.field("codebook").type, pa.list_(pa.int64())) self.assertEqual(schema.field("origin_codebook").type, pa.list_(pa.int64())) @@ -122,14 +119,13 @@ def test_output_is_list_int64(self) -> None: def test_keep_original_when_only_origin_candidate(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") - # candidate == origin last -> skipped -> nothing to place. + # candidate == origin last -> skipped -> nothing to place; kept over cap. _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) - stats = self._run( - inp, out, max_items_per_codebook=2, unassigned_policy="keep_original" - ) + stats = self._run(inp, out, max_items_per_codebook=2) self.assertEqual(stats.relocated_count, 0) self.assertEqual(stats.unresolved_count, 3) d = self._read_parquet(out) + self.assertEqual(len(d["item_id"]), 5) # every item preserved self.assertTrue(all(tuple(cb) == (0, 0) for cb in d["codebook"])) _, resolved_path = self._group_paths(out) resolved = self._read_parquet(resolved_path) @@ -137,37 +133,6 @@ def test_keep_original_when_only_origin_candidate(self) -> None: for item_id, index in zip(d["item_id"], d["index"]): self.assertEqual(resolved["itemids"][0][index - 1], item_id) - def test_error_policy_raises_on_overflow(self) -> None: - inp = os.path.join(self.test_dir, "in.parquet") - out = os.path.join(self.test_dir, "out") - _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) - with self.assertRaises(RuntimeError): - self._run( - inp, - out, - max_items_per_codebook=2, - ) - original_path, resolved_path = self._group_paths(out) - self.assertFalse(os.path.exists(out)) - self.assertFalse(os.path.exists(original_path)) - self.assertFalse(os.path.exists(resolved_path)) - - def test_drop_policy_excludes_unassigned(self) -> None: - inp = os.path.join(self.test_dir, "in.parquet") - out = os.path.join(self.test_dir, "out") - _parquet(inp, list(range(5)), [[0, 0]] * 5, [[[0, 0]]] * 5) - stats = self._run(inp, out, max_items_per_codebook=2, unassigned_policy="drop") - self.assertEqual(stats.unresolved_count, 3) - d = self._read_parquet(out) - self.assertEqual(len(d["item_id"]), 2) - original_path, resolved_path = self._group_paths(out) - original = self._read_parquet(original_path) - resolved = self._read_parquet(resolved_path) - self.assertEqual(len(original["itemids"][0]), 5) - self.assertEqual(len(resolved["itemids"][0]), 2) - for item_id, index in zip(d["item_id"], d["index"]): - self.assertEqual(resolved["itemids"][0][index - 1], item_id) - def test_raises_when_overflow_but_no_candidates(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") @@ -290,7 +255,6 @@ def test_random_reassigns_within_band(self) -> None: out, max_items_per_codebook=2, strategy="random", - unassigned_policy="keep_original", ) self.assertEqual(stats.relocated_count, 3) d = self._read_parquet(out) @@ -324,7 +288,6 @@ def test_single_layer_sid(self) -> None: max_items_per_codebook=1, strategy="random", codebook="16", - unassigned_policy="keep_original", ) self.assertEqual(stats.relocated_count, 3) d = self._read_parquet(out) @@ -352,7 +315,6 @@ def decisions(order): inp, out, max_items_per_codebook=2, - unassigned_policy="keep_original", ) d = self._read_parquet(out) original_path, resolved_path = self._group_paths(out) @@ -399,7 +361,6 @@ def test_csv_backend_within_band(self) -> None: reader_type="CsvReader", writer_type="CsvWriter", max_items_per_codebook=2, - unassigned_policy="keep_original", ) self.assertEqual(stats.relocated_count, 3) result = csv.read_csv(os.path.join(out, "part-0.csv")) @@ -718,7 +679,6 @@ def test_rate_only_skips_write(self) -> None: inp, out, max_items_per_codebook=2, - unassigned_policy="keep_original", rate_only=True, ) self.assertFalse(resolve.call_args.kwargs["collect_grouping"]) diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py index 88bbf655..751612ea 100644 --- a/tzrec/tools/sid/collision_resolution.py +++ b/tzrec/tools/sid/collision_resolution.py @@ -13,14 +13,13 @@ from dataclasses import dataclass from functools import cached_property -from typing import Iterable, Optional +from typing import Iterable import numpy as np _MASK64 = (1 << 64) - 1 _SEED = 2026 _SPLITMIX_INCREMENT = 0x9E3779B97F4A7C15 -_FALLBACK_POLICIES = frozenset({"error", "drop", "keep_original"}) _GROUPING_ROW_CHUNK = 1_000_000 @@ -28,16 +27,16 @@ class CollisionResolutionConfig: """Configuration for within-band SID collision resolution. + Unplaceable overflow items always keep their original SID (over capacity), + so every input item is preserved in the output. + Args: layer_sizes: Cardinality of each SID layer. capacity: Maximum number of retained items in one SID bucket. - fallback_policy: Action when no candidate has a free slot. Supported - values are ``error``, ``drop``, and ``keep_original``. """ layer_sizes: tuple[int, ...] capacity: int - fallback_policy: str def __post_init__(self) -> None: if not self.layer_sizes: @@ -46,11 +45,6 @@ def __post_init__(self) -> None: raise ValueError("all layer_sizes must be positive.") if self.capacity < 1: raise ValueError(f"capacity must be >= 1, got {self.capacity}.") - if self.fallback_policy not in _FALLBACK_POLICIES: - raise ValueError( - "fallback_policy must be one of error, drop, or keep_original, " - f"got {self.fallback_policy!r}." - ) @dataclass(frozen=True) @@ -89,7 +83,7 @@ class CollisionPlan: overflow_origin_last_codes: Origin last-layer code of each overflow row (used to skip a candidate equal to the origin), int64 aligned with ``overflow_rows``. - config: Collision capacity, SID shape, and fallback configuration. + config: Collision capacity and SID shape configuration. """ item_count: int @@ -119,11 +113,10 @@ class CollisionResolutionStats: @dataclass(frozen=True) class CollisionResolutionResult: - """Resolved last codes, indexes, row retention, and diagnostics.""" + """Resolved last codes, slot indexes, and diagnostics.""" resolved_last_codes: np.ndarray slot_indices: np.ndarray - retained_mask: Optional[np.ndarray] unresolved_rows: np.ndarray final_bucket_keys: np.ndarray final_bucket_counts: np.ndarray @@ -263,7 +256,7 @@ def prepare_collision_plan( Args: item_ids: One-dimensional item IDs aligned with ``codes``. codes: Integer SID matrix with shape ``(N, number_of_layers)``. - config: Collision capacity, shape, and fallback configuration. + config: Collision capacity and SID shape configuration. Returns: A compact plan for candidate loading and collision resolution. @@ -368,7 +361,9 @@ def resolve_sid_collisions( """Greedily relocate overflow rows to their first free candidate slot. Candidate rows must be aligned with ``plan.overflow_rows`` and values must - be valid last-layer codebook indices. + be valid last-layer codebook indices. An overflow row that finds no free + candidate slot keeps its original SID (over capacity), so every input row is + preserved in the output. Args: plan: Grouping and overflow plan from :func:`prepare_collision_plan`. @@ -379,11 +374,10 @@ def resolve_sid_collisions( rate-only runs to avoid the grouping arrays and sorts. Returns: - Resolved last-layer codes, slot indices, retention mask, unresolved row - indices, and diagnostics. + Resolved last-layer codes, slot indices, unresolved row indices, and + diagnostics. Raises: - RuntimeError: If a row is unresolved under the ``error`` policy. TypeError: If candidate codes do not use an integer dtype. ValueError: If candidates are not a row-aligned two-dimensional matrix or contain out-of-range last-layer indices. @@ -427,7 +421,6 @@ def resolve_sid_collisions( return CollisionResolutionResult( resolved_last_codes=plan.original_last_codes, slot_indices=plan.initial_slot_indices, - retained_mask=None, unresolved_rows=np.empty(0, dtype=np.int64), final_bucket_keys=final_bucket_keys, final_bucket_counts=final_bucket_counts, @@ -474,23 +467,13 @@ def resolve_sid_collisions( else: unresolved_rows.append(row) - retained_mask = None - fallback_policy = plan.config.fallback_policy - if unresolved_rows and fallback_policy == "error": - preview = ",".join(str(row) for row in unresolved_rows[:10]) - raise RuntimeError( - f"{len(unresolved_rows)} items could not be placed within capacity; " - f"first unresolved row indices: {preview}" - ) - if unresolved_rows and fallback_policy == "drop": - retained_mask = np.ones(plan.item_count, dtype=bool) - retained_mask[unresolved_rows] = False - elif unresolved_rows and fallback_policy == "keep_original": - for row in unresolved_rows: - origin_key = int(plan.bucket_keys[plan.origin_bucket_indices[row]]) - origin_count = get_slot_count(origin_key, 0) + 1 - slot_counts[origin_key] = origin_count - slot_indices[row] = origin_count + # Unplaceable overflow items keep their original SID (over capacity), so + # every input item is preserved in the output. + for row in unresolved_rows: + origin_key = int(plan.bucket_keys[plan.origin_bucket_indices[row]]) + origin_count = get_slot_count(origin_key, 0) + 1 + slot_counts[origin_key] = origin_count + slot_indices[row] = origin_count final_bucket_keys = np.empty(0, dtype=np.int64) final_bucket_counts = np.empty(0, dtype=np.int64) @@ -526,7 +509,6 @@ def resolve_sid_collisions( return CollisionResolutionResult( resolved_last_codes=resolved_last_codes, slot_indices=slot_indices, - retained_mask=retained_mask, unresolved_rows=unresolved_array, final_bucket_keys=final_bucket_keys, final_bucket_counts=final_bucket_counts, @@ -550,7 +532,7 @@ def _scatter_item_grouping( offsets = grouping.offsets if int(offsets[-1]) != row_count: raise RuntimeError( - "bucket counts do not match retained row count: " + "bucket counts do not match grouped row count: " f"{int(offsets[-1])} vs {row_count}." ) written_count = 0 @@ -603,7 +585,7 @@ def row_chunks() -> Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]]: def build_resolved_item_grouping( plan: CollisionPlan, result: CollisionResolutionResult ) -> CodebookItemGrouping: - """Group retained rows by emitted SID and final one-based slot index. + """Group all rows by emitted SID and final one-based slot index. The grouping uses a linear scatter from final one-based slot indices. @@ -612,7 +594,7 @@ def build_resolved_item_grouping( result: Collision output from :func:`resolve_sid_collisions`. Returns: - Sorted final SID keys, bucket counts, and grouped retained row order. + Sorted final SID keys, bucket counts, and grouped row order. Raises: RuntimeError: If grouping metadata was not collected or result bucket @@ -627,13 +609,7 @@ def build_resolved_item_grouping( def row_chunks() -> Iterable[tuple[np.ndarray, np.ndarray, np.ndarray]]: for start in range(0, plan.item_count, _GROUPING_ROW_CHUNK): end = min(start + _GROUPING_ROW_CHUNK, plan.item_count) - if result.retained_mask is None: - rows = np.arange(start, end, dtype=np.int64) - else: - rows = np.flatnonzero(result.retained_mask[start:end]) - rows += start - if rows.size == 0: - continue + rows = np.arange(start, end, dtype=np.int64) emitted_sid_keys = plan.bucket_keys[plan.origin_bucket_indices[rows]] emitted_sid_keys -= plan.original_last_codes[rows] emitted_sid_keys += result.resolved_last_codes[rows] diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index 516db782..b33e73ed 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -43,7 +43,7 @@ def test_golden_candidate_resolution(self) -> None: ], dtype=np.int64, ) - config = CollisionResolutionConfig((2, 4), 2, "error") + config = CollisionResolutionConfig((2, 4), 2) with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): plan = prepare_collision_plan(item_ids, codes, config) @@ -66,7 +66,6 @@ def test_golden_candidate_resolution(self) -> None: np.testing.assert_array_equal(result.final_bucket_keys, [0, 1, 2, 3, 4, 5]) np.testing.assert_array_equal(result.final_bucket_counts, [2, 2, 2, 1, 2, 1]) self.assertTrue(result.grouping_collected) - self.assertIsNone(result.retained_mask) self.assertEqual(result.stats.total_items, 10) self.assertEqual(result.stats.raw_collision_buckets, 2) self.assertEqual(result.stats.final_collision_buckets, 0) @@ -91,38 +90,8 @@ def test_golden_candidate_resolution(self) -> None: resolved_grouping.row_order, [2, 0, 4, 1, 6, 5, 3, 9, 7, 8] ) - def test_error_fallback(self) -> None: - config = CollisionResolutionConfig((2,), 1, "error") - plan = prepare_collision_plan( - np.asarray([0, 1]), np.asarray([[0], [0]]), config - ) - - with self.assertRaisesRegex(RuntimeError, "first unresolved row indices: 1"): - resolve_sid_collisions(plan, np.asarray([[0]], dtype=np.int64)) - - def test_drop_fallback(self) -> None: - config = CollisionResolutionConfig((2,), 1, "drop") - plan = prepare_collision_plan( - np.asarray([0, 1]), np.asarray([[0], [0]]), config - ) - result = resolve_sid_collisions(plan, np.asarray([[0]], dtype=np.int64)) - - np.testing.assert_array_equal(result.resolved_last_codes, [0, 0]) - np.testing.assert_array_equal(result.slot_indices, [1, 2]) - np.testing.assert_array_equal(result.retained_mask, [True, False]) - np.testing.assert_array_equal(result.unresolved_rows, [1]) - self.assertEqual(result.stats.final_collision_buckets, 0) - self.assertEqual(result.stats.unresolved_count, 1) - with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 1): - original_grouping = build_original_item_grouping(plan) - resolved_grouping = build_resolved_item_grouping(plan, result) - np.testing.assert_array_equal(original_grouping.counts, [2]) - np.testing.assert_array_equal(original_grouping.row_order, [0, 1]) - np.testing.assert_array_equal(resolved_grouping.counts, [1]) - np.testing.assert_array_equal(resolved_grouping.row_order, [0]) - def test_keep_original_fallback(self) -> None: - config = CollisionResolutionConfig((2,), 1, "keep_original") + config = CollisionResolutionConfig((2,), 1) plan = prepare_collision_plan( np.asarray([0, 1]), np.asarray([[0], [0]]), config ) @@ -130,7 +99,6 @@ def test_keep_original_fallback(self) -> None: np.testing.assert_array_equal(result.resolved_last_codes, [0, 0]) np.testing.assert_array_equal(result.slot_indices, [1, 2]) - self.assertIsNone(result.retained_mask) np.testing.assert_array_equal(result.unresolved_rows, [1]) self.assertEqual(result.stats.final_collision_buckets, 1) self.assertEqual(result.stats.max_final_bucket_size, 2) @@ -141,7 +109,7 @@ def test_keep_original_fallback(self) -> None: np.testing.assert_array_equal(resolved_grouping.row_order, [0, 1]) def test_no_overflow(self) -> None: - config = CollisionResolutionConfig((2, 4), 2, "error") + config = CollisionResolutionConfig((2, 4), 2) codes = np.asarray([[0, 0], [0, 0], [1, 2]], dtype=np.int64) plan = prepare_collision_plan(np.asarray([0, 1, 2]), codes, config) with ( @@ -161,7 +129,7 @@ def test_no_overflow(self) -> None: self.assertEqual(result.stats.relocated_count, 0) def test_empty_groupings(self) -> None: - config = CollisionResolutionConfig((2, 4), 2, "error") + config = CollisionResolutionConfig((2, 4), 2) plan = prepare_collision_plan( np.empty(0, dtype=np.int64), np.empty((0, 2), dtype=np.int64), @@ -186,7 +154,7 @@ def test_random_candidate_golden_draws(self) -> None: np.testing.assert_array_equal(actual, [[1, 2, 0], [2, 0, 2]]) def test_candidate_matrix_must_be_two_dimensional(self) -> None: - config = CollisionResolutionConfig((2,), 1, "drop") + config = CollisionResolutionConfig((2,), 1) plan = prepare_collision_plan( np.asarray([0, 1]), np.asarray([[0], [0]]), config ) @@ -195,7 +163,7 @@ def test_candidate_matrix_must_be_two_dimensional(self) -> None: resolve_sid_collisions(plan, np.asarray([1], dtype=np.int64)) def test_candidate_matrix_must_align_with_overflow_rows(self) -> None: - config = CollisionResolutionConfig((2,), 1, "drop") + config = CollisionResolutionConfig((2,), 1) plan = prepare_collision_plan( np.asarray([0, 1]), np.asarray([[0], [0]]), config ) @@ -204,7 +172,7 @@ def test_candidate_matrix_must_align_with_overflow_rows(self) -> None: resolve_sid_collisions(plan, np.empty((0, 1), dtype=np.int64)) def test_fixed_width_candidates_can_contend_for_free_slots(self) -> None: - config = CollisionResolutionConfig((6,), 1, "drop") + config = CollisionResolutionConfig((6,), 1) item_ids = np.asarray([0, 1, 2, 3, 4, 5], dtype=np.int64) codes = np.asarray([[0], [0], [1], [3], [4], [4]], dtype=np.int64) plan = prepare_collision_plan(item_ids, codes, config) @@ -225,7 +193,7 @@ def test_fixed_width_candidates_can_contend_for_free_slots(self) -> None: self.assertEqual(result.stats.unresolved_count, 1) def test_rejects_out_of_range_candidate(self) -> None: - config = CollisionResolutionConfig((4,), 1, "drop") + config = CollisionResolutionConfig((4,), 1) item_ids = np.asarray([0, 1, 2], dtype=np.int64) codes = np.asarray([[0], [0], [1]], dtype=np.int64) plan = prepare_collision_plan(item_ids, codes, config) @@ -234,7 +202,7 @@ def test_rejects_out_of_range_candidate(self) -> None: resolve_sid_collisions(plan, np.asarray([[5]], dtype=np.int64)) def test_resolution_can_skip_grouping_metadata(self) -> None: - config = CollisionResolutionConfig((4,), 1, "drop") + config = CollisionResolutionConfig((4,), 1) item_ids = np.asarray([0, 1, 2], dtype=np.int64) codes = np.asarray([[0], [0], [1]], dtype=np.int64) plan = prepare_collision_plan(item_ids, codes, config) @@ -255,7 +223,6 @@ def test_resolution_can_skip_grouping_metadata(self) -> None: result.resolved_last_codes, expected.resolved_last_codes ) np.testing.assert_array_equal(result.slot_indices, expected.slot_indices) - np.testing.assert_array_equal(result.retained_mask, expected.retained_mask) np.testing.assert_array_equal(result.unresolved_rows, expected.unresolved_rows) np.testing.assert_array_equal(result.final_bucket_keys, []) np.testing.assert_array_equal(result.final_bucket_counts, []) From 29ebb32246e5430d294ef0e1ec5d5a7745870b60 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 06:12:30 +0000 Subject: [PATCH 23/29] [refactor] SID collision: fold keep-original into the greedy pass Move the keep-original assignment into the placement loop's else branch (the separate post-loop pass was residue of the removed multi-policy dispatch), fix stale "retained" wording in the resolved-groups CLI help, and bump the version to 1.3.5. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 2 +- tzrec/tools/sid/collision_resolution.py | 13 +++++-------- tzrec/tools/sid/collision_resolution_test.py | 2 +- tzrec/version.py | 2 +- 4 files changed, 8 insertions(+), 11 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 50f0ac13..be893cb2 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -697,7 +697,7 @@ def build_parser() -> argparse.ArgumentParser: "--resolved_sid_groups_output_path", default=None, help=( - "Output grouped retained item IDs keyed by their resolved SID; required " + "Output grouped item IDs keyed by their resolved SID; required " "unless --rate_only is set." ), ) diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py index 751612ea..8370d331 100644 --- a/tzrec/tools/sid/collision_resolution.py +++ b/tzrec/tools/sid/collision_resolution.py @@ -465,15 +465,12 @@ def resolve_sid_collisions( relocated_count += 1 break else: + # No free candidate slot: keep the original SID (over capacity). unresolved_rows.append(row) - - # Unplaceable overflow items keep their original SID (over capacity), so - # every input item is preserved in the output. - for row in unresolved_rows: - origin_key = int(plan.bucket_keys[plan.origin_bucket_indices[row]]) - origin_count = get_slot_count(origin_key, 0) + 1 - slot_counts[origin_key] = origin_count - slot_indices[row] = origin_count + origin_key = int(plan.bucket_keys[plan.origin_bucket_indices[row]]) + origin_count = get_slot_count(origin_key, 0) + 1 + slot_counts[origin_key] = origin_count + slot_indices[row] = origin_count final_bucket_keys = np.empty(0, dtype=np.int64) final_bucket_counts = np.empty(0, dtype=np.int64) diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index b33e73ed..f9b02d1d 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -90,7 +90,7 @@ def test_golden_candidate_resolution(self) -> None: resolved_grouping.row_order, [2, 0, 4, 1, 6, 5, 3, 9, 7, 8] ) - def test_keep_original_fallback(self) -> None: + def test_keep_original_over_capacity(self) -> None: config = CollisionResolutionConfig((2,), 1) plan = prepare_collision_plan( np.asarray([0, 1]), np.asarray([[0], [0]]), config diff --git a/tzrec/version.py b/tzrec/version.py index f963e610..3f74222d 100644 --- a/tzrec/version.py +++ b/tzrec/version.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.3.4" +__version__ = "1.3.5" From b8e2499ade37aed8043e87e43786fb49e76b30c9 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 07:38:01 +0000 Subject: [PATCH 24/29] [refactor] flatten SID candidate_codes to a single list wire format The SID models emitted candidate_codes as nested list> (topk x n_layers), which pyarrow's CSV writer cannot serialize and which forced the offline collision tool to carry two bespoke decoders plus dedicated "|"/";" separators. Flatten it to a single list of topk*n_layers codes in _sid_predictions, so candidates share the same comma-separated wire format and the same _codes_matrix decoder as the primary codes column; the tool recovers per-candidate structure by splitting on the passthrough --codebook length. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/sid_model.py | 2 +- tzrec/models/sid_rqkmeans_test.py | 4 +- tzrec/models/sid_rqvae_test.py | 8 +- tzrec/tools/sid/collision_prevention.py | 96 ++++++++------------ tzrec/tools/sid/collision_prevention_test.py | 22 ++--- 5 files changed, 58 insertions(+), 74 deletions(-) diff --git a/tzrec/models/sid_model.py b/tzrec/models/sid_model.py index 78ba9cb7..684b1daa 100644 --- a/tzrec/models/sid_model.py +++ b/tzrec/models/sid_model.py @@ -109,7 +109,7 @@ def _sid_predictions( """Build SID prediction tensors from a quantizer output.""" predictions = {"codes": quant.cluster_ids} if quant.candidate_codes is not None and quant.candidate_scores is not None: - predictions["candidate_codes"] = quant.candidate_codes + predictions["candidate_codes"] = quant.candidate_codes.flatten(1) predictions["candidate_scores"] = quant.candidate_scores return predictions diff --git a/tzrec/models/sid_rqkmeans_test.py b/tzrec/models/sid_rqkmeans_test.py index 34c7d175..f6cd72f7 100644 --- a/tzrec/models/sid_rqkmeans_test.py +++ b/tzrec/models/sid_rqkmeans_test.py @@ -288,9 +288,9 @@ def test_inference_candidate_output_opt_in(self) -> None: set(preds.keys()), {"codes", "candidate_codes", "candidate_scores"} ) self.assertEqual(preds["codes"].shape, (B, 2)) - self.assertEqual(preds["candidate_codes"].shape, (B, 3, 2)) + self.assertEqual(preds["candidate_codes"].shape, (B, 6)) # (B, topk*n_layers) self.assertEqual(preds["candidate_scores"].shape, (B, 3)) - torch.testing.assert_close(preds["candidate_codes"][:, 0, :], preds["codes"]) + torch.testing.assert_close(preds["candidate_codes"][:, :2], preds["codes"]) def test_candidate_output_topk_exceeds_codebook_raises(self) -> None: """Reject topk above the last-layer codebook size at construction.""" diff --git a/tzrec/models/sid_rqvae_test.py b/tzrec/models/sid_rqvae_test.py index 1717fb1e..b93be676 100644 --- a/tzrec/models/sid_rqvae_test.py +++ b/tzrec/models/sid_rqvae_test.py @@ -247,10 +247,12 @@ def test_rqvae_export_trace_keeps_candidate_output(self) -> None: {"codes", "candidate_codes", "candidate_scores"}, ) self.assertEqual(set(traced_predictions.keys()), set(predictions.keys())) - self.assertEqual(traced_predictions["candidate_codes"].shape, (B, 3, 2)) + self.assertEqual( + traced_predictions["candidate_codes"].shape, (B, 6) + ) # (B, topk*n_layers) self.assertEqual(traced_predictions["candidate_scores"].shape, (B, 3)) torch.testing.assert_close( - traced_predictions["candidate_codes"][:, 0, :], + traced_predictions["candidate_codes"][:, :2], traced_predictions["codes"], ) @@ -268,7 +270,7 @@ def test_rqvae_candidate_scores_sorted(self) -> None: self.assertEqual( set(preds.keys()), {"codes", "candidate_codes", "candidate_scores"} ) - self.assertEqual(preds["candidate_codes"].shape, (B, 3, 2)) + self.assertEqual(preds["candidate_codes"].shape, (B, 6)) # (B, topk*n_layers) self.assertEqual(preds["candidate_scores"].shape, (B, 3)) scores = preds["candidate_scores"] # Slot 0 is the best-scored (nearest) neighbor: minimum score, sorted. diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index be893cb2..0346c26d 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -14,8 +14,10 @@ The runner caps each SID bucket, loads fixed-width last-layer candidates only for overflow items, delegates collision resolution to the pure NumPy core, and writes item-level and grouped SID results through TorchEasyRec readers and -writers. CSV uses compact SID strings and JSON item-ID arrays because Arrow's -CSV writer cannot serialize list columns. +writers. CSV encodes each SID/candidate column as comma-separated codes and +item-ID groups as JSON arrays because Arrow's CSV writer cannot serialize list +columns; ``candidate_codes`` is a flat ``topk * n_layers`` run split by the +``--codebook`` length. The random strategy intentionally preserves the legacy deterministic baseline: it draws with replacement from the full last-layer space. Placement skips an @@ -24,7 +26,8 @@ It is a single-process tool -- launch it with ``python -m`` (no torchrun / process group). The input is a Semantic-ID table from ``tzrec.predict`` (an ``item_id`` column, a ``codes`` ``list`` column, and -- for the default -``--strategy candidate`` -- a ``candidate_codes`` ``list>`` column). +``--strategy candidate`` -- a flat ``candidate_codes`` ``list`` column +(``topk * n_layers`` codes per item)). Example:: @@ -75,8 +78,6 @@ ) from tzrec.utils.logging_util import logger -_CODE_SEP = "|" -_CAND_SEP = ";" _MAP_WRITE_ROWS = 1_000_000 _GROUP_WRITE_ITEMS = 1_000_000 _ARROW_LIST_OFFSET_MAX = int(np.iinfo(np.int32).max) @@ -172,6 +173,8 @@ def __post_init__(self) -> None: if previous_name is not None: raise ValueError(f"{name} must differ from {previous_name}: {path!r}.") seen_paths[path_identity] = name + # Eagerly build the resolution config so its validation (layer_sizes, + # capacity) fails fast here rather than lazily inside run(). _ = self.resolution_config @classmethod @@ -289,7 +292,7 @@ def _codes_matrix(cls, values: pa.Array) -> np.ndarray: array = values.to_numpy(zero_copy_only=False).astype(np.int64, copy=False) return array.reshape(-1, 1) else: - parts = pc.split_pattern(values, _CODE_SEP) + parts = pc.split_pattern(values, ",") row_count = len(parts) flat = pc.cast(parts.flatten(), pa.int64()).to_numpy(zero_copy_only=False) @@ -359,52 +362,31 @@ def _candidate_last_codes(self, plan: CollisionPlan) -> np.ndarray: ) return self._load_candidate_last_codes(plan.overflow_item_ids) - def _candidate_list_matrix(self, values: pa.Array) -> np.ndarray: - """Decode a nested Arrow candidate column into an ``(N, K)`` matrix.""" - row_count = len(values) - if row_count == 0: - return np.empty((0, 0), dtype=np.int64) - inner = values.flatten() - candidate_count = len(inner) // row_count - if len(inner) != row_count * candidate_count: - raise ValueError("ragged candidate_codes: all items must share topk.") - if candidate_count == 0: - return np.empty((row_count, 0), dtype=np.int64) - flat = ( - inner.flatten().to_numpy(zero_copy_only=False).astype(np.int64, copy=False) - ) - tuple_count = row_count * candidate_count - layer_count = flat.shape[0] // tuple_count - if layer_count < 1 or flat.shape[0] != tuple_count * layer_count: + def _candidate_last_matrix(self, values: pa.Array) -> np.ndarray: + """Decode a flat candidate column into its per-candidate last codes. + + ``values`` is decoded like ``codes`` (list, integer, or comma string) into + an ``(N, topk * n_layers)`` matrix, then split into ``topk`` groups of + ``n_layers`` (from ``--codebook``), keeping each group's last-layer code. + + Args: + values: One batch's flat ``candidate_codes`` column. + + Returns: + The last-layer code of every candidate, shape ``(N, topk)``. + + Raises: + ValueError: If the per-row width is not a multiple of ``n_layers``. + """ + flat = self._codes_matrix(values) + per_row = flat.shape[1] + n_layers = len(self._config.layer_sizes) + if per_row % n_layers != 0: raise ValueError( - "ragged candidate_codes: every candidate must share n_layers." + f"candidate_codes width {per_row} is not a multiple of n_layers " + f"{n_layers}." ) - return flat.reshape(row_count, candidate_count, layer_count)[:, :, -1] - - @staticmethod - def _candidate_string_matrix(values: pa.Array) -> np.ndarray: - """Decode compact CSV candidate strings into an ``(N, K)`` matrix.""" - rows: List[List[int]] = [] - for value in values.to_pylist(): - last_codes: List[int] = [] - if value: - for candidate in str(value).split(_CAND_SEP): - parts = [ - int(part) - for part in candidate.strip().split(_CODE_SEP) - if part.strip() - ] - if parts: - last_codes.append(parts[-1]) - rows.append(last_codes) - - widths = {len(row) for row in rows} - if len(widths) > 1: - raise ValueError("ragged candidate_codes: all items must share topk.") - width = widths.pop() if widths else 0 - if width == 0: - return np.empty((len(rows), 0), dtype=np.int64) - return np.asarray(rows, dtype=np.int64).reshape(len(rows), width) + return flat[:, n_layers - 1 :: n_layers] def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarray: """Load fixed-width candidates aligned to ``overflow_item_ids``.""" @@ -417,6 +399,11 @@ def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarra reader = self._make_reader([self._config.item_id_field, field]) for batch in reader.to_batches(): if field not in batch: + logger.warning( + "candidate_codes field %r not present in input batch; " + "stopping candidate scan.", + field, + ) break batch_ids = batch[self._config.item_id_field].to_numpy(zero_copy_only=False) source_rows, target_rows = item_id_lookup.match(batch_ids) @@ -426,10 +413,7 @@ def _load_candidate_last_codes(self, overflow_item_ids: np.ndarray) -> np.ndarra 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())) - if self._is_list_type(selected.type): - batch_candidates = self._candidate_list_matrix(selected) - else: - batch_candidates = self._candidate_string_matrix(selected) + batch_candidates = self._candidate_last_matrix(selected) if candidates is None: candidates = np.empty( (item_count, batch_candidates.shape[1]), dtype=np.int64 @@ -519,9 +503,7 @@ def _codes_column(codes: np.ndarray, is_csv: bool) -> pa.Array: dtype=np.int64, ) ) - return pc.binary_join( - pa.LargeListArray.from_arrays(offsets, values), _CODE_SEP - ) + return pc.binary_join(pa.LargeListArray.from_arrays(offsets, values), ",") if row_count > _ARROW_LIST_OFFSET_MAX // layer_count: raise ValueError("SID output exceeds Arrow list offset capacity.") values = pa.array(codes.reshape(-1)) diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index f87023d4..2382efb3 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -54,9 +54,9 @@ def _parquet(path, item_ids, codes, candidate_codes=None): "codes": pa.array(codes, type=pa.list_(pa.int64())), } if candidate_codes is not None: - cols["candidate_codes"] = pa.array( - candidate_codes, type=pa.list_(pa.list_(pa.int64())) - ) + # flatten each row's [[c0, ..], ..] (topk x n_layers) to a flat list + flat = [[code for cand in row for code in cand] for row in candidate_codes] + cols["candidate_codes"] = pa.array(flat, type=pa.list_(pa.int64())) parquet.write_table(pa.table(cols), path) @@ -340,7 +340,7 @@ def test_csv_code_encoder_joins_all_layers(self) -> None: encoded = CollisionRunner._codes_column(codes, is_csv=True) - self.assertEqual(encoded.to_pylist(), ["1|-2|3", "4|5|6"]) + self.assertEqual(encoded.to_pylist(), ["1,-2,3", "4,5,6"]) def test_csv_backend_within_band(self) -> None: inp = os.path.join(self.test_dir, "in.csv") @@ -349,8 +349,8 @@ def test_csv_backend_within_band(self) -> None: pa.table( { "item_id": ["0", "1", "2", "3", "4"], - "codes": ["0|0"] * 5, - "candidate_codes": ["0|1;0|2;0|3"] * 5, + "codes": ["0,0"] * 5, + "candidate_codes": ["0,1,0,2,0,3"] * 5, } ), inp, @@ -365,7 +365,7 @@ def test_csv_backend_within_band(self) -> None: self.assertEqual(stats.relocated_count, 3) result = csv.read_csv(os.path.join(out, "part-0.csv")) self.assertEqual(result.schema.field("codebook").type, pa.string()) - self.assertTrue(all(s.startswith("0|") for s in result["codebook"].to_pylist())) + self.assertTrue(all(s.startswith("0,") for s in result["codebook"].to_pylist())) original_path, resolved_path = self._group_paths(out) for path in (original_path, resolved_path): @@ -386,7 +386,7 @@ def test_csv_group_itemids_use_json_escaping(self) -> None: pa.table( { "item_id": item_ids, - "codes": ["0|0"] * len(item_ids), + "codes": ["0,0"] * len(item_ids), } ), inp, @@ -417,7 +417,7 @@ def test_parquet_input_with_csv_output(self) -> None: ) result = csv.read_csv(os.path.join(out, "part-0.csv")) self.assertEqual(result.schema.field("codebook").type, pa.string()) - self.assertTrue(all("|" in sid for sid in result["codebook"].to_pylist())) + self.assertTrue(all("," in sid for sid in result["codebook"].to_pylist())) original_path, resolved_path = self._group_paths(out) for path in (original_path, resolved_path): schema = csv.read_csv(os.path.join(path, "part-0.csv")).schema @@ -431,8 +431,8 @@ def test_csv_input_with_parquet_output(self) -> None: pa.table( { "item_id": ["a", "b", "c"], - "codes": ["0|0"] * 3, - "candidate_codes": ["0|1"] * 3, + "codes": ["0,0"] * 3, + "candidate_codes": ["0,1"] * 3, } ), inp, From ff7787241463fd0101c00dd0b7b00c532ab02264 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 09:26:00 +0000 Subject: [PATCH 25/29] [bugfix] SID collision tool: validate inputs against the declared codebook The tool trusted the operator-supplied --codebook without checking the input codes against it. Because a bucket key is band_id * last_size + last_code, an out-of-range last_code aliases into the adjacent band and silently merges two distinct SIDs; a too-small middle-layer size poisons the mixed-radix prefix packing the same way. prepare_collision_plan now range-checks codes against layer_sizes, CollisionResolutionConfig rejects a codebook whose product overflows int64, _load_codes rejects duplicate item IDs, and --codebook parsing rejects empty fields instead of silently dropping them -- each formerly-silent corruption path is now a clear error. Separately, resolve_sid_collisions seeded a Python slot map over every occupied bucket (~one per item) whenever a single overflow row existed; it now scopes that map to overflow-touched bands, keeping cost proportional to the overflow set rather than the whole table. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 17 ++++- tzrec/tools/sid/collision_prevention_test.py | 14 ++++ tzrec/tools/sid/collision_resolution.py | 62 +++++++++++---- tzrec/tools/sid/collision_resolution_test.py | 80 +++++++++++--------- 4 files changed, 120 insertions(+), 53 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index 0346c26d..a0255f91 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -180,9 +180,13 @@ def __post_init__(self) -> None: @classmethod def from_namespace(cls, args: argparse.Namespace) -> "CollisionPreventionConfig": """Build a validated configuration from parsed CLI arguments.""" - layer_sizes = tuple( - int(value) for value in args.codebook.split(",") if value.strip() - ) + try: + layer_sizes = tuple(int(value) for value in args.codebook.split(",")) + except ValueError as err: + raise ValueError( + "--codebook must be 'int,int,...' with no empty fields, got " + f"{args.codebook!r}." + ) from err return cls( input_path=args.input_path, output_path=args.output_path, @@ -345,6 +349,8 @@ 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.") @@ -663,7 +669,10 @@ def _write_sid_groups( def build_parser() -> argparse.ArgumentParser: """Build the collision-prevention command-line parser.""" parser = argparse.ArgumentParser( - description="Prevent SID codebook collisions (vectorized, within-band)." + description=( + "Prevent SID codebook collisions by relocating overflow items " + "within a band." + ) ) parser.add_argument("--input_path", required=True) parser.add_argument("--output_path", required=True) diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 2382efb3..0acbaab7 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -712,6 +712,20 @@ def test_empty_input_raises(self) -> None: with self.assertRaises(ValueError): self._run(inp, out, max_items_per_codebook=2) + def test_duplicate_item_id_raises(self) -> 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) + + def test_empty_codebook_token_raises(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, [0, 1], [[0, 0], [0, 1]]) + with self.assertRaisesRegex(ValueError, "codebook"): + self._run(inp, out, codebook="8,,8") + if __name__ == "__main__": unittest.main() diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py index 8370d331..5050ccd0 100644 --- a/tzrec/tools/sid/collision_resolution.py +++ b/tzrec/tools/sid/collision_resolution.py @@ -11,12 +11,14 @@ """NumPy collision-resolution core for semantic IDs.""" +import math from dataclasses import dataclass from functools import cached_property from typing import Iterable import numpy as np +_INT64_MAX = int(np.iinfo(np.int64).max) _MASK64 = (1 << 64) - 1 _SEED = 2026 _SPLITMIX_INCREMENT = 0x9E3779B97F4A7C15 @@ -45,6 +47,11 @@ def __post_init__(self) -> None: raise ValueError("all layer_sizes must be positive.") if self.capacity < 1: raise ValueError(f"capacity must be >= 1, got {self.capacity}.") + if math.prod(self.layer_sizes) > _INT64_MAX: + raise ValueError( + f"prod(layer_sizes) exceeds int64 range: {self.layer_sizes}; " + "band keys would overflow." + ) @dataclass(frozen=True) @@ -281,6 +288,14 @@ def prepare_collision_plan( f"codes have {layer_count} layers but layer_sizes has " f"{len(config.layer_sizes)}." ) + if codes.size and ( + codes.min() < 0 or np.any(codes >= np.asarray(config.layer_sizes)) + ): + raise ValueError( + "codes contain out-of-range values for layer_sizes " + f"{config.layer_sizes}; check that --codebook matches the model that " + "produced the SID table." + ) original_last_codes = codes[:, -1].astype(np.int64, copy=False) band_ids = _band_ids(codes, config.layer_sizes) @@ -437,7 +452,20 @@ def resolve_sid_collisions( capacity = plan.config.capacity initial_counts = np.minimum(plan.bucket_counts, capacity) - slot_counts = dict(zip(map(int, plan.bucket_keys), map(int, initial_counts))) + # Relocation only ever reads or writes buckets in a band that holds an + # overflow row (destinations stay within the row's band; unplaceable rows + # pile onto their own origin), so seed the mutable slot map from just those + # bands -- every other bucket keeps its capped initial count untouched. This + # bounds the map by overflow-band buckets, not the whole table, which + # dominates cost for sparse-collision inputs at scale. + overflow_band_ids = plan.overflow_bucket_key_prefixes // last_size + in_overflow_band = np.isin(plan.bucket_keys // last_size, overflow_band_ids) + slot_counts = dict( + zip( + plan.bucket_keys[in_overflow_band].tolist(), + initial_counts[in_overflow_band].tolist(), + ) + ) resolved_last_codes = plan.original_last_codes.copy() slot_indices = plan.initial_slot_indices.copy() relocated_count = 0 @@ -474,26 +502,32 @@ def resolve_sid_collisions( final_bucket_keys = np.empty(0, dtype=np.int64) final_bucket_counts = np.empty(0, dtype=np.int64) + untouched_mask = ~in_overflow_band + untouched_counts = initial_counts[untouched_mask] if collect_grouping: - occupancy_counts = np.fromiter( + untouched_keys = plan.bucket_keys[untouched_mask] + band_keys = np.fromiter(slot_counts, dtype=np.int64, count=len(slot_counts)) + band_counts = np.fromiter( slot_counts.values(), dtype=np.int64, count=len(slot_counts) ) - final_collision_buckets = int((occupancy_counts > capacity).sum()) - max_final_bucket_size = ( - int(occupancy_counts.max()) if occupancy_counts.size else 0 - ) - final_bucket_keys = np.fromiter( - slot_counts, dtype=np.int64, count=len(slot_counts) - ) - occupancy_order = np.argsort(final_bucket_keys, kind="stable") - final_bucket_keys = final_bucket_keys[occupancy_order] - final_bucket_counts = occupancy_counts[occupancy_order] + all_keys = np.concatenate((untouched_keys, band_keys)) + all_counts = np.concatenate((untouched_counts, band_counts)) + final_collision_buckets = int((all_counts > capacity).sum()) + max_final_bucket_size = int(all_counts.max()) if all_counts.size else 0 + occupancy_order = np.argsort(all_keys, kind="stable") + final_bucket_keys = all_keys[occupancy_order] + final_bucket_counts = all_counts[occupancy_order] else: + # Untouched buckets never exceed capacity, so only the mutated bands add + # final collisions; the global max still spans both groups. Iterate the + # (overflow-band-scoped) map in Python to skip the grouping arrays/sorts. final_collision_buckets = 0 - max_final_bucket_size = 0 + max_band = 0 for count in slot_counts.values(): final_collision_buckets += count > capacity - max_final_bucket_size = max(max_final_bucket_size, count) + max_band = max(max_band, count) + max_untouched = int(untouched_counts.max()) if untouched_counts.size else 0 + max_final_bucket_size = max(max_untouched, max_band) unresolved_array = np.asarray(unresolved_rows, dtype=np.int64) stats = CollisionResolutionStats( total_items=plan.item_count, diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index f9b02d1d..70cae6f4 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -25,6 +25,14 @@ ) +def _plan(layer_sizes, capacity, item_ids, codes): + return prepare_collision_plan( + np.asarray(item_ids, dtype=np.int64), + np.asarray(codes, dtype=np.int64), + CollisionResolutionConfig(layer_sizes, capacity), + ) + + class CollisionResolutionTest(unittest.TestCase): def test_golden_candidate_resolution(self) -> None: item_ids = np.arange(10, dtype=np.int64) @@ -43,9 +51,8 @@ def test_golden_candidate_resolution(self) -> None: ], dtype=np.int64, ) - config = CollisionResolutionConfig((2, 4), 2) with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): - plan = prepare_collision_plan(item_ids, codes, config) + plan = _plan((2, 4), 2, item_ids, codes) np.testing.assert_array_equal(plan.overflow_rows, [1, 3, 8]) np.testing.assert_array_equal(plan.overflow_item_ids, [1, 3, 8]) @@ -91,10 +98,7 @@ def test_golden_candidate_resolution(self) -> None: ) def test_keep_original_over_capacity(self) -> None: - config = CollisionResolutionConfig((2,), 1) - plan = prepare_collision_plan( - np.asarray([0, 1]), np.asarray([[0], [0]]), config - ) + plan = _plan((2,), 1, [0, 1], [[0], [0]]) result = resolve_sid_collisions(plan, np.asarray([[0]], dtype=np.int64)) np.testing.assert_array_equal(result.resolved_last_codes, [0, 0]) @@ -109,9 +113,7 @@ def test_keep_original_over_capacity(self) -> None: np.testing.assert_array_equal(resolved_grouping.row_order, [0, 1]) def test_no_overflow(self) -> None: - config = CollisionResolutionConfig((2, 4), 2) - codes = np.asarray([[0, 0], [0, 0], [1, 2]], dtype=np.int64) - plan = prepare_collision_plan(np.asarray([0, 1, 2]), codes, config) + plan = _plan((2, 4), 2, [0, 1, 2], [[0, 0], [0, 0], [1, 2]]) with ( mock.patch.object(collision_resolution.np, "fromiter") as fromiter, mock.patch.object(collision_resolution.np, "argsort") as argsort, @@ -129,12 +131,7 @@ def test_no_overflow(self) -> None: self.assertEqual(result.stats.relocated_count, 0) def test_empty_groupings(self) -> None: - config = CollisionResolutionConfig((2, 4), 2) - plan = prepare_collision_plan( - np.empty(0, dtype=np.int64), - np.empty((0, 2), dtype=np.int64), - config, - ) + plan = _plan((2, 4), 2, [], np.empty((0, 2))) result = resolve_sid_collisions(plan, np.empty((0, 0), dtype=np.int64)) for grouping in ( @@ -154,28 +151,19 @@ def test_random_candidate_golden_draws(self) -> None: np.testing.assert_array_equal(actual, [[1, 2, 0], [2, 0, 2]]) def test_candidate_matrix_must_be_two_dimensional(self) -> None: - config = CollisionResolutionConfig((2,), 1) - plan = prepare_collision_plan( - np.asarray([0, 1]), np.asarray([[0], [0]]), config - ) + plan = _plan((2,), 1, [0, 1], [[0], [0]]) with self.assertRaisesRegex(ValueError, "must be 2-D"): resolve_sid_collisions(plan, np.asarray([1], dtype=np.int64)) def test_candidate_matrix_must_align_with_overflow_rows(self) -> None: - config = CollisionResolutionConfig((2,), 1) - plan = prepare_collision_plan( - np.asarray([0, 1]), np.asarray([[0], [0]]), config - ) + plan = _plan((2,), 1, [0, 1], [[0], [0]]) with self.assertRaisesRegex(ValueError, "row-aligned"): resolve_sid_collisions(plan, np.empty((0, 1), dtype=np.int64)) def test_fixed_width_candidates_can_contend_for_free_slots(self) -> None: - config = CollisionResolutionConfig((6,), 1) - item_ids = np.asarray([0, 1, 2, 3, 4, 5], dtype=np.int64) - codes = np.asarray([[0], [0], [1], [3], [4], [4]], dtype=np.int64) - plan = prepare_collision_plan(item_ids, codes, config) + plan = _plan((6,), 1, [0, 1, 2, 3, 4, 5], [[0], [0], [1], [3], [4], [4]]) np.testing.assert_array_equal(plan.overflow_origin_last_codes, [0, 4]) candidates = np.asarray( [ @@ -193,19 +181,13 @@ def test_fixed_width_candidates_can_contend_for_free_slots(self) -> None: self.assertEqual(result.stats.unresolved_count, 1) def test_rejects_out_of_range_candidate(self) -> None: - config = CollisionResolutionConfig((4,), 1) - item_ids = np.asarray([0, 1, 2], dtype=np.int64) - codes = np.asarray([[0], [0], [1]], dtype=np.int64) - plan = prepare_collision_plan(item_ids, codes, config) + plan = _plan((4,), 1, [0, 1, 2], [[0], [0], [1]]) with self.assertRaisesRegex(ValueError, r"must be in \[0, 4\)"): resolve_sid_collisions(plan, np.asarray([[5]], dtype=np.int64)) def test_resolution_can_skip_grouping_metadata(self) -> None: - config = CollisionResolutionConfig((4,), 1) - item_ids = np.asarray([0, 1, 2], dtype=np.int64) - codes = np.asarray([[0], [0], [1]], dtype=np.int64) - plan = prepare_collision_plan(item_ids, codes, config) + plan = _plan((4,), 1, [0, 1, 2], [[0], [0], [1]]) candidates = np.asarray([[2]], dtype=np.int64) expected = resolve_sid_collisions(plan, candidates) @@ -232,6 +214,34 @@ def test_resolution_can_skip_grouping_metadata(self) -> None: with self.assertRaisesRegex(RuntimeError, "metadata was not collected"): build_resolved_item_grouping(plan, result) + def test_rejects_out_of_range_codes(self) -> None: + for bad in ([[0, 0], [0, 4]], [[0, 0], [2, 0]], [[0, 0], [-1, 0]]): + with self.assertRaisesRegex(ValueError, "out-of-range"): + _plan((2, 4), 2, [0, 1], bad) + + def test_rejects_codebook_exceeding_int64(self) -> None: + with self.assertRaisesRegex(ValueError, "int64"): + CollisionResolutionConfig((2**32, 2**32), 1) + + def test_grouping_includes_untouched_bands(self) -> None: + plan = _plan( + (3, 4), 2, [0, 1, 2, 3, 4], [[0, 0], [0, 0], [0, 0], [1, 0], [2, 1]] + ) + candidates = np.asarray([[1, 2, 3]], dtype=np.int64) + + result = resolve_sid_collisions(plan, candidates) + + # Bands 1 (key 4) and 2 (key 9) hold no overflow row; they must still + # appear in the grouping output with their original counts, and the + # relocation must add band-0 key 1. + np.testing.assert_array_equal(result.final_bucket_keys, [0, 1, 4, 9]) + np.testing.assert_array_equal(result.final_bucket_counts, [2, 1, 1, 1]) + self.assertEqual(result.stats.relocated_count, 1) + self.assertEqual(result.stats.unresolved_count, 0) + + rate_only = resolve_sid_collisions(plan, candidates, collect_grouping=False) + self.assertEqual(rate_only.stats, result.stats) + if __name__ == "__main__": unittest.main() From d4646275abbd540ccc2113397838652f3b6b7965 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 09:44:33 +0000 Subject: [PATCH 26/29] [chore] SID collision tool: cover 3-layer band fold and rate-only overflow Add tests for paths the suite left uncovered: a >=3-layer mixed-radix band fold (a wrong middle radix or dropped middle layer now changes the asserted bucket partition), a direct check that the flat candidate column splits at stride n_layers, and a rate-only run with unplaceable overflow whose stats must match the grouping branch. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention_test.py | 44 +++++++++++++++++++ tzrec/tools/sid/collision_resolution_test.py | 46 ++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 0acbaab7..0ea8e70e 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -107,6 +107,50 @@ def test_candidate_reassigns_within_band(self) -> None: # ... and no bucket exceeds the cap. self.assertLessEqual(max(Counter(map(tuple, d["codebook"])).values()), 2) + def test_three_layer_candidate_reassigns_within_band(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + # 3-layer SIDs: band = (code_0, code_1). Three items collide in bucket + # (0, 0, 0); candidates keep the (0, 0) band and vary the last layer, so + # the flat topk*3 candidate column must be split at stride 3. + _parquet( + inp, + [0, 1, 2], + [[0, 0, 0]] * 3, + [[[0, 0, 1], [0, 0, 2], [0, 0, 3]]] * 3, + ) + stats = self._run(inp, out, codebook="2,3,4", max_items_per_codebook=2) + self.assertEqual(stats.total_items, 3) + self.assertEqual(stats.relocated_count, 1) + self.assertEqual(stats.unresolved_count, 0) + self.assertEqual(stats.max_final_bucket_size, 2) + d = self._read_parquet(out) + # The relocated item took the first free candidate last code (1); every + # SID keeps its (0, 0) band prefix. + self.assertEqual( + sorted(tuple(cb) for cb in d["codebook"]), + [(0, 0, 0), (0, 0, 0), (0, 0, 1)], + ) + + def test_candidate_last_matrix_splits_at_n_layers(self) -> None: + # A flat topk*n_layers candidate column splits into topk groups of + # n_layers, keeping each group's LAST code (stride n_layers), not a + # contiguous tail. codebook 2,3,4 -> n_layers=3, so the flat run + # [0,0,1, 0,0,2, 0,0,3] decodes to last codes [1, 2, 3]. + runner = CollisionRunner( + Namespace( + **{ + **_DEFAULTS, + "codebook": "2,3,4", + "rate_only": True, + "input_path": "x", + "output_path": "y", + } + ) + ) + flat = pa.array([[0, 0, 1, 0, 0, 2, 0, 0, 3]], type=pa.list_(pa.int64())) + self.assertEqual(runner._candidate_last_matrix(flat).tolist(), [[1, 2, 3]]) + def test_output_is_list_int64(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index 70cae6f4..d58dad31 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -242,6 +242,52 @@ def test_grouping_includes_untouched_bands(self) -> None: rate_only = resolve_sid_collisions(plan, candidates, collect_grouping=False) self.assertEqual(rate_only.stats, result.stats) + def test_three_layer_band_fold_and_relocation(self) -> None: + # Band = (code_0, code_1) via a mixed-radix fold with middle radix 3. + # Items 0-2 share band (0, 0); item 3 sits at (0, 2) and item 4 at + # (1, 0) -- distinct bands (raw 2 vs 3). A too-small middle radix would + # fold both to 2 and merge them; dropping the middle layer would merge + # item 3 into band (0, 0). Either regression changes the partition below. + plan = _plan( + (2, 3, 4), + 2, + [0, 1, 2, 3, 4], + [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 2, 0], [1, 0, 0]], + ) + # Band (0, 0) holds exactly 3 items (not 4): the middle code kept item 3 + # in its own bucket (key 4) distinct from items 0-2 (key 0). + np.testing.assert_array_equal(plan.bucket_keys, [0, 4, 8]) + np.testing.assert_array_equal(plan.bucket_counts, [3, 1, 1]) + + result = resolve_sid_collisions(plan, np.asarray([[1, 2, 3]], dtype=np.int64)) + + # The single overflow item relocates within band (0, 0) -- only its last + # code changes -- and the two non-overflow bands stay put. + self.assertEqual(result.stats.relocated_count, 1) + self.assertEqual(result.stats.unresolved_count, 0) + self.assertEqual(int(result.resolved_last_codes.sum()), 1) + self.assertEqual(result.resolved_last_codes[3], 0) + self.assertEqual(result.resolved_last_codes[4], 0) + np.testing.assert_array_equal(result.final_bucket_keys, [0, 1, 4, 8]) + np.testing.assert_array_equal(result.final_bucket_counts, [2, 1, 1, 1]) + + def test_rate_only_matches_grouping_with_unplaceable(self) -> None: + # One unplaceable overflow (its only candidate equals the origin) keeps + # bucket (0, 0) over capacity, alongside an untouched band (1, 0). The + # rate-only branch must report the same collision stats as the grouping + # branch -- those numbers are exactly what --rate_only exists to compute. + plan = _plan((2, 2), 1, [0, 1, 2], [[0, 0], [0, 0], [1, 0]]) + candidates = np.asarray([[0]], dtype=np.int64) # only candidate == origin + + grouped = resolve_sid_collisions(plan, candidates, collect_grouping=True) + rate_only = resolve_sid_collisions(plan, candidates, collect_grouping=False) + + self.assertEqual(grouped.stats.unresolved_count, 1) + self.assertEqual(grouped.stats.final_collision_buckets, 1) + self.assertEqual(grouped.stats.max_final_bucket_size, 2) + self.assertEqual(rate_only.stats, grouped.stats) + self.assertFalse(rate_only.grouping_collected) + if __name__ == "__main__": unittest.main() From d822597faf68b4cfbd66dca25022ff242be4ef5e Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 11:25:13 +0000 Subject: [PATCH 27/29] [chore] SID collision tool: compact repeated test code rows Collapse the vertical row-per-line fixtures into row-repetition form (e.g. [[0, 0]] * 4 + [[0, 1]] + [[0, 2]] * 2 + [[1, 0]] * 3). Data is unchanged; the multiplicities now read directly as the asserted bucket counts. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention_test.py | 14 ++---------- tzrec/tools/sid/collision_resolution_test.py | 23 +++----------------- 2 files changed, 5 insertions(+), 32 deletions(-) diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index 0ea8e70e..e092a097 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -208,18 +208,8 @@ def test_group_outputs_follow_sid_and_slot_order(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out") item_ids = list(range(10)) - codes = [ - [0, 0], - [0, 0], - [0, 0], - [0, 0], - [0, 1], - [0, 2], - [0, 2], - [1, 0], - [1, 0], - [1, 0], - ] + # 4 items in bucket (0,0), 1 in (0,1), 2 in (0,2), 3 in (1,0). + codes = [[0, 0]] * 4 + [[0, 1]] + [[0, 2]] * 2 + [[1, 0]] * 3 candidates = [[[0, 0]] for _ in item_ids] candidates[1] = [[0, 0], [0, 1], [0, 2]] candidates[3] = [[0, 1], [0, 2], [0, 3]] diff --git a/tzrec/tools/sid/collision_resolution_test.py b/tzrec/tools/sid/collision_resolution_test.py index d58dad31..8e1017cf 100644 --- a/tzrec/tools/sid/collision_resolution_test.py +++ b/tzrec/tools/sid/collision_resolution_test.py @@ -36,20 +36,9 @@ def _plan(layer_sizes, capacity, item_ids, codes): class CollisionResolutionTest(unittest.TestCase): def test_golden_candidate_resolution(self) -> None: item_ids = np.arange(10, dtype=np.int64) + # 4 items in bucket (0,0), 1 in (0,1), 2 in (0,2), 3 in (1,0). codes = np.asarray( - [ - [0, 0], - [0, 0], - [0, 0], - [0, 0], - [0, 1], - [0, 2], - [0, 2], - [1, 0], - [1, 0], - [1, 0], - ], - dtype=np.int64, + [[0, 0]] * 4 + [[0, 1]] + [[0, 2]] * 2 + [[1, 0]] * 3, dtype=np.int64 ) with mock.patch.object(collision_resolution, "_GROUPING_ROW_CHUNK", 2): plan = _plan((2, 4), 2, item_ids, codes) @@ -165,13 +154,7 @@ def test_candidate_matrix_must_align_with_overflow_rows(self) -> None: def test_fixed_width_candidates_can_contend_for_free_slots(self) -> None: plan = _plan((6,), 1, [0, 1, 2, 3, 4, 5], [[0], [0], [1], [3], [4], [4]]) np.testing.assert_array_equal(plan.overflow_origin_last_codes, [0, 4]) - candidates = np.asarray( - [ - [0, 1, 2, 5], - [4, 2, 1, 3], - ], - dtype=np.int64, - ) + candidates = np.asarray([[0, 1, 2, 5], [4, 2, 1, 3]], dtype=np.int64) result = resolve_sid_collisions(plan, candidates) From 899b9cfd2f0c0f9a129290c23f9977808801ae27 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 16 Jul 2026 11:31:40 +0000 Subject: [PATCH 28/29] [chore] SID collision tool: collapse the _within_bucket_rank return Name the representative-rows expression and return the 5-tuple on one line instead of exploded one-per-line. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_resolution.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tzrec/tools/sid/collision_resolution.py b/tzrec/tools/sid/collision_resolution.py index 5050ccd0..f9e6a4b2 100644 --- a/tzrec/tools/sid/collision_resolution.py +++ b/tzrec/tools/sid/collision_resolution.py @@ -244,13 +244,8 @@ def _within_bucket_rank( bucket_ranks[rows] = ( np.arange(start, end, dtype=np.int64) - first_sorted_rows[sorted_bucket_ids] ) - return ( - bucket_ranks, - bucket_ids, - sorted_rows, - sorted_rows[first_sorted_rows], - bucket_counts, - ) + representatives = sorted_rows[first_sorted_rows] + return bucket_ranks, bucket_ids, sorted_rows, representatives, bucket_counts def prepare_collision_plan( From d54ef58c2297828aa04808731017a6945c4c67d1 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Fri, 17 Jul 2026 02:48:00 +0000 Subject: [PATCH 29/29] [bugfix] SID collision tool: reject output-into-input and multi-process launch Two footguns silently corrupt data. (1) The output path was never checked against the input, but TorchEasyRec writers overwrite in place (Parquet/Csv write part-N into the dir, Odps uses overwrite=True), so pointing an output at the input location destroys the source SID map after it is read; __post_init__ now rejects an output whose write-directory matches the input's. (2) Under torchrun every rank reprocessed the whole input, emitting duplicate local shards or racing ODPS overwrite sessions; run() now raises when WORLD_SIZE != 1 or RANK != 0. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/sid/collision_prevention.py | 39 ++++++++++++++++++++ tzrec/tools/sid/collision_prevention_test.py | 18 +++++++++ 2 files changed, 57 insertions(+) diff --git a/tzrec/tools/sid/collision_prevention.py b/tzrec/tools/sid/collision_prevention.py index a0255f91..c853dab9 100644 --- a/tzrec/tools/sid/collision_prevention.py +++ b/tzrec/tools/sid/collision_prevention.py @@ -95,6 +95,33 @@ def _output_path_identity(path: str) -> _PathIdentity: return "odps", project, schema, table_name, partition_spec +def _input_write_identity(input_path: str) -> _PathIdentity: + """Destination identity a writer must avoid for the given input path. + + Writers emit part files into a directory, so the collision unit is a + directory: the input's own for a local file/glob, else the path itself. + """ + if input_path.startswith("odps://"): + return _output_path_identity(input_path) + base = input_path + if any(ch in input_path for ch in "*?[") or os.path.isfile(input_path): + base = os.path.dirname(input_path) or "." + return os.path.realpath(base) + + +def _require_single_process() -> None: + """Reject multi-process launches; the tool writes one complete map itself.""" + world_size = int(os.environ.get("WORLD_SIZE", "1")) + rank = int(os.environ.get("RANK", "0")) + if world_size != 1 or rank != 0: + raise RuntimeError( + "collision_prevention is single-process; launch it with `python -m` " + f"(not torchrun): got WORLD_SIZE={world_size}, RANK={rank}. Under " + "multiple ranks every rank reprocesses the full input and emits " + "duplicate shards (local) or racing overwrite sessions (ODPS)." + ) + + class _ItemIdLookup: """Map streamed item IDs to positions in a fixed requested-ID array.""" @@ -173,6 +200,17 @@ def __post_init__(self) -> None: if previous_name is not None: raise ValueError(f"{name} must differ from {previous_name}: {path!r}.") seen_paths[path_identity] = name + + # Writers overwrite in place after the full read, so an output in the + # input's location silently destroys the source map. + input_identity = _input_write_identity(self.input_path) + for name, path in named_paths.items(): + if path and _output_path_identity(path) == input_identity: + raise ValueError( + f"{name} would overwrite --input_path ({self.input_path!r}); " + "write outputs to a location outside the input." + ) + # Eagerly build the resolution config so its validation (layer_sizes, # capacity) fails fast here rather than lazily inside run(). _ = self.resolution_config @@ -236,6 +274,7 @@ def __init__( def run(self) -> CollisionResolutionStats: """Read, resolve collisions, and write the resulting SID map.""" + _require_single_process() item_ids, codes = self._load_codes() plan = prepare_collision_plan( item_ids, diff --git a/tzrec/tools/sid/collision_prevention_test.py b/tzrec/tools/sid/collision_prevention_test.py index e092a097..7a63c5b2 100644 --- a/tzrec/tools/sid/collision_prevention_test.py +++ b/tzrec/tools/sid/collision_prevention_test.py @@ -184,6 +184,24 @@ def test_raises_when_overflow_but_no_candidates(self) -> None: with self.assertRaises(ValueError): self._run(inp, out, max_items_per_codebook=2) + def test_output_into_input_location_rejected(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + _parquet(inp, list(range(3)), [[0, 0]] * 3, [[[0, 1]]] * 3) + # output_path is the directory holding the input file -> would overwrite it + with self.assertRaises(ValueError): + self._run(inp, self.test_dir) + + def test_multi_process_launch_rejected(self) -> None: + inp = os.path.join(self.test_dir, "in.parquet") + out = os.path.join(self.test_dir, "out") + _parquet(inp, list(range(3)), [[0, 0]] * 3, [[[0, 1]]] * 3) + os.environ["WORLD_SIZE"] = "2" + try: + with self.assertRaises(RuntimeError): + self._run(inp, out) + finally: + os.environ.pop("WORLD_SIZE", None) + def test_no_overflow_does_not_require_candidates(self) -> None: inp = os.path.join(self.test_dir, "in.parquet") out = os.path.join(self.test_dir, "out")