scripts/diag_zero.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["duckdb", "pyarrow", "numpy"]
# ///
"""If the 19 event-less rows get a zero vector instead of NULL: does cosine throw,
and where do those rows land in a DESC ranking?"""
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)); mat /= np.linalg.norm(mat, axis=1, keepdims=True)
mat[1] = 0.0                      # the "absent" row
arr = pa.FixedSizeListArray.from_arrays(
    pa.array(mat.reshape(-1).astype(np.float16), type=pa.float16()), D)
t = pa.table({"id": pa.array([f"r{i}" for i in range(4)]), "emb": arr})
p = "/tmp/dz.parquet"; pq.write_table(t, p)
con = duckdb.connect()
print("arrow round-trip nulls:", pq.read_table(p).column("emb").null_count)
q = mat[0].astype(np.float32).tolist()
r = con.execute(f"SELECT id, list_cosine_similarity(emb, ?::FLOAT[]) s FROM '{p}' ORDER BY s DESC", [q]).fetchall()
print("ranked DESC:", r)
print("is r1 NaN?", [x for x in r if x[0]=='r1'])
print("self-join:", con.execute(
    f"SELECT a.id, list_cosine_similarity(a.emb,b.emb) FROM '{p}' a JOIN '{p}' b USING(id) ORDER BY 1").fetchall())
print("NaN filterable:", con.execute(
    f"SELECT id FROM (SELECT id, list_cosine_similarity(emb, ?::FLOAT[]) s FROM '{p}') WHERE isnan(s)", [q]).fetchall())