LanceDB_向量索引
更新时间:2026-07-29
简介
把 (sample_id, embedding, meta) 批量写入 LanceDB 表,并可选创建 IVF_PQ 向量索引;返回每行的顺序累计写入行号,作为向量检索链路的落库步骤。
功能描述
• 批量写入:按批写入 (sample_id, embedding, meta_json),返回顺序累计的 int64 行号
• 可选建索引:写入完成且累计行数 ≥ min_rows_for_index 时自动创建 IVF_PQ 索引
• 维度校验:以首行向量维度为准,维度不一致的行跳过
• 空值稳健:embedding 或 sample_id 为空的行返回 -1
算子参数
输入
| 输入 | 含义 |
|---|---|
| sample_ids | 每条样本的唯一键 str。 |
| embeddings | L2 归一化向量,元素类型为 list |
| meta_jsons | 结构化 meta 的 JSON 字符串,可为空。 |
输出
| 输出 | 含义 |
|---|---|
| result | int64 写入行号(顺序累计);无效行返回 -1。 |
参数
| 参数名称 | 类型 | 默认值 | 描述 |
|---|---|---|---|
| uri | str | (必填) | LanceDB 目录 URI,本地路径或 s3://、bos:// 均可 |
| table_name | str | 'p6_samples' | 表名 默认值:"p6_samples" |
| vector_column | str | 'vector' | 向量列名 默认值:"vector" |
| create_ivf_pq | bool | True | 写入完成后是否创建 IVF_PQ 索引 默认值:True |
| num_partitions | int | 64000 | IVF 分区数 默认值:64000 |
| num_sub_vectors | int | 64 | PQ 子向量数 默认值:64 |
| metric | str | 'cosine' | 距离度量 可选值:["cosine", "l2", "dot"] 默认值:"cosine" |
| min_rows_for_index | int | 20000 | 建索引最小行数,不足则仅写入不建索引 默认值:20000 |
| overwrite | bool | False | 是否重建表 默认值:False |
| fail_on_error | bool | False | 出错是否抛异常 默认值:False |
调用示例
Python
1from __future__ import annotations
2
3import json
4import os
5
6import daft
7from daft import col
8
9from daft.aihc.common.udf import aihc_udf
10from daft.aihc.functions.index.lance_db_vector_index import LanceDBVectorIndex
11
12if __name__ == "__main__":
13 if os.getenv("DAFT_RUNNER", "native") == "ray":
14 import ray
15 ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
16 daft.set_runner_ray()
17 daft.set_execution_config(min_cpu_per_task=0)
18
19 # uri 为 mock 路径,请替换为实际 LanceDB 目录(本地路径 / s3:// / bos://)
20 LANCEDB_URI = "bos://your-bucket/index/p6_lancedb" # MOCK
21
22 n, dim = 8, 16
23 samples = {
24 "sample_id": [f"id_{i}" for i in range(n)],
25 "embedding": [[round(((i + j) % 5) * 0.1, 3) for j in range(dim)] for i in range(n)],
26 "meta_json": [json.dumps({"i": i}) for i in range(n)],
27 }
28 ds = daft.from_pydict(samples)
29 ds = ds.with_column(
30 "row",
31 aihc_udf(
32 LanceDBVectorIndex,
33 construct_args={
34 "uri": LANCEDB_URI,
35 "table_name": "p6_samples",
36 "metric": "cosine",
37 "min_rows_for_index": 20000, # 示例数据量小,不足则仅写入不建索引
38 "overwrite": True,
39 },
40 num_cpus=1,
41 concurrency=1,
42 batch_size=None,
43 )(col("sample_id"), col("embedding"), col("meta_json")),
44 )
45 ds.show()
评价此篇文章
