Buckets:

skaltab03
/

STRING-bucket

Files

xet

skaltab03/STRING-bucket / scripts /prepare_hf_dataset.py

skaltab03

4 months ago

download

raw

23.4 kB

#!/usr/bin/env python3
"""Convert raw STRING downloads into Hugging Face Data Viewer-friendly Parquet.
The converter streams the raw files, assigns rows to deterministic
train/validation/test splits, and writes sharded Parquet files under:
    data/<config>/<split>-00000.parquet
It is designed for very large STRING files: chunks are parsed in a bounded
process pool and written incrementally, so the complete dataset is never loaded
into memory.
"""

import argparse
import gzip
import io
import json
import logging
import shutil
import subprocess
import sys
import time
import zlib
from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
from contextlib import contextmanager
from pathlib import Path
from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Set, Tuple


TABLE_ORDER = (
    "species",
    "protein_info",
    "protein_aliases",
    "protein_sequences",
    "protein_links",
)

RAW_FILES = {
    "species": "species.v12.0.txt",
    "protein_info": "protein.info.v12.0.txt.gz",
    "protein_aliases": "protein.aliases.v12.0.txt.gz",
    "protein_sequences": "protein.sequences.v12.0.fa.gz",
    "protein_links": "protein.links.full.v12.0.txt.gz",
}

COLUMNS = {
    "species": (
        "taxon_id",
        "string_type",
        "string_name_compact",
        "official_name_ncbi",
        "domain",
    ),
    "protein_info": (
        "string_protein_id",
        "taxon_id",
        "preferred_name",
        "protein_size",
        "annotation",
    ),
    "protein_aliases": (
        "string_protein_id",
        "taxon_id",
        "alias",
        "source",
    ),
    "protein_sequences": (
        "string_protein_id",
        "taxon_id",
        "sequence",
        "sequence_length",
    ),
    "protein_links": (
        "protein1",
        "protein2",
        "taxon_id",
        "neighborhood",
        "neighborhood_transferred",
        "fusion",
        "cooccurence",
        "homology",
        "coexpression",
        "coexpression_transferred",
        "experiments",
        "experiments_transferred",
        "database",
        "database_transferred",
        "textmining",
        "textmining_transferred",
        "combined_score",
    ),
}

TYPE_NAMES = {
    "species": (
        "int32",
        "string",
        "string",
        "string",
        "string",
    ),
    "protein_info": (
        "string",
        "int32",
        "string",
        "int32",
        "string",
    ),
    "protein_aliases": (
        "string",
        "int32",
        "string",
        "string",
    ),
    "protein_sequences": (
        "string",
        "int32",
        "string",
        "int32",
    ),
    "protein_links": (
        "string",
        "string",
        "int32",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
        "int16",
    ),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Post-process STRING v12.0 raw files into sharded Parquet configs."
    )
    parser.add_argument("--raw-dir", type=Path, default=Path("v12.0"))
    parser.add_argument("--output-dir", type=Path, default=Path("data"))
    parser.add_argument(
        "--tables",
        default="all",
        help="Comma-separated table list, or 'all'. Valid tables: %s" % ", ".join(TABLE_ORDER),
    )
    parser.add_argument("--num-proc", type=int, default=32)
    parser.add_argument(
        "--max-in-flight",
        type=int,
        default=None,
        help="Maximum parser chunks queued at once. Defaults to --num-proc.",
    )
    parser.add_argument(
        "--rows-per-chunk",
        type=int,
        default=100_000,
        help="Rows sent to one worker at a time. Lower this if RAM is tight.",
    )
    parser.add_argument("--train-ratio", type=float, default=0.98)
    parser.add_argument("--validation-ratio", type=float, default=0.01)
    parser.add_argument("--test-ratio", type=float, default=0.01)
    parser.add_argument(
        "--split-seed",
        default="string-v12.0",
        help="Stable salt used by the row hash split assignment.",
    )
    parser.add_argument(
        "--link-min-combined-score",
        type=int,
        default=None,
        help="Optional filter for protein_links. Keeps rows with combined_score >= this value.",
    )
    parser.add_argument(
        "--max-rows-per-table",
        type=int,
        default=None,
        help="Optional raw-row cap per table, useful for creating a small preview dataset.",
    )
    parser.add_argument(
        "--compression",
        default="zstd",
        choices=("zstd", "snappy", "gzip", "brotli", "none"),
        help="Parquet compression codec.",
    )
    parser.add_argument(
        "--decompressor",
        default="auto",
        choices=("auto", "python", "gzip", "pigz"),
        help="How to stream .gz files. auto uses pigz when installed, otherwise Python gzip.",
    )
    parser.add_argument(
        "--decompressor-proc",
        type=int,
        default=4,
        help="Threads passed to pigz when --decompressor is pigz/auto and pigz is available.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Delete existing output subdirectories for selected tables before writing.",
    )
    parser.add_argument("--log-every", type=int, default=25)
    return parser.parse_args()


def resolve_tables(value: str) -> List[str]:
    if value == "all":
        return list(TABLE_ORDER)
    selected = [item.strip() for item in value.split(",") if item.strip()]
    unknown = sorted(set(selected) - set(TABLE_ORDER))
    if unknown:
        raise ValueError("Unknown table(s): %s" % ", ".join(unknown))
    return selected


def split_cutoffs(train_ratio: float, validation_ratio: float, test_ratio: float) -> Tuple[int, int]:
    total = train_ratio + validation_ratio + test_ratio
    if abs(total - 1.0) > 1e-9:
        raise ValueError("Split ratios must sum to 1.0, got %.8f" % total)
    train_cutoff = int(train_ratio * (2**32))
    validation_cutoff = int((train_ratio + validation_ratio) * (2**32))
    return train_cutoff, validation_cutoff


def seed_crc(seed: str) -> int:
    return zlib.crc32(seed.encode("utf-8")) & 0xFFFFFFFF


def assign_split(key: str, seed_value: int, cutoffs: Tuple[int, int]) -> str:
    value = zlib.crc32(key.encode("utf-8"), seed_value) & 0xFFFFFFFF
    if value < cutoffs[0]:
        return "train"
    if value < cutoffs[1]:
        return "validation"
    return "test"


def taxon_from_protein_id(protein_id: str) -> int:
    return int(protein_id.split(".", 1)[0])


@contextmanager
def open_text(path: Path, decompressor: str, decompressor_proc: int) -> Iterator[io.TextIOBase]:
    if path.suffix != ".gz":
        with path.open("rt", encoding="utf-8", newline="") as handle:
            yield handle
        return

    if decompressor == "pigz" and not shutil.which("pigz"):
        raise RuntimeError("Requested --decompressor pigz, but pigz is not installed")

    use_pigz = decompressor == "pigz" or (decompressor == "auto" and shutil.which("pigz"))
    use_gzip_cmd = decompressor == "gzip"

    if use_pigz:
        cmd = ["pigz", "-dc", "-p", str(max(1, decompressor_proc)), str(path)]
    elif use_gzip_cmd:
        cmd = ["gzip", "-cd", str(path)]
    else:
        with gzip.open(path, "rt", encoding="utf-8", newline="") as handle:
            yield handle
        return

    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    assert proc.stdout is not None
    handle = io.TextIOWrapper(proc.stdout, encoding="utf-8", newline="")
    try:
        yield handle
    finally:
        handle.close()
        return_code = proc.wait()
        if return_code != 0:
            raise RuntimeError("Decompressor failed with exit code %s: %s" % (return_code, " ".join(cmd)))


def iter_text_chunks(
    path: Path,
    rows_per_chunk: int,
    max_rows: Optional[int],
    decompressor: str,
    decompressor_proc: int,
) -> Iterator[Tuple[int, List[str]]]:
    with open_text(path, decompressor, decompressor_proc) as handle:
        header = next(handle, None)
        if header is None:
            return

        chunk: List[str] = []
        chunk_id = 0
        seen = 0
        for line in handle:
            if max_rows is not None and seen >= max_rows:
                break
            chunk.append(line)
            seen += 1
            if len(chunk) >= rows_per_chunk:
                yield chunk_id, chunk
                chunk_id += 1
                chunk = []

        if chunk:
            yield chunk_id, chunk


def iter_fasta_records(
    path: Path,
    max_rows: Optional[int],
    decompressor: str,
    decompressor_proc: int,
) -> Iterator[Tuple[str, str]]:
    yielded = 0
    protein_id: Optional[str] = None
    parts: List[str] = []

    with open_text(path, decompressor, decompressor_proc) as handle:
        for raw_line in handle:
            line = raw_line.strip()
            if not line:
                continue
            if line.startswith(">"):
                if protein_id is not None:
                    yield protein_id, "".join(parts)
                    yielded += 1
                    if max_rows is not None and yielded >= max_rows:
                        return
                protein_id = line[1:].split(None, 1)[0]
                parts = []
            else:
                parts.append(line)

        if protein_id is not None and (max_rows is None or yielded < max_rows):
            yield protein_id, "".join(parts)


def iter_sequence_chunks(
    path: Path,
    rows_per_chunk: int,
    max_rows: Optional[int],
    decompressor: str,
    decompressor_proc: int,
) -> Iterator[Tuple[int, List[Tuple[str, str]]]]:
    chunk: List[Tuple[str, str]] = []
    chunk_id = 0
    for record in iter_fasta_records(path, max_rows, decompressor, decompressor_proc):
        chunk.append(record)
        if len(chunk) >= rows_per_chunk:
            yield chunk_id, chunk
            chunk_id += 1
            chunk = []
    if chunk:
        yield chunk_id, chunk


def empty_buckets(table: str) -> Dict[str, List[List[object]]]:
    width = len(COLUMNS[table])
    return {
        "train": [[] for _ in range(width)],
        "validation": [[] for _ in range(width)],
        "test": [[] for _ in range(width)],
    }


def append_row(bucket: List[List[object]], values: Sequence[object]) -> None:
    for index, value in enumerate(values):
        bucket[index].append(value)


def parse_species_line(line: str) -> Tuple[str, Tuple[object, ...]]:
    fields = line.rstrip("\n").split("\t")
    if len(fields) != 5:
        raise ValueError("expected 5 fields")
    taxon_id = int(fields[0])
    return str(taxon_id), (taxon_id, fields[1], fields[2], fields[3], fields[4])


def parse_protein_info_line(line: str) -> Tuple[str, Tuple[object, ...]]:
    fields = line.rstrip("\n").split("\t", 3)
    if len(fields) != 4:
        raise ValueError("expected 4 fields")
    protein_id = fields[0]
    taxon_id = taxon_from_protein_id(protein_id)
    return protein_id, (protein_id, taxon_id, fields[1], int(fields[2]), fields[3])


def parse_alias_line(line: str) -> Tuple[str, Tuple[object, ...]]:
    fields = line.rstrip("\n").split("\t", 2)
    if len(fields) != 3:
        raise ValueError("expected 3 fields")
    protein_id = fields[0]
    taxon_id = taxon_from_protein_id(protein_id)
    key = "%s\t%s\t%s" % (protein_id, fields[1], fields[2])
    return key, (protein_id, taxon_id, fields[1], fields[2])


def parse_link_line(line: str, min_combined_score: Optional[int]) -> Optional[Tuple[str, Tuple[object, ...]]]:
    fields = line.rstrip("\n").split()
    if len(fields) != 16:
        raise ValueError("expected 16 fields")
    combined_score = int(fields[15])
    if min_combined_score is not None and combined_score < min_combined_score:
        return None
    protein1 = fields[0]
    protein2 = fields[1]
    taxon_id = taxon_from_protein_id(protein1)
    scores = tuple(int(value) for value in fields[2:15])
    key = "%s\t%s" % (protein1, protein2)
    return key, (protein1, protein2, taxon_id) + scores + (combined_score,)


def process_text_chunk(
    table: str,
    chunk_id: int,
    lines: Sequence[str],
    seed_value: int,
    cutoffs: Tuple[int, int],
    min_combined_score: Optional[int],
) -> Dict[str, object]:
    buckets = empty_buckets(table)
    bad_rows = 0
    filtered_rows = 0
    parsed_rows = 0

    for line in lines:
        if not line.strip():
            continue
        try:
            if table == "species":
                parsed = parse_species_line(line)
            elif table == "protein_info":
                parsed = parse_protein_info_line(line)
            elif table == "protein_aliases":
                parsed = parse_alias_line(line)
            elif table == "protein_links":
                parsed = parse_link_line(line, min_combined_score)
            else:
                raise ValueError("unsupported text table %s" % table)
        except Exception:
            bad_rows += 1
            continue

        if parsed is None:
            filtered_rows += 1
            continue

        key, values = parsed
        split = assign_split(key, seed_value, cutoffs)
        append_row(buckets[split], values)
        parsed_rows += 1

    return {
        "table": table,
        "chunk_id": chunk_id,
        "buckets": buckets,
        "parsed_rows": parsed_rows,
        "bad_rows": bad_rows,
        "filtered_rows": filtered_rows,
    }


def process_sequence_chunk(
    chunk_id: int,
    records: Sequence[Tuple[str, str]],
    seed_value: int,
    cutoffs: Tuple[int, int],
) -> Dict[str, object]:
    table = "protein_sequences"
    buckets = empty_buckets(table)
    bad_rows = 0
    parsed_rows = 0

    for protein_id, sequence in records:
        try:
            taxon_id = taxon_from_protein_id(protein_id)
        except Exception:
            bad_rows += 1
            continue
        values = (protein_id, taxon_id, sequence, len(sequence))
        split = assign_split(protein_id, seed_value, cutoffs)
        append_row(buckets[split], values)
        parsed_rows += 1

    return {
        "table": table,
        "chunk_id": chunk_id,
        "buckets": buckets,
        "parsed_rows": parsed_rows,
        "bad_rows": bad_rows,
        "filtered_rows": 0,
    }


def import_arrow():
    try:
        import pyarrow as pa
        import pyarrow.parquet as pq
    except ImportError as exc:
        raise SystemExit(
            "pyarrow is required. Install dependencies with: python -m pip install -r requirements.txt"
        ) from exc
    return pa, pq


def arrow_type(pa, type_name: str):
    if type_name == "string":
        return pa.string()
    if type_name == "int16":
        return pa.int16()
    if type_name == "int32":
        return pa.int32()
    raise ValueError("unsupported type %s" % type_name)


def arrow_schema(pa, table: str):
    return pa.schema(
        [pa.field(name, arrow_type(pa, type_name)) for name, type_name in zip(COLUMNS[table], TYPE_NAMES[table])]
    )


class ParquetSink:
    def __init__(self, output_dir: Path, compression: str):
        self.output_dir = output_dir
        self.compression = None if compression == "none" else compression
        self.pa, self.pq = import_arrow()
        self.schemas = {table: arrow_schema(self.pa, table) for table in TABLE_ORDER}
        self.shard_counts: Dict[str, Dict[str, int]] = {
            table: {"train": 0, "validation": 0, "test": 0} for table in TABLE_ORDER
        }
        self.row_counts: Dict[str, Dict[str, int]] = {
            table: {"train": 0, "validation": 0, "test": 0} for table in TABLE_ORDER
        }

    def write_buckets(self, table: str, buckets: Mapping[str, Sequence[Sequence[object]]]) -> None:
        table_dir = self.output_dir / table
        table_dir.mkdir(parents=True, exist_ok=True)

        for split in ("train", "validation", "test"):
            column_lists = buckets[split]
            row_count = len(column_lists[0]) if column_lists else 0
            if row_count == 0:
                continue

            data = {name: column_lists[index] for index, name in enumerate(COLUMNS[table])}
            arrow_table = self.pa.Table.from_pydict(data, schema=self.schemas[table])
            shard_id = self.shard_counts[table][split]
            output_path = table_dir / ("%s-%05d.parquet" % (split, shard_id))
            self.pq.write_table(
                arrow_table,
                output_path,
                compression=self.compression,
                use_dictionary=True,
                write_statistics=True,
            )
            self.shard_counts[table][split] += 1
            self.row_counts[table][split] += row_count


def prepare_output_dirs(output_dir: Path, tables: Sequence[str], overwrite: bool) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    for table in tables:
        table_dir = output_dir / table
        if not table_dir.exists():
            continue
        if not overwrite:
            raise SystemExit(
                "Output directory already exists: %s. Re-run with --overwrite or choose another --output-dir."
                % table_dir
            )
        shutil.rmtree(table_dir)


def handle_result(result: Mapping[str, object], sink: ParquetSink, stats: Dict[str, Dict[str, int]]) -> None:
    table = str(result["table"])
    sink.write_buckets(table, result["buckets"])  # type: ignore[arg-type]
    stats[table]["parsed_rows"] += int(result["parsed_rows"])
    stats[table]["bad_rows"] += int(result["bad_rows"])
    stats[table]["filtered_rows"] += int(result["filtered_rows"])
    stats[table]["chunks"] += 1


def drain_one(
    pending: Set[object],
    sink: ParquetSink,
    stats: Dict[str, Dict[str, int]],
) -> Set[object]:
    done, still_pending = wait(pending, return_when=FIRST_COMPLETED)
    for future in done:
        handle_result(future.result(), sink, stats)
    return set(still_pending)


def drain_all(
    pending: Set[object],
    sink: ParquetSink,
    stats: Dict[str, Dict[str, int]],
) -> None:
    while pending:
        pending = drain_one(pending, sink, stats)


def chunk_iterator(
    table: str,
    path: Path,
    rows_per_chunk: int,
    max_rows: Optional[int],
    decompressor: str,
    decompressor_proc: int,
) -> Iterator[Tuple[int, object]]:
    if max_rows is not None and decompressor == "auto":
        decompressor = "python"
    if table == "protein_sequences":
        yield from iter_sequence_chunks(path, rows_per_chunk, max_rows, decompressor, decompressor_proc)
    else:
        yield from iter_text_chunks(path, rows_per_chunk, max_rows, decompressor, decompressor_proc)


def process_table(
    executor: ProcessPoolExecutor,
    sink: ParquetSink,
    table: str,
    raw_path: Path,
    args: argparse.Namespace,
    seed_value: int,
    cutoffs: Tuple[int, int],
    stats: Dict[str, Dict[str, int]],
) -> None:
    logging.info("Processing %s from %s", table, raw_path)
    started = time.time()
    pending: Set[object] = set()
    max_in_flight = args.max_in_flight or args.num_proc

    for chunk_id, chunk in chunk_iterator(
        table,
        raw_path,
        args.rows_per_chunk,
        args.max_rows_per_table,
        args.decompressor,
        args.decompressor_proc,
    ):
        if table == "protein_sequences":
            future = executor.submit(process_sequence_chunk, chunk_id, chunk, seed_value, cutoffs)
        else:
            future = executor.submit(
                process_text_chunk,
                table,
                chunk_id,
                chunk,
                seed_value,
                cutoffs,
                args.link_min_combined_score,
            )
        pending.add(future)

        if len(pending) >= max_in_flight:
            pending = drain_one(pending, sink, stats)

        submitted = chunk_id + 1
        if args.log_every and submitted % args.log_every == 0:
            logging.info("%s: submitted %d chunks", table, submitted)

    drain_all(pending, sink, stats)
    elapsed = time.time() - started
    logging.info(
        "Finished %s in %.1fs: %s rows, %s bad rows, %s filtered rows",
        table,
        elapsed,
        stats[table]["parsed_rows"],
        stats[table]["bad_rows"],
        stats[table]["filtered_rows"],
    )


def write_summary(
    output_dir: Path,
    tables: Sequence[str],
    args: argparse.Namespace,
    sink: ParquetSink,
    stats: Mapping[str, Mapping[str, int]],
) -> None:
    summary = {
        "raw_dir": str(args.raw_dir),
        "output_dir": str(args.output_dir),
        "tables": list(tables),
        "num_proc": args.num_proc,
        "rows_per_chunk": args.rows_per_chunk,
        "split_ratios": {
            "train": args.train_ratio,
            "validation": args.validation_ratio,
            "test": args.test_ratio,
        },
        "split_seed": args.split_seed,
        "link_min_combined_score": args.link_min_combined_score,
        "max_rows_per_table": args.max_rows_per_table,
        "compression": args.compression,
        "stats": {
            table: {
                "parser": dict(stats[table]),
                "rows": dict(sink.row_counts[table]),
                "shards": dict(sink.shard_counts[table]),
            }
            for table in tables
        },
    }
    path = output_dir / "processing_summary.json"
    path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    logging.info("Wrote %s", path)


def main() -> int:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        datefmt="%H:%M:%S",
    )
    args = parse_args()
    tables = resolve_tables(args.tables)

    if args.num_proc < 1:
        raise SystemExit("--num-proc must be >= 1")
    if args.rows_per_chunk < 1:
        raise SystemExit("--rows-per-chunk must be >= 1")
    if args.max_in_flight is not None and args.max_in_flight < 1:
        raise SystemExit("--max-in-flight must be >= 1")

    cutoffs = split_cutoffs(args.train_ratio, args.validation_ratio, args.test_ratio)
    seed_value = seed_crc(args.split_seed)

    for table in tables:
        raw_path = args.raw_dir / RAW_FILES[table]
        if not raw_path.exists():
            raise SystemExit("Missing raw file for %s: %s" % (table, raw_path))

    prepare_output_dirs(args.output_dir, tables, args.overwrite)
    sink = ParquetSink(args.output_dir, args.compression)

    try:
        sink.pa.set_cpu_count(args.num_proc)
    except AttributeError:
        pass

    stats: Dict[str, Dict[str, int]] = {
        table: {"chunks": 0, "parsed_rows": 0, "bad_rows": 0, "filtered_rows": 0} for table in TABLE_ORDER
    }

    logging.info("Using %d worker processes", args.num_proc)
    with ProcessPoolExecutor(max_workers=args.num_proc) as executor:
        for table in tables:
            raw_path = args.raw_dir / RAW_FILES[table]
            process_table(executor, sink, table, raw_path, args, seed_value, cutoffs, stats)

    write_summary(args.output_dir, tables, args, sink, stats)
    return 0


if __name__ == "__main__":
    sys.exit(main())

Xet Storage Details

Size:: 23.4 kB
Xet hash:: 43d5248094f8ac0bd3345e966bcb92be351b4f087e9f4f13f09c1b4b3b318c0c

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	#!/usr/bin/env python3
	"""Convert raw STRING downloads into Hugging Face Data Viewer-friendly Parquet.

	The converter streams the raw files, assigns rows to deterministic
	train/validation/test splits, and writes sharded Parquet files under:

	data/<config>/<split>-00000.parquet

	It is designed for very large STRING files: chunks are parsed in a bounded
	process pool and written incrementally, so the complete dataset is never loaded
	into memory.
	"""

	import argparse
	import gzip
	import io
	import json
	import logging
	import shutil
	import subprocess
	import sys
	import time
	import zlib
	from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
	from contextlib import contextmanager
	from pathlib import Path
	from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Set, Tuple


	TABLE_ORDER = (
	"species",
	"protein_info",
	"protein_aliases",
	"protein_sequences",
	"protein_links",
	)

	RAW_FILES = {
	"species": "species.v12.0.txt",
	"protein_info": "protein.info.v12.0.txt.gz",
	"protein_aliases": "protein.aliases.v12.0.txt.gz",
	"protein_sequences": "protein.sequences.v12.0.fa.gz",
	"protein_links": "protein.links.full.v12.0.txt.gz",
	}

	COLUMNS = {
	"species": (
	"taxon_id",
	"string_type",
	"string_name_compact",
	"official_name_ncbi",
	"domain",
	),
	"protein_info": (
	"string_protein_id",
	"taxon_id",
	"preferred_name",
	"protein_size",
	"annotation",
	),
	"protein_aliases": (
	"string_protein_id",
	"taxon_id",
	"alias",
	"source",
	),
	"protein_sequences": (
	"string_protein_id",
	"taxon_id",
	"sequence",
	"sequence_length",
	),
	"protein_links": (
	"protein1",
	"protein2",
	"taxon_id",
	"neighborhood",
	"neighborhood_transferred",
	"fusion",
	"cooccurence",
	"homology",
	"coexpression",
	"coexpression_transferred",
	"experiments",
	"experiments_transferred",
	"database",
	"database_transferred",
	"textmining",
	"textmining_transferred",
	"combined_score",
	),
	}

	TYPE_NAMES = {
	"species": (
	"int32",
	"string",
	"string",
	"string",
	"string",
	),
	"protein_info": (
	"string",
	"int32",
	"string",
	"int32",
	"string",
	),
	"protein_aliases": (
	"string",
	"int32",
	"string",
	"string",
	),
	"protein_sequences": (
	"string",
	"int32",
	"string",
	"int32",
	),
	"protein_links": (
	"string",
	"string",
	"int32",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	"int16",
	),
	}


	def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(
	description="Post-process STRING v12.0 raw files into sharded Parquet configs."
	)
	parser.add_argument("--raw-dir", type=Path, default=Path("v12.0"))
	parser.add_argument("--output-dir", type=Path, default=Path("data"))
	parser.add_argument(
	"--tables",
	default="all",
	help="Comma-separated table list, or 'all'. Valid tables: %s" % ", ".join(TABLE_ORDER),
	)
	parser.add_argument("--num-proc", type=int, default=32)
	parser.add_argument(
	"--max-in-flight",
	type=int,
	default=None,
	help="Maximum parser chunks queued at once. Defaults to --num-proc.",
	)
	parser.add_argument(
	"--rows-per-chunk",
	type=int,
	default=100_000,
	help="Rows sent to one worker at a time. Lower this if RAM is tight.",
	)
	parser.add_argument("--train-ratio", type=float, default=0.98)
	parser.add_argument("--validation-ratio", type=float, default=0.01)
	parser.add_argument("--test-ratio", type=float, default=0.01)
	parser.add_argument(
	"--split-seed",
	default="string-v12.0",
	help="Stable salt used by the row hash split assignment.",
	)
	parser.add_argument(
	"--link-min-combined-score",
	type=int,
	default=None,
	help="Optional filter for protein_links. Keeps rows with combined_score >= this value.",
	)
	parser.add_argument(
	"--max-rows-per-table",
	type=int,
	default=None,
	help="Optional raw-row cap per table, useful for creating a small preview dataset.",
	)
	parser.add_argument(
	"--compression",
	default="zstd",
	choices=("zstd", "snappy", "gzip", "brotli", "none"),
	help="Parquet compression codec.",
	)
	parser.add_argument(
	"--decompressor",
	default="auto",
	choices=("auto", "python", "gzip", "pigz"),
	help="How to stream .gz files. auto uses pigz when installed, otherwise Python gzip.",
	)
	parser.add_argument(
	"--decompressor-proc",
	type=int,
	default=4,
	help="Threads passed to pigz when --decompressor is pigz/auto and pigz is available.",
	)
	parser.add_argument(
	"--overwrite",
	action="store_true",
	help="Delete existing output subdirectories for selected tables before writing.",
	)
	parser.add_argument("--log-every", type=int, default=25)
	return parser.parse_args()


	def resolve_tables(value: str) -> List[str]:
	if value == "all":
	return list(TABLE_ORDER)
	selected = [item.strip() for item in value.split(",") if item.strip()]
	unknown = sorted(set(selected) - set(TABLE_ORDER))
	if unknown:
	raise ValueError("Unknown table(s): %s" % ", ".join(unknown))
	return selected


	def split_cutoffs(train_ratio: float, validation_ratio: float, test_ratio: float) -> Tuple[int, int]:
	total = train_ratio + validation_ratio + test_ratio
	if abs(total - 1.0) > 1e-9:
	raise ValueError("Split ratios must sum to 1.0, got %.8f" % total)
	train_cutoff = int(train_ratio * (2**32))
	validation_cutoff = int((train_ratio + validation_ratio) * (2**32))
	return train_cutoff, validation_cutoff


	def seed_crc(seed: str) -> int:
	return zlib.crc32(seed.encode("utf-8")) & 0xFFFFFFFF


	def assign_split(key: str, seed_value: int, cutoffs: Tuple[int, int]) -> str:
	value = zlib.crc32(key.encode("utf-8"), seed_value) & 0xFFFFFFFF
	if value < cutoffs[0]:
	return "train"
	if value < cutoffs[1]:
	return "validation"
	return "test"


	def taxon_from_protein_id(protein_id: str) -> int:
	return int(protein_id.split(".", 1)[0])


	@contextmanager
	def open_text(path: Path, decompressor: str, decompressor_proc: int) -> Iterator[io.TextIOBase]:
	if path.suffix != ".gz":
	with path.open("rt", encoding="utf-8", newline="") as handle:
	yield handle
	return

	if decompressor == "pigz" and not shutil.which("pigz"):
	raise RuntimeError("Requested --decompressor pigz, but pigz is not installed")

	use_pigz = decompressor == "pigz" or (decompressor == "auto" and shutil.which("pigz"))
	use_gzip_cmd = decompressor == "gzip"

	if use_pigz:
	cmd = ["pigz", "-dc", "-p", str(max(1, decompressor_proc)), str(path)]
	elif use_gzip_cmd:
	cmd = ["gzip", "-cd", str(path)]
	else:
	with gzip.open(path, "rt", encoding="utf-8", newline="") as handle:
	yield handle
	return

	proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
	assert proc.stdout is not None
	handle = io.TextIOWrapper(proc.stdout, encoding="utf-8", newline="")
	try:
	yield handle
	finally:
	handle.close()
	return_code = proc.wait()
	if return_code != 0:
	raise RuntimeError("Decompressor failed with exit code %s: %s" % (return_code, " ".join(cmd)))


	def iter_text_chunks(
	path: Path,
	rows_per_chunk: int,
	max_rows: Optional[int],
	decompressor: str,
	decompressor_proc: int,
	) -> Iterator[Tuple[int, List[str]]]:
	with open_text(path, decompressor, decompressor_proc) as handle:
	header = next(handle, None)
	if header is None:
	return

	chunk: List[str] = []
	chunk_id = 0
	seen = 0
	for line in handle:
	if max_rows is not None and seen >= max_rows:
	break
	chunk.append(line)
	seen += 1
	if len(chunk) >= rows_per_chunk:
	yield chunk_id, chunk
	chunk_id += 1
	chunk = []

	if chunk:
	yield chunk_id, chunk


	def iter_fasta_records(
	path: Path,
	max_rows: Optional[int],
	decompressor: str,
	decompressor_proc: int,
	) -> Iterator[Tuple[str, str]]:
	yielded = 0
	protein_id: Optional[str] = None
	parts: List[str] = []

	with open_text(path, decompressor, decompressor_proc) as handle:
	for raw_line in handle:
	line = raw_line.strip()
	if not line:
	continue
	if line.startswith(">"):
	if protein_id is not None:
	yield protein_id, "".join(parts)
	yielded += 1
	if max_rows is not None and yielded >= max_rows:
	return
	protein_id = line[1:].split(None, 1)[0]
	parts = []
	else:
	parts.append(line)

	if protein_id is not None and (max_rows is None or yielded < max_rows):
	yield protein_id, "".join(parts)


	def iter_sequence_chunks(
	path: Path,
	rows_per_chunk: int,
	max_rows: Optional[int],
	decompressor: str,
	decompressor_proc: int,
	) -> Iterator[Tuple[int, List[Tuple[str, str]]]]:
	chunk: List[Tuple[str, str]] = []
	chunk_id = 0
	for record in iter_fasta_records(path, max_rows, decompressor, decompressor_proc):
	chunk.append(record)
	if len(chunk) >= rows_per_chunk:
	yield chunk_id, chunk
	chunk_id += 1
	chunk = []
	if chunk:
	yield chunk_id, chunk


	def empty_buckets(table: str) -> Dict[str, List[List[object]]]:
	width = len(COLUMNS[table])
	return {
	"train": [[] for _ in range(width)],
	"validation": [[] for _ in range(width)],
	"test": [[] for _ in range(width)],
	}


	def append_row(bucket: List[List[object]], values: Sequence[object]) -> None:
	for index, value in enumerate(values):
	bucket[index].append(value)


	def parse_species_line(line: str) -> Tuple[str, Tuple[object, ...]]:
	fields = line.rstrip("\n").split("\t")
	if len(fields) != 5:
	raise ValueError("expected 5 fields")
	taxon_id = int(fields[0])
	return str(taxon_id), (taxon_id, fields[1], fields[2], fields[3], fields[4])


	def parse_protein_info_line(line: str) -> Tuple[str, Tuple[object, ...]]:
	fields = line.rstrip("\n").split("\t", 3)
	if len(fields) != 4:
	raise ValueError("expected 4 fields")
	protein_id = fields[0]
	taxon_id = taxon_from_protein_id(protein_id)
	return protein_id, (protein_id, taxon_id, fields[1], int(fields[2]), fields[3])


	def parse_alias_line(line: str) -> Tuple[str, Tuple[object, ...]]:
	fields = line.rstrip("\n").split("\t", 2)
	if len(fields) != 3:
	raise ValueError("expected 3 fields")
	protein_id = fields[0]
	taxon_id = taxon_from_protein_id(protein_id)
	key = "%s\t%s\t%s" % (protein_id, fields[1], fields[2])
	return key, (protein_id, taxon_id, fields[1], fields[2])


	def parse_link_line(line: str, min_combined_score: Optional[int]) -> Optional[Tuple[str, Tuple[object, ...]]]:
	fields = line.rstrip("\n").split()
	if len(fields) != 16:
	raise ValueError("expected 16 fields")
	combined_score = int(fields[15])
	if min_combined_score is not None and combined_score < min_combined_score:
	return None
	protein1 = fields[0]
	protein2 = fields[1]
	taxon_id = taxon_from_protein_id(protein1)
	scores = tuple(int(value) for value in fields[2:15])
	key = "%s\t%s" % (protein1, protein2)
	return key, (protein1, protein2, taxon_id) + scores + (combined_score,)


	def process_text_chunk(
	table: str,
	chunk_id: int,
	lines: Sequence[str],
	seed_value: int,
	cutoffs: Tuple[int, int],
	min_combined_score: Optional[int],
	) -> Dict[str, object]:
	buckets = empty_buckets(table)
	bad_rows = 0
	filtered_rows = 0
	parsed_rows = 0

	for line in lines:
	if not line.strip():
	continue
	try:
	if table == "species":
	parsed = parse_species_line(line)
	elif table == "protein_info":
	parsed = parse_protein_info_line(line)
	elif table == "protein_aliases":
	parsed = parse_alias_line(line)
	elif table == "protein_links":
	parsed = parse_link_line(line, min_combined_score)
	else:
	raise ValueError("unsupported text table %s" % table)
	except Exception:
	bad_rows += 1
	continue

	if parsed is None:
	filtered_rows += 1
	continue

	key, values = parsed
	split = assign_split(key, seed_value, cutoffs)
	append_row(buckets[split], values)
	parsed_rows += 1

	return {
	"table": table,
	"chunk_id": chunk_id,
	"buckets": buckets,
	"parsed_rows": parsed_rows,
	"bad_rows": bad_rows,
	"filtered_rows": filtered_rows,
	}


	def process_sequence_chunk(
	chunk_id: int,
	records: Sequence[Tuple[str, str]],
	seed_value: int,
	cutoffs: Tuple[int, int],
	) -> Dict[str, object]:
	table = "protein_sequences"
	buckets = empty_buckets(table)
	bad_rows = 0
	parsed_rows = 0

	for protein_id, sequence in records:
	try:
	taxon_id = taxon_from_protein_id(protein_id)
	except Exception:
	bad_rows += 1
	continue
	values = (protein_id, taxon_id, sequence, len(sequence))
	split = assign_split(protein_id, seed_value, cutoffs)
	append_row(buckets[split], values)
	parsed_rows += 1

	return {
	"table": table,
	"chunk_id": chunk_id,
	"buckets": buckets,
	"parsed_rows": parsed_rows,
	"bad_rows": bad_rows,
	"filtered_rows": 0,
	}


	def import_arrow():
	try:
	import pyarrow as pa
	import pyarrow.parquet as pq
	except ImportError as exc:
	raise SystemExit(
	"pyarrow is required. Install dependencies with: python -m pip install -r requirements.txt"
	) from exc
	return pa, pq


	def arrow_type(pa, type_name: str):
	if type_name == "string":
	return pa.string()
	if type_name == "int16":
	return pa.int16()
	if type_name == "int32":
	return pa.int32()
	raise ValueError("unsupported type %s" % type_name)


	def arrow_schema(pa, table: str):
	return pa.schema(
	[pa.field(name, arrow_type(pa, type_name)) for name, type_name in zip(COLUMNS[table], TYPE_NAMES[table])]
	)


	class ParquetSink:
	def __init__(self, output_dir: Path, compression: str):
	self.output_dir = output_dir
	self.compression = None if compression == "none" else compression
	self.pa, self.pq = import_arrow()
	self.schemas = {table: arrow_schema(self.pa, table) for table in TABLE_ORDER}
	self.shard_counts: Dict[str, Dict[str, int]] = {
	table: {"train": 0, "validation": 0, "test": 0} for table in TABLE_ORDER
	}
	self.row_counts: Dict[str, Dict[str, int]] = {
	table: {"train": 0, "validation": 0, "test": 0} for table in TABLE_ORDER
	}

	def write_buckets(self, table: str, buckets: Mapping[str, Sequence[Sequence[object]]]) -> None:
	table_dir = self.output_dir / table
	table_dir.mkdir(parents=True, exist_ok=True)

	for split in ("train", "validation", "test"):
	column_lists = buckets[split]
	row_count = len(column_lists[0]) if column_lists else 0
	if row_count == 0:
	continue

	data = {name: column_lists[index] for index, name in enumerate(COLUMNS[table])}
	arrow_table = self.pa.Table.from_pydict(data, schema=self.schemas[table])
	shard_id = self.shard_counts[table][split]
	output_path = table_dir / ("%s-%05d.parquet" % (split, shard_id))
	self.pq.write_table(
	arrow_table,
	output_path,
	compression=self.compression,
	use_dictionary=True,
	write_statistics=True,
	)
	self.shard_counts[table][split] += 1
	self.row_counts[table][split] += row_count


	def prepare_output_dirs(output_dir: Path, tables: Sequence[str], overwrite: bool) -> None:
	output_dir.mkdir(parents=True, exist_ok=True)
	for table in tables:
	table_dir = output_dir / table
	if not table_dir.exists():
	continue
	if not overwrite:
	raise SystemExit(
	"Output directory already exists: %s. Re-run with --overwrite or choose another --output-dir."
	% table_dir
	)
	shutil.rmtree(table_dir)


	def handle_result(result: Mapping[str, object], sink: ParquetSink, stats: Dict[str, Dict[str, int]]) -> None:
	table = str(result["table"])
	sink.write_buckets(table, result["buckets"]) # type: ignore[arg-type]
	stats[table]["parsed_rows"] += int(result["parsed_rows"])
	stats[table]["bad_rows"] += int(result["bad_rows"])
	stats[table]["filtered_rows"] += int(result["filtered_rows"])
	stats[table]["chunks"] += 1


	def drain_one(
	pending: Set[object],
	sink: ParquetSink,
	stats: Dict[str, Dict[str, int]],
	) -> Set[object]:
	done, still_pending = wait(pending, return_when=FIRST_COMPLETED)
	for future in done:
	handle_result(future.result(), sink, stats)
	return set(still_pending)


	def drain_all(
	pending: Set[object],
	sink: ParquetSink,
	stats: Dict[str, Dict[str, int]],
	) -> None:
	while pending:
	pending = drain_one(pending, sink, stats)


	def chunk_iterator(
	table: str,
	path: Path,
	rows_per_chunk: int,
	max_rows: Optional[int],
	decompressor: str,
	decompressor_proc: int,
	) -> Iterator[Tuple[int, object]]:
	if max_rows is not None and decompressor == "auto":
	decompressor = "python"
	if table == "protein_sequences":
	yield from iter_sequence_chunks(path, rows_per_chunk, max_rows, decompressor, decompressor_proc)
	else:
	yield from iter_text_chunks(path, rows_per_chunk, max_rows, decompressor, decompressor_proc)


	def process_table(
	executor: ProcessPoolExecutor,
	sink: ParquetSink,
	table: str,
	raw_path: Path,
	args: argparse.Namespace,
	seed_value: int,
	cutoffs: Tuple[int, int],
	stats: Dict[str, Dict[str, int]],
	) -> None:
	logging.info("Processing %s from %s", table, raw_path)
	started = time.time()
	pending: Set[object] = set()
	max_in_flight = args.max_in_flight or args.num_proc

	for chunk_id, chunk in chunk_iterator(
	table,
	raw_path,
	args.rows_per_chunk,
	args.max_rows_per_table,
	args.decompressor,
	args.decompressor_proc,
	):
	if table == "protein_sequences":
	future = executor.submit(process_sequence_chunk, chunk_id, chunk, seed_value, cutoffs)
	else:
	future = executor.submit(
	process_text_chunk,
	table,
	chunk_id,
	chunk,
	seed_value,
	cutoffs,
	args.link_min_combined_score,
	)
	pending.add(future)

	if len(pending) >= max_in_flight:
	pending = drain_one(pending, sink, stats)

	submitted = chunk_id + 1
	if args.log_every and submitted % args.log_every == 0:
	logging.info("%s: submitted %d chunks", table, submitted)

	drain_all(pending, sink, stats)
	elapsed = time.time() - started
	logging.info(
	"Finished %s in %.1fs: %s rows, %s bad rows, %s filtered rows",
	table,
	elapsed,
	stats[table]["parsed_rows"],
	stats[table]["bad_rows"],
	stats[table]["filtered_rows"],
	)


	def write_summary(
	output_dir: Path,
	tables: Sequence[str],
	args: argparse.Namespace,
	sink: ParquetSink,
	stats: Mapping[str, Mapping[str, int]],
	) -> None:
	summary = {
	"raw_dir": str(args.raw_dir),
	"output_dir": str(args.output_dir),
	"tables": list(tables),
	"num_proc": args.num_proc,
	"rows_per_chunk": args.rows_per_chunk,
	"split_ratios": {
	"train": args.train_ratio,
	"validation": args.validation_ratio,
	"test": args.test_ratio,
	},
	"split_seed": args.split_seed,
	"link_min_combined_score": args.link_min_combined_score,
	"max_rows_per_table": args.max_rows_per_table,
	"compression": args.compression,
	"stats": {
	table: {
	"parser": dict(stats[table]),
	"rows": dict(sink.row_counts[table]),
	"shards": dict(sink.shard_counts[table]),
	}
	for table in tables
	},
	}
	path = output_dir / "processing_summary.json"
	path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
	logging.info("Wrote %s", path)


	def main() -> int:
	logging.basicConfig(
	level=logging.INFO,
	format="%(asctime)s %(levelname)s %(message)s",
	datefmt="%H:%M:%S",
	)
	args = parse_args()
	tables = resolve_tables(args.tables)

	if args.num_proc < 1:
	raise SystemExit("--num-proc must be >= 1")
	if args.rows_per_chunk < 1:
	raise SystemExit("--rows-per-chunk must be >= 1")
	if args.max_in_flight is not None and args.max_in_flight < 1:
	raise SystemExit("--max-in-flight must be >= 1")

	cutoffs = split_cutoffs(args.train_ratio, args.validation_ratio, args.test_ratio)
	seed_value = seed_crc(args.split_seed)

	for table in tables:
	raw_path = args.raw_dir / RAW_FILES[table]
	if not raw_path.exists():
	raise SystemExit("Missing raw file for %s: %s" % (table, raw_path))

	prepare_output_dirs(args.output_dir, tables, args.overwrite)
	sink = ParquetSink(args.output_dir, args.compression)

	try:
	sink.pa.set_cpu_count(args.num_proc)
	except AttributeError:
	pass

	stats: Dict[str, Dict[str, int]] = {
	table: {"chunks": 0, "parsed_rows": 0, "bad_rows": 0, "filtered_rows": 0} for table in TABLE_ORDER
	}

	logging.info("Using %d worker processes", args.num_proc)
	with ProcessPoolExecutor(max_workers=args.num_proc) as executor:
	for table in tables:
	raw_path = args.raw_dir / RAW_FILES[table]
	process_table(executor, sink, table, raw_path, args, seed_value, cutoffs, stats)

	write_summary(args.output_dir, tables, args, sink, stats)
	return 0


	if __name__ == "__main__":
	sys.exit(main())