scripts/test_merge_logic.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["pyarrow", "numpy", "fsspec"]
# ///
"""Exercise load_column() against a synthetic local mirror of the saturate output
shape, including the failure it exists to catch (a NULL inside a vector)."""
import importlib.util, os, shutil, sys
import numpy as np, pyarrow as pa, pyarrow.parquet as pq
import fsspec

spec = importlib.util.spec_from_file_location("m", "merge_embeddings.py")
m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m)

ROOT = "/tmp/fake_emb"
DIM = m.DIM

def build(poison=False, err_row=False):
    shutil.rmtree(ROOT, ignore_errors=True)
    d = f"{ROOT}/datasets/davanstrien/prelinger-moments-emb-caption/data"
    os.makedirs(d)
    rng = np.random.default_rng(1)
    ids_all = [f"film@{i}.0" for i in range(70)]
    recs = []
    for b in range(0, 70, 32):
        chunk = ids_all[b:b+32]
        vecs = rng.normal(size=(len(chunk), DIM)).tolist()
        if poison and b == 0:
            vecs[0][5] = None
        recs.append({"id": f"b-{chunk[0]}", "ids": chunk,
                     "embeddings": vecs, "n_texts": len(chunk),
                     "prompt_tokens": 100, "error": None})
    if err_row:
        recs.append({"id": "b-boom", "ids": None, "embeddings": None,
                     "n_texts": None, "prompt_tokens": None, "error": "502"})
    pq.write_table(pa.Table.from_pylist(recs), f"{d}/part-0.parquet")

fs = fsspec.filesystem("file")
class Local:  # stand in for HfFileSystem: same glob/open, rooted at ROOT
    def glob(self, p): return sorted(fs.glob(f"{ROOT}/{p}"))
    def open(self, p, mode): return fs.open(p, mode)

build()
ids, vecs, rep = m.load_column(Local(), "caption")
print("clean:", rep, vecs.shape, vecs.dtype)
assert len(ids) == 70 and vecs.shape == (70, DIM)

build(err_row=True)
ids, vecs, rep = m.load_column(Local(), "caption")
print("with error row:", rep, vecs.shape)
assert rep["error_rows"] == 1 and len(ids) == 70

build(poison=True)
try:
    m.load_column(Local(), "caption")
    print("POISON NOT CAUGHT — bug")
    sys.exit(1)
except SystemExit as e:
    print("poison caught:", str(e)[:90])