用ClickHouse玩转向量搜索

Nauu

发布于 2023-09-02 10:56:06

1.1K0

发布于 2023-09-02 10:56:06

文章被收录于专栏：ClickHouse的秘密基地ClickHouse的秘密基地

ChatGPT火了，顺带着把向量数据库也带火了。各种向量数据库如雨后春笋般的出现在了众人眼前。

那 ClickHouse 能玩向量搜索吗？

答案：那必须是可以啦， ClickHouse 可是能把数组玩出花的存在啊！！！

在 ClickHouse 中，可以使用浮点类型的数组保存向量

Array(Float32)

然后用内置的距离函数，得出两组向量之间的相似度.

目前距离函数支持主流的余弦相似度和欧几里得距离:

#余弦相似度
cosineDistance(vector1, vector2)

#欧几里得距离
L2Distance(vector1, vector2)

接下来我用一个简单示例演示:

1.准备一些测试文档数据，利用 azure openai的接口帮我们 Embedding成向量

input_data = [
    "ClickHouse是一款高性能的列式数据库，特别适合处理海量数据。",
    "它支持实时分析和查询，可以在不影响性能的情况下轻松搞定大规模数据。",
    "ClickHouse的可扩展性强，可以处理PB级别的数据，而且性能表现非常优异。",
    "数据存储和压缩方面也非常高效，可支持多种数据格式。",
    "ClickHouse提供完整的SQL支持，方便用户进行数据查询和分析。",
    "它可以与各种数据源集成，包括Hadoop、Kafka、Elasticsearch等等。",
    "ClickHouse还支持自定义指标和聚合函数，可以根据用户需求进行定制化操作。",
    "作为一款开源免费的数据库，ClickHouse的社区非常活跃，拥有强大的生态系统。",
    "它还提供了丰富的工具和插件，包括可视化工具、ETL工具等等，方便用户进行数据分析和处理。",
    "总的来说，ClickHouse是一款非常优秀的数据库产品，可以满足各种数据处理和分析的需求。"
]

embeddings_data = []
for i in range(len(input_data)):
    data = openai.Embedding.create(input=input_data[i], engine=embedding_model).data[0]["embedding"]
    embeddings_data.append([i, input_data[i], data])
    # print(data)

2.建一张MergeTree，用浮点数组保存向量数据

client = clickhouse_connect.get_client(host='ch9.nauu.com', username='default')

client.command('DROP TABLE IF EXISTS test_embedding_ch')
client.command(
    """
    CREATE TABLE test_embedding_ch
    (
        `key` UInt32,
        `content`String,
        `text_embedding` Array(Float32)
    )
    ENGINE = MergeTree
    ORDER BY key
    """)
result = client.query('DESCRIBE TABLE test_embedding_ch')
print (result.column_names[0:2])
print (result.result_columns[0:2])

client.insert('test_embedding_ch', embeddings_data, column_names=['key', 'content','text_embedding'])

3. 将用户提问同样向量化

q = ""
while q != 'quit':
    q = input("请提问: ")
    if q!="quit":
        # print("q: "+q)
        data = openai.Embedding.create(input=q, engine=embedding_model).data[0]["embedding"]

        result = query(data)
        for i in range(len(result.result_rows)):
            print("no:{} {} score:{} ".format(result.result_rows[i][0],result.result_rows[i][1],result.result_rows[i][2]))

4. 利用距离函数，轻松通过SQL实现向量的相似度查询

def query(d):
    sql = "SELECT key,content,L2Distance(text_embedding,{embeddings}) AS score FROM test_embedding_ch ORDER BY score ASC LIMIT 5".format(embeddings=d)
    print(sql)
    result = client.query(sql)
    return result

最终效果: