scripts/probe_fp16.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["pyarrow", "duckdb", "numpy", "datasets"]
# ///
"""Which fp16 list layout does DuckDB actually query?

The card + Space snippets are written against `list_cosine_similarity(col,
?::FLOAT[])`. fixed_size_list may surface as DuckDB ARRAY instead of LIST, in
which case that binding breaks — decide before publishing, not after.
"""
import duckdb
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

rng = np.random.default_rng(0)
vecs = rng.normal(size=(8, 1024)).astype(np.float32)
vecs /= np.linalg.norm(vecs, axis=1, keepdims=True)

for name, typ in [("varlist", pa.list_(pa.float16())),
                  ("fixedlist", pa.list_(pa.float16(), 1024))]:
    arr = pa.array([v.astype(np.float16).tolist() for v in vecs], type=typ)
    t = pa.table({"id": pa.array([f"r{i}" for i in range(8)]), "emb": arr})
    path = f"/tmp/probe_{name}.parquet"
    pq.write_table(t, path, compression="zstd", row_group_size=4)
    con = duckdb.connect()
    print(f"\n=== {name}: arrow {typ}")
    print(" duckdb type:", con.execute(
        f"SELECT typeof(emb) FROM '{path}' LIMIT 1").fetchone())
    q = rng.normal(size=1024).astype(np.float32)
    q /= np.linalg.norm(q)
    for fn in ("list_cosine_similarity", "array_cosine_similarity"):
        try:
            r = con.execute(
                f"SELECT id, {fn}(emb, ?::FLOAT[]) s FROM '{path}' ORDER BY s DESC LIMIT 2",
                [q.tolist()]).fetchall()
            print(f"  {fn}: OK {r[0]}")
        except Exception as e:
            print(f"  {fn}: FAIL {type(e).__name__}: {str(e)[:120]}")
    # self-join sanity: identical id must score 1.0
    try:
        r = con.execute(f"""
            SELECT a.id, list_cosine_similarity(a.emb, b.emb) s
            FROM '{path}' a JOIN '{path}' b USING (id) LIMIT 3""").fetchall()
        print("  self-join:", r)
    except Exception as e:
        print("  self-join FAIL:", str(e)[:120])
    print("  parquet physical:", pq.read_schema(path).field("emb").type)

# what does datasets' Sequence(float16, length=1024) actually produce?
from datasets import Features, Sequence, Value
from datasets.features.features import get_nested_type
f = Features({"emb": Sequence(Value("float16"), length=1024)})
print("\ndatasets Sequence(float16, length=1024) ->", get_nested_type(f["emb"]))

# fp16 round-trip cosine error vs fp32
a = vecs[0]
b = a.astype(np.float16).astype(np.float32)
print("fp16 cosine error:", 1 - float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b))))