Relify: Parquet-native vector indexes with DataFusion #24131

import relify
from relify.datafusion import col, functions

session = relify.connect("./relify-data")

session.register_parquet(
    "documents",
    relify.datasets.uri("documents"),
)
documents = session.table("documents")

documents.create_index(
    "documents_embedding",
    column="embedding",
    key=["document_id"],
    config=relify.IVF(nlist=3),
)
documents.wait_for_index("documents_embedding")

query = (
    documents.search([0.2, 0.0], column="embedding")
    .where("tenant_id = 42 AND status = 'published'")
    .nprobes(3)
    .limit(3)
    .select(["document_id", "title"])
)

hits = session.to_dataframe(query)
document_stats = session.read_parquet(
    relify.datasets.uri("document_stats")
)

result = (
    hits.join(document_stats, on="document_id")
    .aggregate(
        "category",
        [
            functions.count(col("document_id")).alias("matches"),
            functions.avg(col("_distance")).alias("avg_distance"),
            functions.max(col("popularity")).alias("max_popularity"),
        ],
    )
    .sort("category")
)

print(result.to_pydict())