pgvector ivfflat index build-time parameters

Pgvector has a lists parameter that should be set as follows: For datasets with less than one million rows, use lists = rows / 1000. For datasets with more than one million rows, use lists = sqrt(rows). It is generally advisable to have at least 10 clusters.

You can use the following code to simplify creating ivfflat indexes:

def create_ivfflat_index(conn, table_name, column_name, query_operator="<=>"):
    index_method = "invalid"
    if query_operator == "<->":
        index_method = "vector_l2_ops"
    elif query_operator == "<#>":
        index_method = "vector_ip_ops"
    elif query_operator == "<=>":
        index_method = "vector_cosine_ops"
    else:
        raise ValueError(f"unrecognized operator {query_operator}")

    with conn.cursor() as cur:
        cur.execute(f"SELECT COUNT(*) as cnt FROM {table_name};")
        num_records = cur.fetchone()[0]

        num_lists = num_records / 1000
        if num_lists < 10:
            num_lists = 10
        if num_records > 1000000:
            num_lists = math.sqrt(num_records)

        cur.execute(f'CREATE INDEX ON {table_name} USING ivfflat ({column_name} {index_method}) WITH (lists = {num_lists});')
        conn.commit()