Interrupted first build leaves a partial event table which later runs use without a warning #1295

"""Ctrl-C during the first build of a dataset does not leave a partial cache.

The first time a dataset is used, PyHealth builds its event table and caches it
in ``cache_dir``. This script starts that build in a child process, presses
Ctrl-C (SIGINT to the child's process group) as soon as the first part of the
table is on disk, then runs the same build again on the same ``cache_dir`` and
counts the patients. The table is published only once it is complete, so the
second run builds it again and counts all 3000 patients. (Before, the second
run reused the partial table without a warning.)

Synthetic data only: three CSV tables of 60,000 rows. POSIX only, since it
sends Ctrl-C to a process group.

Usage:
    python examples/dataset_cache_interrupt.py
"""

import glob
import os
import signal
import subprocess
import sys
import tempfile
import time

import numpy as np
import pandas as pd

N_PATIENTS = 3000
ROWS_PER_TABLE = 60_000
TABLES = ["t1", "t2", "t3"]


def write_data(root: str) -> None:
    """Writes three CSV tables of random events and a config for them."""
    rng = np.random.default_rng(0)
    config = ['version: "1.0"', "tables:"]
    for table in TABLES:
        patients = rng.integers(0, N_PATIENTS, size=ROWS_PER_TABLE)
        seconds = rng.integers(0, 10**8, size=ROWS_PER_TABLE)
        times = pd.Timestamp("2020-01-01") + pd.to_timedelta(seconds, unit="s")
        pd.DataFrame(
            {
                "patient_id": [f"p{p:05d}" for p in patients],
                "time": times.strftime("%Y-%m-%d %H:%M:%S"),
                "code": "c",
            }
        ).to_csv(f"{root}/{table}.csv", index=False)
        config += [
            f"  {table}:",
            f'    file_path: "{table}.csv"',
            '    patient_id: "patient_id"',
            '    timestamp: "time"',
            '    timestamp_format: "%Y-%m-%d %H:%M:%S"',
            "    attributes:",
            '      - "code"',
        ]
    with open(f"{root}/config.yaml", "w") as f:
        f.write("\n".join(config) + "\n")


def build(root: str, cache: str) -> None:
    """Builds (or reuses) the event table and prints its patient count."""
    import polars as pl

    from pyhealth.datasets import BaseDataset

    dataset = BaseDataset(
        root=root,
        tables=TABLES,
        dataset_name="interrupt_demo",
        config_path=f"{root}/config.yaml",
        cache_dir=cache,
    )
    events = dataset.global_event_df
    n = events.select(pl.col("patient_id").n_unique()).collect().item()
    print(f"patients in the event table: {n} (expected {N_PATIENTS})", flush=True)


def first_part_written(cache: str) -> bool:
    """True once one complete parquet part of the event table is on disk."""
    for path in glob.glob(f"{cache}/*/global_event_df.parquet*/*.parquet"):
        try:
            with open(path, "rb") as f:
                f.seek(-4, os.SEEK_END)
                if f.read(4) == b"PAR1":  # the last bytes of a complete file
                    return True
        except OSError:
            pass  # empty or just removed
    return False


def main():
    with tempfile.TemporaryDirectory() as tmp:
        root, cache = f"{tmp}/data", f"{tmp}/cache"
        os.makedirs(root)
        write_data(root)
        command = [sys.executable, __file__, "build", root, cache]

        print("run 1: build, then Ctrl-C once the first part is written", flush=True)
        run = subprocess.Popen(
            command, start_new_session=True, stderr=subprocess.DEVNULL
        )
        while run.poll() is None and not first_part_written(cache):
            time.sleep(0.005)
        if run.poll() is None:
            os.killpg(run.pid, signal.SIGINT)
        run.wait()

        print("run 2: same cache_dir, no interruption", flush=True)
        subprocess.run(command, check=True)


if __name__ == "__main__":
    if sys.argv[1:2] == ["build"]:
        build(*sys.argv[2:4])
    else:
        main()