Partition Key：从一个社区提问走出来的新功能

说起今年科技圈里最热闹的话题，一定非大模型莫属：似乎每天都会涌现无数基于大模型开发的应用、各类知识库、记忆助手、智能 Agent ……当然，也包括向量数据库这一原本小众的领域。

向量数据库不仅承担着“大模型记忆体”的职能，也是 AIGC 应用开发新范式的重要组成部分。Milvus 作为向量数据库赛道的领先者，自 2019 年正式开源以来，已经成长为全球最大、最活跃的向量数据库开源项目与开发者社区。

随着 Milvus 社区的不断壮大，用户的需求也越来越多样化。今年 4 月的某天，有位基于大模型做知识库开发的社区用户提出了这样一个问题：

跟这位朋友详细沟通后，了解到他们的具体场景如下：

• 向量维度: 1536

• 租户：10K - 20K 个

• 每个租户数据量 1G

• 数据总容量: 10 - 20T

总结下来上述场景有 2 个核心特点：一是租户数量多，二是单个租户数据少。

01.三种解决方案

这个问题提出的时候，Milvus 的最新版本是 2.2.8，我们做个角色互换，在当时站在这个用户的角度，留在我们面前的选择有这么几个：

• 为每个租户创建一个 collection

• 为每个租户创建一个 partition

• 创建一个租户名称的标量字段

接下来，我们依次分析下这三种方案的可行性：

方案 1：为每个租户创建一个 collection。

这是我们最自然想到的方式，非常直观，使用也最简单，但是它有一个致命缺点，Milvus 的一个集群里面最多只能创建 65536 个集合。之所以有这个限制，是因为 Milvus 里的集合是和消息系统（Pulsar/Kafka）的 topic 绑定的，Pulsar/Kafka 的 topic 有数量上限，集合数量过多之后，topic 的复用率也会很高，会导致严重的读放大问题。因为我们有 10K - 20K 个租户，所以每个租户一个集合的方式走不通了。

不过好消息是，社区里面已经在筹划引入一些更轻量的消息系统（NATS），集合数量有望在未来达到更高的水平。假如集合数量的问题能够解决，能达到像 MySQL 那样上亿的表数目上限，每个租户一个 collection 肯定是最佳方式。

方案 2：为每个租户创建一个 partition。

这个方案和第一种类似，它也会受到 partition 个数的限制，Milvus 的一个集合最多只能创建 4096 个 partition。数量限制的原因也跟前文讲的类似，每个 partition 也是和消息系统里的 topic 绑定的。除了这个缺点之外，Milvus 2.2 版本的 partition，不具备动态加载释放的能力，假如之前创建并 load 过 partition A，现在新建一个 partition B，并且要把它 load 起来，必须要把之前 load 过的 partition A 释放掉，才能 load 这个新的 partition B。

简单描述这个操作就是：release A，load AB。当你 load 了几百个分区以后，再去新建分区加载，操作会非常复杂，基本不具备可行性。不过给大家预告一下，分区动态加载释放的需求，也是社区里面呼声很高的功能，这项功能已经明确会在 Milvus 2.3 里面提供支持，社区的朋友们可以期待一下。

方案 3：创建一个租户名称的标量字段。

初看这个方案还比较合理，在某些场景里，这种方案确实是可以满足要求的。但是采用这种做法，我们的每一次搜索都会进行全表的过滤，假如 1 个 G 的知识会生成 1 万条向量（实际情况下，生成的向量肯定会更多），那么 10K 个租户，加起来就会有 1 亿条向量。对于 1 亿条数据做全量扫描，符合条件的仅占万分之一，整个搜索性能肯定不会太好，并且还会浪费大量的算力。如果对性能要求不高，并且机器资源也比较充裕，这种方案也是可以 work 的。但是这种方式，使用起来总是不那么优雅，并且对于一些对性能要求比较高的场景也不能满足需求。

分析下来，当前的这几种方案，性能好的受租户数目限制（collection/partition），不受租户数目限制的性能不好（标量字段过滤）。

02.有没有两全其美的方案？

行文至此，不禁要问，有没有一种方案既能享受 collection/partition 方案的高性能，又能兼备标量过滤的无租户数限制？

这个问题很快被 Milvus 社区的 Maintainer 注意到并迅速商讨出了解决方案：

[Feature]: Implement logical partition keys · Issue #23553 · milvus-io/milvus

由此，引出了我们今天的主角：Partition Key，一种既具备 partition 的高性能又兼备标量过滤无租户数限制的方案。

说它具备 partition 的高性能，因为 partition key 是在物理 partition 的基础上再做了一层逻辑的 partition，每个逻辑的 partition 就是一个 partition key，一个物理 partition 对应多个 partition key。它的底层存储还是走的 partition 那套数据分片管理的逻辑，每次搜索的时候也是在一些特定的 partition 中进行，减少无关数据的计算从而保证搜索的高性能。同时，因为 partition key 是逻辑分区，不会受限于物理 partition 数目的限制，创建百万数目的 partition key 都没有问题。

Parition key 在使用上和标量过滤的方式非常相似，不过有一点需要注意，如果在集合中启用了 partiton key 的功能，那么 partition 的相关功能就会禁用。下面通过一个简单的示例代码来给大家演示 partition key 的使用方法。

• 创建集合时指定 book_name 作为 partiton key，并指定使用的物理 partition 数目为 100。

field1 = FieldSchema(name="text_id", dtype=DataType.INT64, is_primary=True)
field2 = FieldSchema(name="text_vector", dtype=DataType.FLOAT_VECTOR, dim=dim)
field3 = FieldSchema(name="book_name", dtype=DataType.VARCHAR, max_length=256, is_partition_key=True)
schema = CollectionSchema(fields=[field1, field2, field3])
collection = Collection(name="book", schema=schema, num_partitions=100)

• partition key 模式开启后，禁止创建 partition。

try: # throw exception, not support manually specifying the partition names if partition key mode is used
    collection.create_partition(partition_name="aaa")
except Exception as e:
    print(e)
    
 <MilvusException: (code=1, message=disable create partition if partition key mode is used)>

• 以书名作为 partition key，准备数据。

books = ['西游记'] * 5 
books += ['红楼梦'] * 5
books += ['水浒传'] * 5
books += ['三国演义'] * 5

data = [
    [i for i in range(row_count)],  # ID
    [[random.random() for _ in range(dim)] for _ in range(row_count)],  # vectors
    books,  # partiitonKey
]

• partition key 模式开启后，禁止指定 partition name。

try: # throw exception, not support manually specifying the partition names if partition key mode is used
    collection.insert(data, partition_name="_default_0")
except Exception as e:
    print(e)
    
<MilvusException: (code=1, message=not support manually specifying the partition names if partition key mode is used)>

# ok to insert
collection.insert(data)
print("succeed to insert {} entities".format(row_count))

通过 expr 表达式指定 partition key 进行搜索，支持如下两种表达式：

expr='<partition_key>=="xxxx"'

expr='<partition_key> in ["xxx", "xxx"]'

results = collection.search(data=search_vectors, anns_field="text_vector",
                            param={"metric_type": "L2", "params": {}},
                            limit=3, expr="book_name=='西游记'",output_fields=["book_name"])

到这里，文章开头提到的那个问题可以解决了，把租户名称作为 partition key，同一个租户下的数据使用同一个 partition key，10K - 20K 租户数的需求可以完美解决。

除了高性能和无租户数限制，partition key 还有另一个值得一提的地方。前文讲到 2.2 版本里的 partition 没有动态加载释放的功能，当我们的 partition 数目过多之后，partition 的管理使用是非常麻烦的，需要频繁地对集合和分区做加载释放，使用 partition key 将完全摆脱这些问题，你只用关心自己的业务，有新的租户过来直接指定一个新的 partition key 插入集合即可。

不过，现在的 partition key 也并非十全十美，当我们想要删除某个租户的数据时，由于存在 partition key 无法作为主键的限制，必须先用 query 接口根据 partition key 找到主键，然后再根据主键来做删除。不能像 collection 或者 partition 管理租户，可以直接通过删除 collection 或者 partition 的方式来删除某个租户的数据，未来还有优化空间。

打开 Milvus 官网的 Release Notes，我们可以看到在今年 6 月份发布的 Milvus 2.2.9 版本，也就是社区提出这个问题时的下一个 Milvus 版本，partition key 功能就已经上线了。

作为 Milvus 社区的 Committer，笔者借此也希望社区的每一位朋友都能慷慨地去分享你在使用 Milvus 过程中遇到的问题，以及对 Milvus 期望的功能。说不定你今天提的 feature，就在下一个版本中有了呢？