diff --git a/README.md b/README.md index 531e8b29..959bbad1 100644 --- a/README.md +++ b/README.md @@ -366,7 +366,7 @@ git config core.hooksPath .githooks Workflow: ```bash # 1) edit VERSION -echo 0.51.0 > VERSION +echo 0.51.1 > VERSION # 2) sync manifests uv run python scripts/sync_version.py diff --git a/VERSION b/VERSION index c5d4cee3..1e0c609c 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.51.0 +0.51.1 diff --git a/docs/development.md b/docs/development.md index 631d6269..4a480e55 100644 --- a/docs/development.md +++ b/docs/development.md @@ -69,7 +69,7 @@ Version bump workflow: ```bash # 1) edit VERSION -echo 0.51.0 > VERSION +echo 0.51.1 > VERSION # 2) sync manifests uv run python scripts/sync_version.py diff --git a/pyproject.toml b/pyproject.toml index 816f05bb..2c2edd43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,7 @@ dev = [ "ipykernel", ] rust = [ - "s2and-rust==0.51.0", + "s2and-rust==0.51.1", ] [tool.setuptools.packages.find] diff --git a/s2and/arrow_service_io.py b/s2and/arrow_service_io.py new file mode 100644 index 00000000..210abf6b --- /dev/null +++ b/s2and/arrow_service_io.py @@ -0,0 +1,468 @@ +"""Packaged FeatureBlock conversion helpers for writing Arrow IPC tables. + +These let a service (e.g. the author-disambiguation service) produce the Arrow IPC +artifacts that ``Clusterer.predict_from_arrow_paths`` consumes, without depending on the +unpackaged ``scripts/`` tree. + +- ``write_feature_block_arrow_from_rows`` is the preferred production path: build the tables + straight from service-native rows, no ``ANDData``. +- ``write_feature_block_arrow_from_anddata`` is a compatibility bridge for callers that already + hold an ``ANDData`` (also re-exported by ``scripts/arrow_conversion_helpers.py``). +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import numpy as np + +from s2and.incremental_linking.feature_block_arrow import ( + RAW_PLANNER_ARROW_MAX_RECORD_BATCH_ROWS, + _record_batch_limit_for_table, + write_arrow_ipc_table, + write_feature_block_arrow_tables, +) +from s2and.incremental_linking.feature_block_contract import ( + FeatureBlock, + FeatureBlockPaper, + FeatureBlockPaperAuthor, + FeatureBlockSignature, + _optional_bool, + _optional_int, + _optional_str, + _strict_string_tuple, + build_feature_block_arrow_tables, + filter_cluster_seed_disallows_for_signature_subset, + normalize_cluster_seed_disallow_pairs, + validate_feature_block_columns, +) + + +def write_feature_block_arrow_from_rows( + output_dir: str | Path, + *, + signatures: Sequence[Mapping[str, Any]], + papers: Sequence[Mapping[str, Any]], + paper_authors: Sequence[Mapping[str, Any]], + specter_embeddings: Mapping[str, Sequence[float]] | None = None, + cluster_seeds_require: Mapping[Any, Any] | None = None, + cluster_seeds_disallow: Iterable[tuple[Any, Any]] | None = None, + query_signature_ids: Sequence[Any] = (), + include_empty_cluster_seeds: bool = False, + max_record_batch_rows: Mapping[str, int] | int | None = RAW_PLANNER_ARROW_MAX_RECORD_BATCH_ROWS, + overwrite: bool = True, +) -> dict[str, str]: + """Build the FeatureBlock Arrow tables directly from service-native rows and write them. + + Builds the Arrow tables column-by-column (no per-row FeatureBlock dataclass objects) for + throughput on large blocks, while funneling through the shared validator + (``validate_feature_block_columns``) and schema builder (``build_feature_block_arrow_tables``) + so it enforces the same contract and matches the canonical schema. Row mappings use the Arrow + column names; cluster_seeds_require maps signature_id -> cluster_id (already grouped); + cluster_seeds_disallow is an iterable of signature-id pairs. + """ + + signature_columns = { + "signature_id": [str(r["signature_id"]) for r in signatures], + "paper_id": [str(r["paper_id"]) for r in signatures], + "author_first": [_optional_str(r.get("author_first")) for r in signatures], + "author_middle": [_optional_str(r.get("author_middle")) for r in signatures], + "author_last": [_optional_str(r.get("author_last")) for r in signatures], + "author_suffix": [_optional_str(r.get("author_suffix")) for r in signatures], + "author_affiliations": [ + _strict_string_tuple(r.get("author_affiliations"), field_name="signatures.author_affiliations") + for r in signatures + ], + "author_orcid": [_optional_str(r.get("author_orcid")) for r in signatures], + "author_position": [ + _optional_int(r.get("author_position"), field_name="signatures.author_position") for r in signatures + ], + "author_block": [_optional_str(r.get("author_block")) for r in signatures], + "author_email": [_optional_str(r.get("author_email")) for r in signatures], + "source_author_ids": [ + _strict_string_tuple(r.get("source_author_ids"), field_name="signatures.source_author_ids", skip_none=True) + for r in signatures + ], + } + paper_columns = { + "paper_id": [str(r["paper_id"]) for r in papers], + "title": [_optional_str(r.get("title")) for r in papers], + "abstract": ["Has Abstract" if bool(r.get("has_abstract", False)) else "" for r in papers], + "venue": [_optional_str(r.get("venue")) for r in papers], + "journal_name": [_optional_str(r.get("journal_name")) for r in papers], + "year": [_optional_int(r.get("year"), field_name="papers.year") for r in papers], + "predicted_language": [_optional_str(r.get("predicted_language")) for r in papers], + "is_reliable": [_optional_bool(r.get("is_reliable"), field_name="papers.is_reliable") for r in papers], + } + paper_author_columns = { + "paper_id": [str(r["paper_id"]) for r in paper_authors], + "position": [ + (_optional_int(r.get("position"), field_name="paper_authors.position") or 0) for r in paper_authors + ], + "author_name": [str(r.get("author_name", "") or "") for r in paper_authors], + } + + signature_id_set = set(signature_columns["signature_id"]) + require_pairs = tuple( + (str(sid), str(cid)) for sid, cid in (cluster_seeds_require or {}).items() if str(sid) in signature_id_set + ) + disallow_pairs = filter_cluster_seed_disallows_for_signature_subset( + normalize_cluster_seed_disallow_pairs(cluster_seeds_disallow or ()), + signature_id_set, + ) + specter_paper_ids, specter_matrix = _specter_columns_from_paper_rows(papers, specter_embeddings) + + validate_feature_block_columns( + signature_ids=signature_columns["signature_id"], + signature_paper_ids=signature_columns["paper_id"], + paper_ids=paper_columns["paper_id"], + paper_author_position_keys=zip( + paper_author_columns["paper_id"], paper_author_columns["position"], strict=False + ), + query_signature_ids=[str(value) for value in query_signature_ids], + cluster_seeds_require=require_pairs, + ) + + tables = build_feature_block_arrow_tables( + signature_columns=signature_columns, + paper_columns=paper_columns, + paper_author_columns=paper_author_columns, + cluster_seeds_require=require_pairs, + cluster_seeds_disallow=disallow_pairs, + specter_paper_ids=specter_paper_ids, + specter_embeddings=specter_matrix, + ) + + output_path = Path(output_dir) + paths: dict[str, str] = {} + for name, table in tables.items(): + if ( + name in {"cluster_seeds", "cluster_seed_disallows"} + and table.num_rows == 0 + and not include_empty_cluster_seeds + ): + continue + path = output_path / f"{name}.arrow" + if overwrite or not path.exists(): + write_arrow_ipc_table( + table, path, max_record_batch_rows=_record_batch_limit_for_table(name, max_record_batch_rows) + ) + paths[name] = str(path) + return paths + + +def _specter_columns_from_paper_rows( + paper_rows: Sequence[Mapping[str, Any]], + specter_embeddings: Mapping[str, Sequence[float]] | None, +) -> tuple[tuple[str, ...], np.ndarray | None]: + if not specter_embeddings: + return (), None + paper_ids: list[str] = [] + vectors: list[np.ndarray] = [] + expected_dim: int | None = None + for row in paper_rows: + paper_id = str(row["paper_id"]) + vector = specter_embeddings.get(paper_id) + if vector is None: + continue + array = np.asarray(vector, dtype=np.float32) + if array.ndim != 1: + raise ValueError(f"SPECTER vector for paper_id={paper_id!r} must be 1D, got shape={array.shape}") + if expected_dim is None: + expected_dim = int(array.shape[0]) + elif int(array.shape[0]) != expected_dim: + raise ValueError( + "SPECTER vectors must have equal dimensions: " + f"expected {expected_dim}, got {array.shape[0]} for paper_id={paper_id!r}" + ) + paper_ids.append(paper_id) + vectors.append(array) + if not vectors: + return (), None + return tuple(paper_ids), np.ascontiguousarray(np.vstack(vectors), dtype=np.float32) + + +def feature_block_from_rows( + *, + signatures: Sequence[Mapping[str, Any]], + papers: Sequence[Mapping[str, Any]], + paper_authors: Sequence[Mapping[str, Any]], + specter_embeddings: Mapping[str, Sequence[float]] | None = None, + cluster_seeds_require: Mapping[Any, Any] | None = None, + cluster_seeds_disallow: Iterable[tuple[Any, Any]] | None = None, + query_signature_ids: Sequence[Any] = (), +) -> FeatureBlock: + """Build a FeatureBlock view from service-native row mappings (no ANDData).""" + + block_signatures = tuple(_feature_block_signature_from_row(row) for row in signatures) + signature_id_set = {sig.signature_id for sig in block_signatures} + block_papers = tuple(_feature_block_paper_from_row(row) for row in papers) + block_paper_authors = tuple(_feature_block_paper_author_from_row(row) for row in paper_authors) + + require_pairs = tuple( + (str(signature_id), str(component_id)) + for signature_id, component_id in (cluster_seeds_require or {}).items() + if str(signature_id) in signature_id_set + ) + disallow_pairs = filter_cluster_seed_disallows_for_signature_subset( + normalize_cluster_seed_disallow_pairs(cluster_seeds_disallow or ()), + signature_id_set, + ) + specter_paper_ids, specter_matrix = _specter_from_mapping(block_papers, specter_embeddings) + return FeatureBlock( + signatures=block_signatures, + papers=block_papers, + paper_authors=block_paper_authors, + cluster_seeds_require=require_pairs, + cluster_seeds_disallow=disallow_pairs, + query_signature_ids=tuple(str(value) for value in query_signature_ids), + specter_paper_ids=specter_paper_ids, + specter_embeddings=specter_matrix, + ) + + +def _feature_block_signature_from_row(row: Mapping[str, Any]) -> FeatureBlockSignature: + return FeatureBlockSignature( + signature_id=str(row["signature_id"]), + paper_id=str(row["paper_id"]), + author_first=_optional_str(row.get("author_first")), + author_middle=_optional_str(row.get("author_middle")), + author_last=_optional_str(row.get("author_last")), + author_suffix=_optional_str(row.get("author_suffix")), + author_affiliations=_strict_string_tuple( + row.get("author_affiliations"), field_name="signatures.author_affiliations" + ), + author_orcid=_optional_str(row.get("author_orcid")), + author_position=_optional_int(row.get("author_position"), field_name="signatures.author_position"), + author_block=_optional_str(row.get("author_block")), + author_email=_optional_str(row.get("author_email")), + source_author_ids=_strict_string_tuple( + row.get("source_author_ids"), + field_name="signatures.source_author_ids", + skip_none=True, + ), + ) + + +def _feature_block_paper_from_row(row: Mapping[str, Any]) -> FeatureBlockPaper: + return FeatureBlockPaper( + paper_id=str(row["paper_id"]), + title=_optional_str(row.get("title")), + abstract="Has Abstract" if bool(row.get("has_abstract", False)) else "", + venue=_optional_str(row.get("venue")), + journal_name=_optional_str(row.get("journal_name")), + year=_optional_int(row.get("year"), field_name="papers.year"), + predicted_language=_optional_str(row.get("predicted_language")), + is_reliable=_optional_bool(row.get("is_reliable"), field_name="papers.is_reliable"), + ) + + +def _feature_block_paper_author_from_row(row: Mapping[str, Any]) -> FeatureBlockPaperAuthor: + position = _optional_int(row.get("position"), field_name="paper_authors.position") + return FeatureBlockPaperAuthor( + paper_id=str(row["paper_id"]), + position=0 if position is None else position, + author_name=str(row.get("author_name", "") or ""), + ) + + +def write_feature_block_arrow_from_anddata( + dataset: Any, + output_dir: str | Path, + *, + signature_ids: Sequence[Any] | None = None, + query_signature_ids: Sequence[Any] = (), + cluster_seeds_require: Mapping[Any, Any] | None = None, + include_specter: bool = True, + include_empty_cluster_seeds: bool = False, + max_record_batch_rows: Mapping[str, int] | int | None = RAW_PLANNER_ARROW_MAX_RECORD_BATCH_ROWS, + overwrite: bool = True, +) -> dict[str, str]: + """Build a `FeatureBlock` from `ANDData` and write complete Arrow IPC tables.""" + + feature_block = feature_block_from_anddata( + dataset, + signature_ids=signature_ids, + query_signature_ids=query_signature_ids, + cluster_seeds_require=cluster_seeds_require, + include_specter=include_specter, + ) + return write_feature_block_arrow_tables( + feature_block, + output_dir, + include_empty_cluster_seeds=include_empty_cluster_seeds, + max_record_batch_rows=max_record_batch_rows, + overwrite=overwrite, + ) + + +def feature_block_from_anddata( + dataset: Any, + *, + signature_ids: Sequence[Any] | None = None, + query_signature_ids: Sequence[Any] = (), + cluster_seeds_require: Mapping[Any, Any] | None = None, + include_specter: bool = True, +) -> FeatureBlock: + """Build a `FeatureBlock` view from an existing `ANDData`-like object.""" + + resolved_signature_ids = tuple( + str(value) for value in (dataset.signatures.keys() if signature_ids is None else signature_ids) + ) + signatures = tuple( + _feature_block_signature_from_anddata(dataset, signature_id) for signature_id in resolved_signature_ids + ) + papers = _feature_block_papers_from_anddata(dataset, signatures) + paper_authors = _feature_block_paper_authors_from_papers( + papers_by_id={paper.paper_id: paper for paper in papers}, dataset=dataset + ) + + signature_id_set = set(resolved_signature_ids) + source_cluster_seeds = dict( + getattr(dataset, "cluster_seeds_require", {}) if cluster_seeds_require is None else cluster_seeds_require + ) + require_pairs = tuple( + (str(signature_id), str(component_id)) + for signature_id, component_id in source_cluster_seeds.items() + if str(signature_id) in signature_id_set + ) + disallow_pairs = filter_cluster_seed_disallows_for_signature_subset( + getattr(dataset, "cluster_seeds_disallow", set()), + signature_id_set, + ) + specter_paper_ids, specter_embeddings = _feature_block_specter_from_anddata( + dataset, + papers, + include_specter=include_specter, + ) + return FeatureBlock( + signatures=signatures, + papers=papers, + paper_authors=paper_authors, + cluster_seeds_require=require_pairs, + cluster_seeds_disallow=disallow_pairs, + query_signature_ids=tuple(str(value) for value in query_signature_ids), + specter_paper_ids=specter_paper_ids, + specter_embeddings=specter_embeddings, + ) + + +def _feature_block_signature_from_anddata(dataset: Any, signature_id: str) -> FeatureBlockSignature: + signature = dataset.signatures[str(signature_id)] + return FeatureBlockSignature( + signature_id=str(signature_id), + paper_id=str(signature.paper_id), + author_first=_optional_str(getattr(signature, "author_info_first", None)), + author_middle=_optional_str(getattr(signature, "author_info_middle", None)), + author_last=_optional_str(getattr(signature, "author_info_last", None)), + author_suffix=_optional_str(getattr(signature, "author_info_suffix", None)), + author_affiliations=_strict_string_tuple( + getattr(signature, "author_info_affiliations", None), + field_name="signatures.author_info_affiliations", + ), + author_orcid=_optional_str(getattr(signature, "author_info_orcid", None)), + author_position=_optional_int( + getattr(signature, "author_info_position", None), + field_name="signatures.author_info_position", + ), + author_block=_optional_str(getattr(signature, "author_info_block", None)), + author_email=_optional_str(getattr(signature, "author_info_email", None)), + source_author_ids=_strict_string_tuple( + getattr(signature, "sourced_author_ids", None), + field_name="signatures.sourced_author_ids", + skip_none=True, + ), + ) + + +def _feature_block_papers_from_anddata( + dataset: Any, + signatures: Sequence[FeatureBlockSignature], +) -> tuple[FeatureBlockPaper, ...]: + papers: list[FeatureBlockPaper] = [] + seen: set[str] = set() + for signature in signatures: + paper_id = str(signature.paper_id) + if paper_id in seen: + continue + paper = getattr(dataset, "papers", {}).get(paper_id) + if paper is None: + raise ValueError(f"ANDData papers are missing signature paper_id: {paper_id!r}") + seen.add(paper_id) + papers.append( + FeatureBlockPaper( + paper_id=paper_id, + title=_optional_str(getattr(paper, "title", None)), + abstract="Has Abstract" if bool(getattr(paper, "has_abstract", False)) else "", + venue=_optional_str(getattr(paper, "venue", None)), + journal_name=_optional_str(getattr(paper, "journal_name", None)), + year=_optional_int(getattr(paper, "year", None), field_name="papers.year"), + predicted_language=_optional_str(getattr(paper, "predicted_language", None)), + is_reliable=_optional_bool(getattr(paper, "is_reliable", None), field_name="papers.is_reliable"), + ) + ) + return tuple(papers) + + +def _feature_block_paper_authors_from_papers( + *, + papers_by_id: Mapping[str, FeatureBlockPaper], + dataset: Any, +) -> tuple[FeatureBlockPaperAuthor, ...]: + rows: list[FeatureBlockPaperAuthor] = [] + dataset_papers = getattr(dataset, "papers", {}) + for paper_id in papers_by_id: + paper = dataset_papers[str(paper_id)] + for index, author in enumerate(getattr(paper, "authors", None) or ()): + position = _optional_int(getattr(author, "position", index), field_name="papers.authors.position") + rows.append( + FeatureBlockPaperAuthor( + paper_id=paper_id, + position=index if position is None else position, + author_name=str(getattr(author, "author_name", "") or ""), + ) + ) + return tuple(rows) + + +def _feature_block_specter_from_anddata( + dataset: Any, + papers: Sequence[FeatureBlockPaper], + *, + include_specter: bool, +) -> tuple[tuple[str, ...], np.ndarray | None]: + if not include_specter: + return (), None + return _specter_from_mapping(papers, getattr(dataset, "specter_embeddings", None)) + + +def _specter_from_mapping( + papers: Sequence[FeatureBlockPaper], + specter_embeddings: Mapping[str, Sequence[float]] | None, +) -> tuple[tuple[str, ...], np.ndarray | None]: + if not specter_embeddings: + return (), None + paper_ids: list[str] = [] + vectors: list[np.ndarray] = [] + expected_dim: int | None = None + for paper in papers: + vector = specter_embeddings.get(str(paper.paper_id)) + if vector is None: + continue + array = np.asarray(vector, dtype=np.float32) + if array.ndim != 1: + raise ValueError(f"SPECTER vector for paper_id={paper.paper_id!r} must be 1D, got shape={array.shape}") + if expected_dim is None: + expected_dim = int(array.shape[0]) + elif int(array.shape[0]) != expected_dim: + raise ValueError( + "SPECTER vectors in a FeatureBlock must have equal dimensions: " + f"expected {expected_dim}, got {array.shape[0]} for paper_id={paper.paper_id!r}" + ) + paper_ids.append(str(paper.paper_id)) + vectors.append(array) + if not vectors: + return (), None + return tuple(paper_ids), np.ascontiguousarray(np.vstack(vectors), dtype=np.float32) diff --git a/s2and/incremental_linking/feature_block_contract.py b/s2and/incremental_linking/feature_block_contract.py index 0fda0104..c749c35e 100644 --- a/s2and/incremental_linking/feature_block_contract.py +++ b/s2and/incremental_linking/feature_block_contract.py @@ -200,24 +200,25 @@ class FeatureBlock: def __post_init__(self) -> None: signature_ids = tuple(signature.signature_id for signature in self.signatures) - object.__setattr__(self, "_signature_ids_cache", signature_ids) - if len(set(signature_ids)) != len(signature_ids): - raise ValueError("FeatureBlock signatures must have unique signature_id values") + signature_paper_ids = tuple(signature.paper_id for signature in self.signatures) paper_ids = tuple(paper.paper_id for paper in self.papers) - if len(set(paper_ids)) != len(paper_ids): - raise ValueError("FeatureBlock papers must have unique paper_id values") - seen_paper_author_positions: set[tuple[str, int]] = set() - for author in self.paper_authors: - key = (str(author.paper_id), int(author.position)) - if key in seen_paper_author_positions: - raise ValueError(f"FeatureBlock paper_authors contains duplicate (paper_id, position): {key!r}") - seen_paper_author_positions.add(key) + paper_author_position_keys = [(str(author.paper_id), int(author.position)) for author in self.paper_authors] + query_signature_ids = tuple(str(value) for value in self.query_signature_ids) + require_pairs = tuple( + (str(signature_id), str(component_id)) for signature_id, component_id in self.cluster_seeds_require + ) + + validate_feature_block_columns( + signature_ids=signature_ids, + signature_paper_ids=signature_paper_ids, + paper_ids=paper_ids, + paper_author_position_keys=paper_author_position_keys, + query_signature_ids=query_signature_ids, + cluster_seeds_require=require_pairs, + ) signature_id_set = set(signature_ids) - query_signature_ids = tuple(str(value) for value in self.query_signature_ids) - missing_queries = sorted(set(query_signature_ids) - signature_id_set) - if missing_queries: - raise ValueError(f"query_signature_ids are missing from FeatureBlock signatures: {missing_queries}") + object.__setattr__(self, "_signature_ids_cache", signature_ids) object.__setattr__(self, "query_signature_ids", query_signature_ids) signature_order = FeatureBlockSignatureOrder( signature_ids=signature_ids, @@ -229,26 +230,6 @@ def __post_init__(self) -> None: "_signature_id_to_index_cache", MappingProxyType(signature_order.signature_id_to_index), ) - - require_pair_list: list[tuple[str, str]] = [] - seen_require_signature_ids: set[str] = set() - for signature_id, component_id in self.cluster_seeds_require: - signature_key = str(signature_id) - component_key = str(component_id) - if not signature_key: - raise ValueError("cluster_seeds_require cannot contain empty signature_id values") - if not component_key: - raise ValueError(f"cluster_seeds_require cannot contain empty component_id values: {signature_key!r}") - if signature_key in seen_require_signature_ids: - raise ValueError(f"cluster_seeds_require contains duplicate signature_id: {signature_key!r}") - seen_require_signature_ids.add(signature_key) - require_pair_list.append((signature_key, component_key)) - require_pairs = tuple(require_pair_list) - missing_require = sorted( - signature_id for signature_id, _component_id in require_pairs if signature_id not in signature_id_set - ) - if missing_require: - raise ValueError(f"cluster_seeds_require contains signatures missing from FeatureBlock: {missing_require}") object.__setattr__(self, "cluster_seeds_require", require_pairs) disallow_pairs = normalize_cluster_seed_disallow_pairs( @@ -303,87 +284,173 @@ def seed_component_members(self) -> dict[str, tuple[str, ...]]: def to_arrow_tables(self) -> dict[str, Any]: """Return Arrow tables for the current Rust raw-candidate schema.""" - import pyarrow as pa - - tables: dict[str, Any] = { - "signatures": pa.table( - { - "signature_id": pa.array([row.signature_id for row in self.signatures], type=pa.string()), - "paper_id": pa.array([row.paper_id for row in self.signatures], type=pa.string()), - "author_first": pa.array([row.author_first for row in self.signatures], type=pa.string()), - "author_middle": pa.array([row.author_middle for row in self.signatures], type=pa.string()), - "author_last": pa.array([row.author_last for row in self.signatures], type=pa.string()), - "author_suffix": pa.array([row.author_suffix for row in self.signatures], type=pa.string()), - "author_affiliations": pa.array( - [list(row.author_affiliations) for row in self.signatures], - type=pa.list_(pa.string()), - ), - "author_orcid": pa.array([row.author_orcid for row in self.signatures], type=pa.string()), - "author_position": pa.array([row.author_position for row in self.signatures], type=pa.int64()), - "author_block": pa.array([row.author_block for row in self.signatures], type=pa.string()), - "author_email": pa.array([row.author_email for row in self.signatures], type=pa.string()), - "source_author_ids": pa.array( - [list(row.source_author_ids) for row in self.signatures], - type=pa.list_(pa.string()), - ), - } - ), - "papers": pa.table( - { - "paper_id": pa.array([row.paper_id for row in self.papers], type=pa.string()), - "title": pa.array([row.title for row in self.papers], type=pa.string()), - "abstract": pa.array([row.abstract for row in self.papers], type=pa.string()), - "venue": pa.array([row.venue for row in self.papers], type=pa.string()), - "journal_name": pa.array([row.journal_name for row in self.papers], type=pa.string()), - "year": pa.array([row.year for row in self.papers], type=pa.int64()), - "predicted_language": pa.array( - [row.predicted_language for row in self.papers], - type=pa.string(), - ), - "is_reliable": pa.array([row.is_reliable for row in self.papers], type=pa.bool_()), - } - ), - "paper_authors": pa.table( - { - "paper_id": pa.array([row.paper_id for row in self.paper_authors], type=pa.string()), - "position": pa.array([row.position for row in self.paper_authors], type=pa.int64()), - "author_name": pa.array([row.author_name for row in self.paper_authors], type=pa.string()), - } - ), - "cluster_seeds": pa.table( - { - "signature_id": pa.array( - [signature_id for signature_id, _component_id in self.cluster_seeds_require], - type=pa.string(), - ), - "cluster_id": pa.array( - [component_id for _signature_id, component_id in self.cluster_seeds_require], - type=pa.string(), - ), - } - ), - "cluster_seed_disallows": pa.table( - { - "signature_id_1": pa.array( - [left for left, _right in self.cluster_seeds_disallow], - type=pa.string(), - ), - "signature_id_2": pa.array( - [right for _left, right in self.cluster_seeds_disallow], - type=pa.string(), - ), - } - ), - } - if self.specter_embeddings is not None: - flat = pa.array(np.ravel(self.specter_embeddings), type=pa.float32()) - tables["specter"] = pa.table( - { - "paper_id": list(self.specter_paper_ids), - "embedding": pa.FixedSizeListArray.from_arrays(flat, int(self.specter_embeddings.shape[1])), - } - ) - return tables + return build_feature_block_arrow_tables( + signature_columns={ + "signature_id": [row.signature_id for row in self.signatures], + "paper_id": [row.paper_id for row in self.signatures], + "author_first": [row.author_first for row in self.signatures], + "author_middle": [row.author_middle for row in self.signatures], + "author_last": [row.author_last for row in self.signatures], + "author_suffix": [row.author_suffix for row in self.signatures], + "author_affiliations": [list(row.author_affiliations) for row in self.signatures], + "author_orcid": [row.author_orcid for row in self.signatures], + "author_position": [row.author_position for row in self.signatures], + "author_block": [row.author_block for row in self.signatures], + "author_email": [row.author_email for row in self.signatures], + "source_author_ids": [list(row.source_author_ids) for row in self.signatures], + }, + paper_columns={ + "paper_id": [row.paper_id for row in self.papers], + "title": [row.title for row in self.papers], + "abstract": [row.abstract for row in self.papers], + "venue": [row.venue for row in self.papers], + "journal_name": [row.journal_name for row in self.papers], + "year": [row.year for row in self.papers], + "predicted_language": [row.predicted_language for row in self.papers], + "is_reliable": [row.is_reliable for row in self.papers], + }, + paper_author_columns={ + "paper_id": [row.paper_id for row in self.paper_authors], + "position": [row.position for row in self.paper_authors], + "author_name": [row.author_name for row in self.paper_authors], + }, + cluster_seeds_require=self.cluster_seeds_require, + cluster_seeds_disallow=self.cluster_seeds_disallow, + specter_paper_ids=self.specter_paper_ids, + specter_embeddings=self.specter_embeddings, + ) + + +def validate_feature_block_columns( + *, + signature_ids: Sequence[str], + signature_paper_ids: Sequence[str], + paper_ids: Sequence[str], + paper_author_position_keys: Iterable[tuple[str, int]], + query_signature_ids: Sequence[str], + cluster_seeds_require: Sequence[tuple[str, str]], +) -> None: + """Structural FeatureBlock invariants, shared by ``FeatureBlock.__post_init__`` and the + columnar service path so both reject the same malformed inputs. + """ + + for signature_id, signature_paper_id in zip(signature_ids, signature_paper_ids, strict=False): + if not str(signature_id): + raise ValueError("FeatureBlockSignature.signature_id must be non-empty") + if not str(signature_paper_id): + raise ValueError(f"FeatureBlockSignature.paper_id must be non-empty for {signature_id!r}") + if len(set(signature_ids)) != len(signature_ids): + raise ValueError("FeatureBlock signatures must have unique signature_id values") + for paper_id in paper_ids: + if not str(paper_id): + raise ValueError("FeatureBlockPaper.paper_id must be non-empty") + if len(set(paper_ids)) != len(paper_ids): + raise ValueError("FeatureBlock papers must have unique paper_id values") + seen_positions: set[tuple[str, int]] = set() + for key in paper_author_position_keys: + if not str(key[0]): + raise ValueError("FeatureBlockPaperAuthor.paper_id must be non-empty") + if key in seen_positions: + raise ValueError(f"FeatureBlock paper_authors contains duplicate (paper_id, position): {key!r}") + seen_positions.add(key) + signature_id_set = set(signature_ids) + missing_queries = sorted({str(value) for value in query_signature_ids} - signature_id_set) + if missing_queries: + raise ValueError(f"query_signature_ids are missing from FeatureBlock signatures: {missing_queries}") + seen_require: set[str] = set() + for signature_id, component_id in cluster_seeds_require: + signature_key = str(signature_id) + if not signature_key: + raise ValueError("cluster_seeds_require cannot contain empty signature_id values") + if not str(component_id): + raise ValueError(f"cluster_seeds_require cannot contain empty component_id values: {signature_key!r}") + if signature_key in seen_require: + raise ValueError(f"cluster_seeds_require contains duplicate signature_id: {signature_key!r}") + seen_require.add(signature_key) + missing_require = sorted( + str(signature_id) for signature_id, _c in cluster_seeds_require if str(signature_id) not in signature_id_set + ) + if missing_require: + raise ValueError(f"cluster_seeds_require contains signatures missing from FeatureBlock: {missing_require}") + + +def build_feature_block_arrow_tables( + *, + signature_columns: Mapping[str, Sequence[Any]], + paper_columns: Mapping[str, Sequence[Any]], + paper_author_columns: Mapping[str, Sequence[Any]], + cluster_seeds_require: Sequence[tuple[str, str]], + cluster_seeds_disallow: Sequence[tuple[str, str]], + specter_paper_ids: Sequence[str], + specter_embeddings: np.ndarray | None, +) -> dict[str, Any]: + """Single source of truth for the FeatureBlock Arrow schema. + + Builds the Rust raw-candidate Arrow tables from already-prepared column sequences. Used by + both ``FeatureBlock.to_arrow_tables`` (columns extracted from row objects) and the columnar + service path (columns built directly from request rows), so the schema lives in one place. + """ + + import pyarrow as pa + + tables: dict[str, Any] = { + "signatures": pa.table( + { + "signature_id": pa.array(signature_columns["signature_id"], type=pa.string()), + "paper_id": pa.array(signature_columns["paper_id"], type=pa.string()), + "author_first": pa.array(signature_columns["author_first"], type=pa.string()), + "author_middle": pa.array(signature_columns["author_middle"], type=pa.string()), + "author_last": pa.array(signature_columns["author_last"], type=pa.string()), + "author_suffix": pa.array(signature_columns["author_suffix"], type=pa.string()), + "author_affiliations": pa.array(signature_columns["author_affiliations"], type=pa.list_(pa.string())), + "author_orcid": pa.array(signature_columns["author_orcid"], type=pa.string()), + "author_position": pa.array(signature_columns["author_position"], type=pa.int64()), + "author_block": pa.array(signature_columns["author_block"], type=pa.string()), + "author_email": pa.array(signature_columns["author_email"], type=pa.string()), + "source_author_ids": pa.array(signature_columns["source_author_ids"], type=pa.list_(pa.string())), + } + ), + "papers": pa.table( + { + "paper_id": pa.array(paper_columns["paper_id"], type=pa.string()), + "title": pa.array(paper_columns["title"], type=pa.string()), + "abstract": pa.array(paper_columns["abstract"], type=pa.string()), + "venue": pa.array(paper_columns["venue"], type=pa.string()), + "journal_name": pa.array(paper_columns["journal_name"], type=pa.string()), + "year": pa.array(paper_columns["year"], type=pa.int64()), + "predicted_language": pa.array(paper_columns["predicted_language"], type=pa.string()), + "is_reliable": pa.array(paper_columns["is_reliable"], type=pa.bool_()), + } + ), + "paper_authors": pa.table( + { + "paper_id": pa.array(paper_author_columns["paper_id"], type=pa.string()), + "position": pa.array(paper_author_columns["position"], type=pa.int64()), + "author_name": pa.array(paper_author_columns["author_name"], type=pa.string()), + } + ), + "cluster_seeds": pa.table( + { + "signature_id": pa.array([s for s, _c in cluster_seeds_require], type=pa.string()), + "cluster_id": pa.array([c for _s, c in cluster_seeds_require], type=pa.string()), + } + ), + "cluster_seed_disallows": pa.table( + { + "signature_id_1": pa.array([left for left, _r in cluster_seeds_disallow], type=pa.string()), + "signature_id_2": pa.array([right for _l, right in cluster_seeds_disallow], type=pa.string()), + } + ), + } + if specter_embeddings is not None: + flat = pa.array(np.ravel(specter_embeddings), type=pa.float32()) + tables["specter"] = pa.table( + { + "paper_id": list(specter_paper_ids), + "embedding": pa.FixedSizeListArray.from_arrays(flat, int(specter_embeddings.shape[1])), + } + ) + return tables def feature_block_signature_order_from_raw_candidate_plan(plan: Mapping[str, Any]) -> FeatureBlockSignatureOrder: diff --git a/s2and/runtime.py b/s2and/runtime.py index b0ad2aac..903a4621 100644 --- a/s2and/runtime.py +++ b/s2and/runtime.py @@ -19,7 +19,7 @@ _STARTUP_WARNING_EMITTED = False _STARTUP_WARNING_LOCK = threading.Lock() -MIN_SUPPORTED_RUST_EXTENSION_VERSION = (0, 51, 0) +MIN_SUPPORTED_RUST_EXTENSION_VERSION = (0, 51, 1) _CORE_REQUIRED_FEATURIZER_MARKERS = ( "from_arrow_paths", "signature_ids", diff --git a/s2and_rust/Cargo.lock b/s2and_rust/Cargo.lock index 035f6b8a..446c4ed7 100644 --- a/s2and_rust/Cargo.lock +++ b/s2and_rust/Cargo.lock @@ -31,7 +31,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" dependencies = [ - "memchr 2.8.1", + "memchr 2.8.2", ] [[package]] @@ -205,10 +205,10 @@ dependencies = [ "arrow-data", "arrow-schema", "arrow-select", - "memchr 2.8.1", + "memchr 2.8.2", "num", - "regex 1.12.3", - "regex-syntax 0.8.10", + "regex 1.12.4", + "regex-syntax 0.8.11", ] [[package]] @@ -246,15 +246,15 @@ checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" [[package]] name = "cc" -version = "1.2.62" +version = "1.2.65" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98" +checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", "shlex", @@ -277,9 +277,9 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "chrono" -version = "0.4.44" +version = "0.4.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" dependencies = [ "iana-time-zone", "num-traits", @@ -513,13 +513,12 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "js-sys" -version = "0.3.99" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "142bc4740e452c1e57ade0cbc129f139c9093e354346f0872ef985f4f5cf5f11" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" dependencies = [ "cfg-if", "futures-util", - "once_cell", "wasm-bindgen", ] @@ -610,9 +609,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "log" -version = "0.4.30" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "616ec5685824bcc94416c6d4a7a446eea774a31efd7062c8480ba6fd06d7a6e5" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "matrixmultiply" @@ -635,15 +634,15 @@ dependencies = [ [[package]] name = "memchr" -version = "2.8.1" +version = "2.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b947ae49db0d222b1dbc6b113ce7248a3fc3a6ca21b696717bfc000ba4484d8" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" [[package]] name = "memmap2" -version = "0.9.10" +version = "0.9.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", ] @@ -863,9 +862,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.45" +version = "1.0.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" dependencies = [ "proc-macro2", ] @@ -917,14 +916,14 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.3" +version = "1.12.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" dependencies = [ "aho-corasick 1.1.4", - "memchr 2.8.1", + "memchr 2.8.2", "regex-automata", - "regex-syntax 0.8.10", + "regex-syntax 0.8.11", ] [[package]] @@ -934,8 +933,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" dependencies = [ "aho-corasick 1.1.4", - "memchr 2.8.1", - "regex-syntax 0.8.10", + "memchr 2.8.2", + "regex-syntax 0.8.11", ] [[package]] @@ -946,9 +945,9 @@ checksum = "f9ec002c35e86791825ed294b50008eea9ddfc8def4420124fbc6b08db834957" [[package]] name = "regex-syntax" -version = "0.8.10" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "rustc-hash" @@ -985,7 +984,7 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "s2and_rust" -version = "0.51.0" +version = "0.51.1" dependencies = [ "arrow", "cld2", @@ -1042,7 +1041,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ "itoa", - "memchr 2.8.1", + "memchr 2.8.2", "serde", "serde_core", "zmij", @@ -1050,9 +1049,9 @@ dependencies = [ [[package]] name = "shlex" -version = "1.3.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" [[package]] name = "slab" @@ -1062,9 +1061,9 @@ checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" [[package]] name = "syn" -version = "2.0.117" +version = "2.0.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" dependencies = [ "proc-macro2", "quote", @@ -1152,18 +1151,18 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] name = "wasip2" -version = "1.0.3+wasi-0.2.9" +version = "1.0.4+wasi-0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" dependencies = [ "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ed04576f974d2b2fba0f38c51dbc5518011e38c36bf1143164be765528fd409" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" dependencies = [ "cfg-if", "once_cell", @@ -1174,9 +1173,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "916151b09da36bd82f6615cbf3a419e2f0ba23a03c6160e8e92eb6bd4aa1dec6" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1184,9 +1183,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "299047362ccbfce148b67ab7e73349f77748e00c8296f9542adfad2ad82c5c5e" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" dependencies = [ "bumpalo", "proc-macro2", @@ -1197,9 +1196,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a929b2c61f11ba3e9bc35b50c1f25cb38e0e892c0c231ae2b8cf78d5dad4437" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" dependencies = [ "unicode-ident", ] @@ -1283,18 +1282,18 @@ checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" [[package]] name = "zerocopy" -version = "0.8.49" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bce33a6288fa3f072a8c2c7d0f2fdbb90e28298f0135c1f99b96c3db2efcc60b" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.49" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd425244944f4ab65ccff928e7323354c5a018c75838362fdce749dfad2ee1e" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" dependencies = [ "proc-macro2", "quote", diff --git a/s2and_rust/Cargo.toml b/s2and_rust/Cargo.toml index 90759321..33de510d 100644 --- a/s2and_rust/Cargo.toml +++ b/s2and_rust/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "s2and_rust" -version = "0.51.0" +version = "0.51.1" edition = "2021" [lib] diff --git a/s2and_rust/pyproject.toml b/s2and_rust/pyproject.toml index f15fd11a..bbcc02ea 100644 --- a/s2and_rust/pyproject.toml +++ b/s2and_rust/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "s2and-rust" -version = "0.51.0" +version = "0.51.1" description = "Rust extension for s2and" readme = "README.md" requires-python = ">=3.11,<3.14" diff --git a/scripts/arrow_conversion_helpers.py b/scripts/arrow_conversion_helpers.py index 52d8ef10..68231dea 100644 --- a/scripts/arrow_conversion_helpers.py +++ b/scripts/arrow_conversion_helpers.py @@ -1,219 +1,18 @@ -"""FeatureBlock conversion helpers for writing Arrow tables from ANDData.""" +"""FeatureBlock conversion helpers for writing Arrow tables from ANDData. -from __future__ import annotations - -from collections.abc import Mapping, Sequence -from pathlib import Path -from typing import Any +The implementation now lives in the packaged module ``s2and.arrow_service_io`` so +services that install s2and (without the ``scripts/`` tree) can import it. This +module re-exports the public helpers for backwards compatibility. +""" -import numpy as np +from __future__ import annotations -from s2and.incremental_linking.feature_block_arrow import ( - RAW_PLANNER_ARROW_MAX_RECORD_BATCH_ROWS, - write_feature_block_arrow_tables, +from s2and.arrow_service_io import ( + feature_block_from_anddata, + write_feature_block_arrow_from_anddata, ) -from s2and.incremental_linking.feature_block_contract import ( - FeatureBlock, - FeatureBlockPaper, - FeatureBlockPaperAuthor, - FeatureBlockSignature, - _optional_bool, - _optional_int, - _optional_str, - _strict_string_tuple, - filter_cluster_seed_disallows_for_signature_subset, -) - - -def write_feature_block_arrow_from_anddata( - dataset: Any, - output_dir: str | Path, - *, - signature_ids: Sequence[Any] | None = None, - query_signature_ids: Sequence[Any] = (), - cluster_seeds_require: Mapping[Any, Any] | None = None, - include_specter: bool = True, - include_empty_cluster_seeds: bool = False, - max_record_batch_rows: Mapping[str, int] | int | None = RAW_PLANNER_ARROW_MAX_RECORD_BATCH_ROWS, - overwrite: bool = True, -) -> dict[str, str]: - """Build a `FeatureBlock` from `ANDData` and write complete Arrow IPC tables.""" - - feature_block = feature_block_from_anddata( - dataset, - signature_ids=signature_ids, - query_signature_ids=query_signature_ids, - cluster_seeds_require=cluster_seeds_require, - include_specter=include_specter, - ) - return write_feature_block_arrow_tables( - feature_block, - output_dir, - include_empty_cluster_seeds=include_empty_cluster_seeds, - max_record_batch_rows=max_record_batch_rows, - overwrite=overwrite, - ) - - -def feature_block_from_anddata( - dataset: Any, - *, - signature_ids: Sequence[Any] | None = None, - query_signature_ids: Sequence[Any] = (), - cluster_seeds_require: Mapping[Any, Any] | None = None, - include_specter: bool = True, -) -> FeatureBlock: - """Build a `FeatureBlock` view from an existing `ANDData`-like object.""" - - resolved_signature_ids = tuple( - str(value) for value in (dataset.signatures.keys() if signature_ids is None else signature_ids) - ) - signatures = tuple( - _feature_block_signature_from_anddata(dataset, signature_id) for signature_id in resolved_signature_ids - ) - papers = _feature_block_papers_from_anddata(dataset, signatures) - paper_authors = _feature_block_paper_authors_from_papers( - papers_by_id={paper.paper_id: paper for paper in papers}, dataset=dataset - ) - - signature_id_set = set(resolved_signature_ids) - source_cluster_seeds = dict( - getattr(dataset, "cluster_seeds_require", {}) if cluster_seeds_require is None else cluster_seeds_require - ) - require_pairs = tuple( - (str(signature_id), str(component_id)) - for signature_id, component_id in source_cluster_seeds.items() - if str(signature_id) in signature_id_set - ) - disallow_pairs = filter_cluster_seed_disallows_for_signature_subset( - getattr(dataset, "cluster_seeds_disallow", set()), - signature_id_set, - ) - specter_paper_ids, specter_embeddings = _feature_block_specter_from_anddata( - dataset, - papers, - include_specter=include_specter, - ) - return FeatureBlock( - signatures=signatures, - papers=papers, - paper_authors=paper_authors, - cluster_seeds_require=require_pairs, - cluster_seeds_disallow=disallow_pairs, - query_signature_ids=tuple(str(value) for value in query_signature_ids), - specter_paper_ids=specter_paper_ids, - specter_embeddings=specter_embeddings, - ) - - -def _feature_block_signature_from_anddata(dataset: Any, signature_id: str) -> FeatureBlockSignature: - signature = dataset.signatures[str(signature_id)] - return FeatureBlockSignature( - signature_id=str(signature_id), - paper_id=str(signature.paper_id), - author_first=_optional_str(getattr(signature, "author_info_first", None)), - author_middle=_optional_str(getattr(signature, "author_info_middle", None)), - author_last=_optional_str(getattr(signature, "author_info_last", None)), - author_suffix=_optional_str(getattr(signature, "author_info_suffix", None)), - author_affiliations=_strict_string_tuple( - getattr(signature, "author_info_affiliations", None), - field_name="signatures.author_info_affiliations", - ), - author_orcid=_optional_str(getattr(signature, "author_info_orcid", None)), - author_position=_optional_int( - getattr(signature, "author_info_position", None), - field_name="signatures.author_info_position", - ), - author_block=_optional_str(getattr(signature, "author_info_block", None)), - author_email=_optional_str(getattr(signature, "author_info_email", None)), - source_author_ids=_strict_string_tuple( - getattr(signature, "sourced_author_ids", None), - field_name="signatures.sourced_author_ids", - skip_none=True, - ), - ) - - -def _feature_block_papers_from_anddata( - dataset: Any, - signatures: Sequence[FeatureBlockSignature], -) -> tuple[FeatureBlockPaper, ...]: - papers: list[FeatureBlockPaper] = [] - seen: set[str] = set() - for signature in signatures: - paper_id = str(signature.paper_id) - if paper_id in seen: - continue - paper = getattr(dataset, "papers", {}).get(paper_id) - if paper is None: - raise ValueError(f"ANDData papers are missing signature paper_id: {paper_id!r}") - seen.add(paper_id) - papers.append( - FeatureBlockPaper( - paper_id=paper_id, - title=_optional_str(getattr(paper, "title", None)), - abstract="Has Abstract" if bool(getattr(paper, "has_abstract", False)) else "", - venue=_optional_str(getattr(paper, "venue", None)), - journal_name=_optional_str(getattr(paper, "journal_name", None)), - year=_optional_int(getattr(paper, "year", None), field_name="papers.year"), - predicted_language=_optional_str(getattr(paper, "predicted_language", None)), - is_reliable=_optional_bool(getattr(paper, "is_reliable", None), field_name="papers.is_reliable"), - ) - ) - return tuple(papers) - - -def _feature_block_paper_authors_from_papers( - *, - papers_by_id: Mapping[str, FeatureBlockPaper], - dataset: Any, -) -> tuple[FeatureBlockPaperAuthor, ...]: - rows: list[FeatureBlockPaperAuthor] = [] - dataset_papers = getattr(dataset, "papers", {}) - for paper_id in papers_by_id: - paper = dataset_papers[str(paper_id)] - for index, author in enumerate(getattr(paper, "authors", None) or ()): - position = _optional_int(getattr(author, "position", index), field_name="papers.authors.position") - rows.append( - FeatureBlockPaperAuthor( - paper_id=paper_id, - position=index if position is None else position, - author_name=str(getattr(author, "author_name", "") or ""), - ) - ) - return tuple(rows) - -def _feature_block_specter_from_anddata( - dataset: Any, - papers: Sequence[FeatureBlockPaper], - *, - include_specter: bool, -) -> tuple[tuple[str, ...], np.ndarray | None]: - if not include_specter: - return (), None - specter = getattr(dataset, "specter_embeddings", None) - if not specter: - return (), None - paper_ids: list[str] = [] - vectors: list[np.ndarray] = [] - expected_dim: int | None = None - for paper in papers: - vector = specter.get(str(paper.paper_id)) - if vector is None: - continue - array = np.asarray(vector, dtype=np.float32) - if array.ndim != 1: - raise ValueError(f"SPECTER vector for paper_id={paper.paper_id!r} must be 1D, got shape={array.shape}") - if expected_dim is None: - expected_dim = int(array.shape[0]) - elif int(array.shape[0]) != expected_dim: - raise ValueError( - "SPECTER vectors in a FeatureBlock must have equal dimensions: " - f"expected {expected_dim}, got {array.shape[0]} for paper_id={paper.paper_id!r}" - ) - paper_ids.append(str(paper.paper_id)) - vectors.append(array) - if not vectors: - return (), None - return tuple(paper_ids), np.ascontiguousarray(np.vstack(vectors), dtype=np.float32) +__all__ = [ + "feature_block_from_anddata", + "write_feature_block_arrow_from_anddata", +] diff --git a/tests/test_arrow_production_boundary.py b/tests/test_arrow_production_boundary.py index 5444a4c4..23898918 100644 --- a/tests/test_arrow_production_boundary.py +++ b/tests/test_arrow_production_boundary.py @@ -8,10 +8,14 @@ import s2and.feature_port as feature_port import s2and.model as model_module +import s2and.runtime as runtime from s2and.arrow_inputs import MissingArrowArtifactError from s2and.featurizer import FeaturizationInfo from s2and.model import Clusterer +# Derived from the guard so fixtures track the lockstep version bump instead of going stale. +_SUPPORTED_VERSION = runtime.min_supported_rust_extension_version_string() + class ArrowOnlyRustFeaturizer: calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] @@ -39,7 +43,7 @@ def update_signature_name_counts(self, signatures: dict[str, Any]) -> int: class ArrowOnlyRustModule: - __version__ = "0.51.0" + __version__ = _SUPPORTED_VERSION RustFeaturizer = ArrowOnlyRustFeaturizer diff --git a/tests/test_feature_block.py b/tests/test_feature_block.py index d68938f2..a2e5b3ea 100644 --- a/tests/test_feature_block.py +++ b/tests/test_feature_block.py @@ -10,6 +10,7 @@ import s2and.incremental_linking.feature_block_arrow as feature_block_arrow_module from s2and.arrow_inputs import MissingArrowArtifactError +from s2and.arrow_service_io import feature_block_from_rows, write_feature_block_arrow_from_rows from s2and.data import ANDData, NameCounts from s2and.featurizer import FeaturizationInfo from s2and.incremental_linking.feature_block import ( @@ -2233,3 +2234,251 @@ def test_raw_candidate_plan_bridge_rejects_missing_schema_version() -> None: raw_plan, feature_block_signature_order=feature_block_signature_order_from_raw_candidate_plan(_raw_plan()), ) + + +def _rows_from_feature_block(fb: Any) -> tuple[list[dict], list[dict], list[dict], dict | None]: + sig_rows = [ + { + "signature_id": s.signature_id, + "paper_id": s.paper_id, + "author_first": s.author_first, + "author_middle": s.author_middle, + "author_last": s.author_last, + "author_suffix": s.author_suffix, + "author_affiliations": list(s.author_affiliations), + "author_orcid": s.author_orcid, + "author_position": s.author_position, + "author_block": s.author_block, + "author_email": s.author_email, + "source_author_ids": list(s.source_author_ids), + } + for s in fb.signatures + ] + paper_rows = [ + { + "paper_id": p.paper_id, + "title": p.title, + "has_abstract": p.abstract == "Has Abstract", + "venue": p.venue, + "journal_name": p.journal_name, + "year": p.year, + "predicted_language": p.predicted_language, + "is_reliable": p.is_reliable, + } + for p in fb.papers + ] + paper_author_rows = [ + {"paper_id": a.paper_id, "position": a.position, "author_name": a.author_name} for a in fb.paper_authors + ] + specter = ( + {pid: fb.specter_embeddings[i] for i, pid in enumerate(fb.specter_paper_ids)} + if fb.specter_embeddings is not None + else None + ) + return sig_rows, paper_rows, paper_author_rows, specter + + +def test_feature_block_from_rows_matches_anddata() -> None: + # Parity (identity): rows extracted from the ANDData-built FeatureBlock rebuild an identical + # FeatureBlock via the row path, so the preferred path is contract-equivalent to the bridge. + dataset = _tiny_anddata() + dataset.cluster_seeds_disallow = set() + fb_anddata = feature_block_from_anddata(dataset, signature_ids=["q", "s1"], query_signature_ids=["q"]) + + sig_rows, paper_rows, paper_author_rows, specter = _rows_from_feature_block(fb_anddata) + fb_rows = feature_block_from_rows( + signatures=sig_rows, + papers=paper_rows, + paper_authors=paper_author_rows, + specter_embeddings=specter, + cluster_seeds_require=dict(fb_anddata.cluster_seeds_require), + cluster_seeds_disallow=fb_anddata.cluster_seeds_disallow, + query_signature_ids=fb_anddata.query_signature_ids, + ) + + assert fb_rows.signatures == fb_anddata.signatures + assert fb_rows.papers == fb_anddata.papers + assert fb_rows.paper_authors == fb_anddata.paper_authors + assert set(fb_rows.cluster_seeds_require) == set(fb_anddata.cluster_seeds_require) + assert set(fb_rows.cluster_seeds_disallow) == set(fb_anddata.cluster_seeds_disallow) + assert fb_rows.query_signature_ids == fb_anddata.query_signature_ids + assert fb_rows.specter_paper_ids == fb_anddata.specter_paper_ids + np.testing.assert_array_equal(fb_rows.specter_embeddings, fb_anddata.specter_embeddings) + + +def test_write_feature_block_arrow_from_rows_round_trips(tmp_path: Path) -> None: + pa = pytest.importorskip("pyarrow") + + sig_rows = [ + { + "signature_id": "q", + "paper_id": "p_q", + "author_first": "Ada", + "author_middle": "", + "author_last": "Lovelace", + "author_suffix": "", + "author_affiliations": ["Lab"], + "author_orcid": "0000-0000-0000-0001", + "author_position": 0, + "author_block": "a lovelace", + "author_email": "", + "source_author_ids": ["src-q"], + }, + { + "signature_id": "s1", + "paper_id": "p1", + "author_first": "Ada", + "author_middle": "", + "author_last": "Lovelace", + "author_suffix": "", + "author_affiliations": [], + "author_orcid": None, + "author_position": 0, + "author_block": "a lovelace", + "author_email": "", + "source_author_ids": [], + }, + ] + paper_rows = [ + {"paper_id": "p_q", "title": "Notes", "has_abstract": True, "venue": "RS", "journal_name": "", "year": 1843}, + {"paper_id": "p1", "title": "Notes", "has_abstract": False, "venue": "RS", "journal_name": "", "year": 1843}, + ] + paper_author_rows = [ + {"paper_id": "p_q", "position": 0, "author_name": "ada lovelace"}, + {"paper_id": "p1", "position": 0, "author_name": "ada lovelace"}, + ] + specter = {"p_q": np.asarray([1.0, 0.0], dtype=np.float32), "p1": np.asarray([0.5, 0.5], dtype=np.float32)} + + table_paths = write_feature_block_arrow_from_rows( + tmp_path, + signatures=sig_rows, + papers=paper_rows, + paper_authors=paper_author_rows, + specter_embeddings=specter, + cluster_seeds_require={"s1": "c_ada"}, + cluster_seeds_disallow=[("q", "s1")], + include_empty_cluster_seeds=True, + ) + indexed_paths, _metrics = feature_block_arrow_module.write_raw_arrow_batch_lookup_indexes(table_paths, tmp_path) + + def read(name: str) -> Any: + with pa.memory_map(table_paths[name]) as source: + return pa.ipc.open_file(source).read_all() + + signatures = read("signatures") + assert signatures.column("signature_id").to_pylist() == ["q", "s1"] + assert signatures.column("author_orcid").to_pylist() == ["0000-0000-0000-0001", None] + papers = read("papers") + assert papers.column("paper_id").to_pylist() == ["p_q", "p1"] + assert papers.column("abstract").to_pylist() == ["Has Abstract", ""] + assert read("paper_authors").num_rows == 2 + assert read("specter").column("paper_id").to_pylist() == ["p_q", "p1"] + cluster_seeds = read("cluster_seeds") + assert cluster_seeds.column("signature_id").to_pylist() == ["s1"] + assert cluster_seeds.column("cluster_id").to_pylist() == ["c_ada"] + assert read("cluster_seed_disallows").num_rows == 1 + for key in ( + "signatures_batch_index", + "papers_batch_index", + "paper_authors_batch_index", + "specter_batch_index", + ): + assert key in indexed_paths and Path(indexed_paths[key]).exists() + + +def _sig_row(signature_id: str, paper_id: str, **overrides: Any) -> dict[str, Any]: + row = { + "signature_id": signature_id, + "paper_id": paper_id, + "author_first": "Ada", + "author_last": "Lovelace", + "author_affiliations": [], + "author_position": 0, + "author_block": "a lovelace", + } + row.update(overrides) + return row + + +def _paper_row(paper_id: str, **overrides: Any) -> dict[str, Any]: + row = {"paper_id": paper_id, "title": "T", "has_abstract": False, "venue": "RS", "year": 2020} + row.update(overrides) + return row + + +def test_write_feature_block_arrow_from_rows_enforces_contract(tmp_path: Path) -> None: + # The columnar row path must reject the same malformed inputs the dataclass path does + # (it funnels through validate_feature_block_columns + the strict coercers). + pytest.importorskip("pyarrow") + + def write(**kwargs: Any) -> Any: + return write_feature_block_arrow_from_rows(tmp_path, paper_authors=[], **kwargs) + + with pytest.raises(ValueError, match="unique signature_id"): + write(signatures=[_sig_row("q", "p1"), _sig_row("q", "p2")], papers=[_paper_row("p1"), _paper_row("p2")]) + with pytest.raises(ValueError, match="unique paper_id"): + write(signatures=[_sig_row("q", "p1")], papers=[_paper_row("p1"), _paper_row("p1")]) + with pytest.raises(ValueError, match="cannot contain null"): + write(signatures=[_sig_row("q", "p1", author_affiliations=["ok", None])], papers=[_paper_row("p1")]) + with pytest.raises(ValueError, match="must be a sequence"): + write(signatures=[_sig_row("q", "p1", author_affiliations="oops")], papers=[_paper_row("p1")]) + with pytest.raises(ValueError, match="query_signature_ids"): + write(signatures=[_sig_row("q", "p1")], papers=[_paper_row("p1")], query_signature_ids=["missing"]) + with pytest.raises(ValueError, match="FeatureBlockSignature.signature_id must be non-empty"): + write(signatures=[_sig_row("", "p1")], papers=[_paper_row("p1")]) + with pytest.raises(ValueError, match="FeatureBlockSignature.paper_id must be non-empty"): + write(signatures=[_sig_row("q", "")], papers=[_paper_row("p1")]) + with pytest.raises(ValueError, match="FeatureBlockPaper.paper_id must be non-empty"): + write(signatures=[_sig_row("q", "p1")], papers=[_paper_row("")]) + with pytest.raises(ValueError, match="FeatureBlockPaperAuthor.paper_id must be non-empty"): + write_feature_block_arrow_from_rows( + tmp_path, + signatures=[_sig_row("q", "p1")], + papers=[_paper_row("p1")], + paper_authors=[{"paper_id": "", "position": 0, "author_name": "a"}], + ) + with pytest.raises(ValueError, match="duplicate"): + write_feature_block_arrow_from_rows( + tmp_path, + signatures=[_sig_row("q", "p1")], + papers=[_paper_row("p1")], + paper_authors=[ + {"paper_id": "p1", "position": 0, "author_name": "a"}, + {"paper_id": "p1", "position": 0, "author_name": "b"}, + ], + ) + + +def test_write_feature_block_arrow_from_rows_matches_to_arrow_tables(tmp_path: Path) -> None: + # Pin the columnar output against the canonical FeatureBlock.to_arrow_tables schema/values. + pa = pytest.importorskip("pyarrow") + + signatures = [_sig_row("q", "p1", author_orcid="0000-0000-0000-0001", source_author_ids=["src-q"])] + papers = [_paper_row("p1", has_abstract=True)] + paper_authors = [{"paper_id": "p1", "position": 0, "author_name": "ada lovelace"}] + specter = {"p1": np.asarray([1.0, 0.0], dtype=np.float32)} + + table_paths = write_feature_block_arrow_from_rows( + tmp_path, + signatures=signatures, + papers=papers, + paper_authors=paper_authors, + specter_embeddings=specter, + ) + canonical = feature_block_from_rows( + signatures=signatures, + papers=papers, + paper_authors=paper_authors, + specter_embeddings=specter, + ).to_arrow_tables() + + def read(path: str) -> Any: + with pa.memory_map(path) as source: + return pa.ipc.open_file(source).read_all() + + for name, table in canonical.items(): + if table.num_rows == 0: + continue + written = read(table_paths[name]) + assert written.schema == table.schema, name + assert written.to_pydict() == table.to_pydict(), name diff --git a/tests/test_feature_port_cache.py b/tests/test_feature_port_cache.py index cda2b875..e5ecd489 100644 --- a/tests/test_feature_port_cache.py +++ b/tests/test_feature_port_cache.py @@ -14,6 +14,9 @@ from s2and.incremental_linking.feature_block import write_name_counts_index from tests.helpers import patch_tiny_name_counts_loader +# Derived from the guard so fixtures track the lockstep version bump instead of going stale. +_SUPPORTED_VERSION = runtime.min_supported_rust_extension_version_string() + def _missing_module(name: str) -> ModuleNotFoundError: return ModuleNotFoundError(f"No module named {name!r}", name=name) @@ -64,7 +67,7 @@ def cluster_seeds_disallow(self): class DummyRustModule: - __version__ = "0.51.0" + __version__ = _SUPPORTED_VERSION RustFeaturizer = DummyRustFeaturizer @@ -270,7 +273,7 @@ def from_dataset(cls, dataset, _require_value, _disallow_value, _num_threads=Non return cls(dataset.name) class EmbeddingRustModule: - __version__ = "0.51.0" + __version__ = _SUPPORTED_VERSION RustFeaturizer = EmbeddingRustFeaturizer monkeypatch.setattr(feature_port, "s2and_rust", EmbeddingRustModule) @@ -437,7 +440,7 @@ def from_arrow_paths(cls, *_args, **_kwargs): raise AssertionError("from_arrow_paths should not be called for invalid paths") class ArrowRustModule: - __version__ = "0.51.0" + __version__ = _SUPPORTED_VERSION RustFeaturizer = ArrowRustFeaturizer monkeypatch.setattr(feature_port, "s2and_rust", ArrowRustModule) @@ -461,7 +464,7 @@ def from_arrow_paths(cls, paths, _signature_ids, _name_tuples, *_args): return cls("arrow") class ArrowRustModule: - __version__ = "0.51.0" + __version__ = _SUPPORTED_VERSION RustFeaturizer = ArrowRustFeaturizer monkeypatch.setattr(feature_port, "s2and_rust", ArrowRustModule) @@ -507,7 +510,7 @@ def from_arrow_paths(cls, *_args, **_kwargs): return cls("arrow") class ArrowRustModule: - __version__ = "0.51.0" + __version__ = _SUPPORTED_VERSION RustFeaturizer = ArrowRustFeaturizer monkeypatch.setattr(feature_port, "s2and_rust", ArrowRustModule) diff --git a/uv.lock b/uv.lock index d048440a..0f4c5d8c 100644 --- a/uv.lock +++ b/uv.lock @@ -1241,7 +1241,7 @@ provides-extras = ["dev", "rust"] [[package]] name = "s2and-rust" -version = "0.51.0" +version = "0.51.1" source = { directory = "s2and_rust" } [[package]]