Polars With Cosine Similarity

example.py

# !uv pip install polars polars_distance numpy

import polars_distance as pld
import polars as pl
import numpy as np
import time

dims = 512
num_rows = 250_000

df = pl.DataFrame(
    {
        "vec1": [np.random.rand(dims) for _ in range(num_rows)],
        "vec2": [np.random.rand(dims) for _ in range(num_rows)],
    },
    schema={
        "vec1": pl.Array(pl.Float64, dims),
        "vec2": pl.Array(pl.Float64, dims),
    },
)

start = time.time()
df.with_columns(pld.col("vec1").dist_arr.cosine("vec2").alias("cosine_similarity"))
end = time.time()
print(f"Time taken: {end - start:0.3f} seconds")

display(df.head())