# /// script
# requires-python = ">=3.11"
# dependencies = ["duckdb", "pyarrow", "numpy"]
# ///
"""What does DuckDB actually see where a vector is absent? fixed_size_list vs list,
NULL-list vs zero-vector."""
import duckdb, numpy as np, pyarrow as pa, pyarrow.parquet as pq
rng = np.random.default_rng(0); D = 8
mat = rng.normal(size=(4, D)).astype(np.float16)
mask = np.array([False, True, False, False])   # row 1 absent

for label, typ in [("fixed", "fixed"), ("var", "var")]:
    flat = pa.array(mat.reshape(-1), type=pa.float16())
    if typ == "fixed":
        arr = pa.FixedSizeListArray.from_arrays(flat, D, mask=pa.array(mask))
    else:
        vals, offs = [], [0]
        for i in range(4):
            if not mask[i]:
                vals.extend(mat[i].tolist())
            offs.append(len(vals))
        arr = pa.ListArray.from_arrays(
            pa.array(offs, type=pa.int32()), pa.array(vals, type=pa.float16()),
            mask=pa.array(mask))
    t = pa.table({"id": pa.array([f"r{i}" for i in range(4)]), "emb": arr})
    p = f"/tmp/dn_{label}.parquet"; pq.write_table(t, p)
    con = duckdb.connect()
    print(f"\n=== {label}: {arr.type}, arrow nulls={arr.null_count}")
    print("  read back arrow nulls:", pq.read_table(p).column("emb").null_count)
    print("  duckdb rows:", con.execute(f"SELECT id, emb IS NULL isnull, len(emb) n FROM '{p}'").fetchall())
    print("  duckdb row1 value:", con.execute(f"SELECT emb FROM '{p}' WHERE id='r1'").fetchone())
    try:
        print("  self-join:", con.execute(
            f"SELECT a.id, list_cosine_similarity(a.emb,b.emb) FROM '{p}' a JOIN '{p}' b USING(id)").fetchall())
    except Exception as e:
        print("  self-join FAIL:", str(e)[:100])
    try:
        print("  guarded self-join:", con.execute(
            f"SELECT a.id, CASE WHEN a.emb IS NULL OR b.emb IS NULL THEN NULL ELSE list_cosine_similarity(a.emb,b.emb) END FROM '{p}' a JOIN '{p}' b USING(id)").fetchall())
    except Exception as e:
        print("  guarded FAIL:", str(e)[:100])