Milvus实战：构建Q&A系统及推荐系统
Milvus简介
全民AI的时代已经在趋势之中，各类应用层出不穷，而想要构建一个完善的AI应用/系统，底层存储是不可缺少的一个组件。
与传统数据库或大数据存储不同的是，这种场景下则需要选择向量数据库，是专门用来存储和查询向量的数据库，其存储的向量来自于对文本、语音、图像、视频等的向量化数据，向量数据库不仅能够完成基本的CRUD（添加、读取查询、更新、删除）等操作，还能够对向量数据进行更快速的相似性搜索。
Milvus是众多向量库中的之一，适用于多个场景，如Questions & Answering系统、推荐系统等，单节点 Milvus 可以在秒内完成十亿级的向量搜索，分布式架构亦能满足用户的水平扩展需求。
参考文档：
Milvus官网
为AI而生的数据库：Milvus详解及实战
实践之问答系统&推荐系统Mix-in
元数据定义
不论是问答还是推荐，它们对上层暴露的接口仅仅是predict(...)/search(...)/query(...)，模式是相同的，因此可以共用一个基本的Schema，固定基本的字段即可。
public class MilvusMeta {
@Getter
private final RecommenderSchema defaultMetricsSchema;
@Getter
private final QASchema defaultQASchema;
public interface Schema {
String getCollectionName();
CreateCollectionParam getCreateCollectionParam();
CreateIndexParam getCreateIndexParam();
}
@AllArgsConstructor
public abstract static class BasicSchema implements Schema {
public static final String SEARCH_PARAM = "{\"nprobe\":10}"; // Params
public static final String INDEX_PARAM = "{\"nlist\":1024}"; // ExtraParam
public static final IndexType INDEX_TYPE_DEFAULT = IndexType.IVF_FLAT;
public static final MetricType METRIC_TYPE_DEFAULT = MetricType.L2; // metric type
public static final String INDEX_NAME_DEFAULT = "ivf_flat";
public static final String PRIMARY_KEY_FIELD_NAME_DEFAULT = "id";
public static final String PARTITION_KEY_FIELD_NAME_DEFAULT = "public";
public static final String FIELD_NAME_DEFAULT = "embeddings";
@Getter
protected final CreateCollectionParam createCollectionParam;
@Getter
protected final CreateIndexParam createIndexParam;
@Override
public String getCollectionName() {
return createCollectionParam.getCollectionName();
}
}
/**
* This schema is designed for storing Zen metrics.
* TODO: Add more fields/features to describe a metric.
*/
public static class RecommenderSchema extends BasicSchema {
private RecommenderSchema(CreateCollectionParam collectionParam, CreateIndexParam indexParam) {
super(collectionParam, indexParam);
}
public static RecommenderSchema create(ZenAiConfig.Storages.MilvusConf conf) {
ZenAiConfig.Storages.Collection collection = conf.getActiveRecommenderCollection();
return new RecommenderSchema(
defaultCollectionParam(collection, collection.getEmbeddingsDimension()),
MilvusUtil.createIndexParam(collection));
}
private static CreateCollectionParam defaultCollectionParam(ZenAiConfig.Storages.Collection collection,
int dimension) {
FieldType pkType = FieldType.newBuilder()
.withName(collection.getPrimaryKey())
.withDataType(DataType.VarChar)
.withPrimaryKey(true)
.withMaxLength(100)
.withAutoID(false)
.build();
// 被embedding的字段
FieldType embeddedFieldType = FieldType.newBuilder()
.withName(collection.getEmbeddedFieldName())
.withDataType(DataType.VarChar)
.withMaxLength(255)
.build();
// embedding vector字段
FieldType embeddingFieldType = FieldType.newBuilder()
.withName(collection.getFieldName())
.withDataType(DataType.FloatVector)
.withDimension(dimension)
.build();
// 指定分区键字段，每一个Collection都需要指定一个分区键，除了能够Hive/Spark那样切分数据外，还能够加速相似查询。
// 虽然Milvus支持多种方案以切分数据，但从管理复杂度、查询效率上来看，一个Collection对应多个数据分区，是最佳的方案。
FieldType partitionKeyType = FieldType.newBuilder()
.withName(collection.getPartitionKey())
.withPartitionKey(true)
.withDataType(DataType.VarChar)
.withMaxLength(100)
.build();
return CreateCollectionParam.newBuilder()
.withCollectionName(collection.getName())
.withDescription(collection.getDescription())
// .withShardsNum(2)
.addFieldType(pkType)
.addFieldType(embeddedFieldType)
.addFieldType(embeddingFieldType)
.addFieldType(partitionKeyType)
// 开启动态字段添加功能
.withEnableDynamicField(true)
.build();
}
}
public static class QASchema extends BasicSchema {
public static final String ANSWER_FIELD_NAME = "answer";
public static final String SCORE_FIELD_NAME = "score";
public static final float SCORE_MAX_DEFAULT = 5.0f;
public static final float SCORE_MIN_DEFAULT = 0.0f;
public static final String INTENTION_FIELD_NAME = "intention";
public static final String QUESTION_OCCURRENCE = "occurrence";
public QASchema(CreateCollectionParam createCollectionParam, CreateIndexParam createIndexParam) {
super(createCollectionParam, createIndexParam);
}
public static QASchema create(ZenAiConfig.Storages.MilvusConf conf) {
ZenAiConfig.Storages.Collection collection = conf.getActiveQACollection();
return new QASchema(
defaultCollectionParam(collection, collection.getEmbeddingsDimension()),
MilvusUtil.createIndexParam(collection));
}
private static CreateCollectionParam defaultCollectionParam(ZenAiConfig.Storages.Collection collection,
int dimension) {
FieldType pkType = FieldType.newBuilder()
.withName(collection.getPrimaryKey())
.withDataType(DataType.VarChar)
.withPrimaryKey(true)
.withMaxLength(100)
.withAutoID(false)
.build();
FieldType embeddedFieldType = FieldType.newBuilder()
.withName(collection.getEmbeddedFieldName())
.withDataType(DataType.VarChar)
.withMaxLength(65535)
.build();
FieldType embeddingFieldType = FieldType.newBuilder()
.withName(collection.getFieldName())
.withDataType(DataType.FloatVector)
.withDimension(dimension)
.build();
FieldType partitionKeyType = FieldType.newBuilder()
.withName(collection.getPartitionKey())
.withPartitionKey(true)
.withDataType(DataType.VarChar)
.withMaxLength(100)
.build();
return CreateCollectionParam.newBuilder()
.withCollectionName(collection.getName())
.withDescription(collection.getDescription())
// .withShardsNum(2)
.addFieldType(pkType)
.addFieldType(embeddedFieldType)
.addFieldType(embeddingFieldType)
.addFieldType(partitionKeyType)
.withEnableDynamicField(true) // enable to insert new fields without modifying the code
.build();
}
}
Milvus可行的操作接口定义
public interface IMilvusOperations {
ZenAiConfig.Storages.Collection getCollection();
MilvusConnection.MultiStatus delete(Filter filter);
MilvusConnection.MultiStatus create(MilvusMeta.Index index);
MilvusConnection.MultiStatus drop(String index);
MilvusConnection.MultiStatus insert(MilvusData.Dataset dataset);
MilvusConnection.MultiStatus insertAndFlush(MilvusData.Dataset dataset);
/**
* Query records by filter on the specified partition, which works like a normal SQL engine.
*
* @param partition which partition to query
* @param filter boolean expression obeys the rules of Milvus
* @param outputFields if empty, the result will contain all the fields, including the dynamic;
* otherwise the result only contains the specified fields.
* @return a nonnull instance, size of which is 0 if no matched records, otherwise is positive.
*/
MilvusData.BasicPredictData queryByPartition(String partition, Filter filter, List<String> outputFields);
List<MilvusData.BasicPredictData> search(List<List<Float>> vectors, Filter filter, int topK,
List<String> outputFields);
}
抽象系统接口定义
/**
* 每个系统可能有不同的embedding的实现，因此需要定义一个接口。
*/
public interface IEmbedding {
ImmutableList<List<Float>> getEmbeddings(List<String> messages);
}
/**
* 通用接口定义，供应用层使用，可以基于sentence返回Milvus相似性结果集。
*/
public interface INlpSystem extends IMilvusOperations, IDataset, IEmbedding {
default MilvusData.BasicPredictData predict(String sentence) {
return predict(sentence, Filter.TRUE);
}
default MilvusData.BasicPredictData predict(String sentence, Filter filter) {
return predict(sentence, filter, getCollection().getOutputFields());
}
default MilvusData.BasicPredictData predict(String sentence, Filter filter,
List<String> outputFields) {
ImmutableList<List<Float>> vectors = getEmbeddings(Lists.newArrayList(sentence));
if (vectors.isEmpty()) {
return MilvusData.BasicPredictData.EMPTY;
}
List<String> mergedOutputFields = Sets.union(
ImmutableSet.copyOf(outputFields),
ImmutableSet.copyOf(getCollection().getOutputFields()))
.immutableCopy().asList();
List<MilvusData.BasicPredictData> res = search(vectors, filter, getCollection().getTopk(), mergedOutputFields);
return res.isEmpty() ? MilvusData.BasicPredictData.EMPTY : res.get(0);
}
default MilvusData.BasicPredictData predictByPartition(String partition, String sentence) {
return predictByPartition(partition, sentence, Filter.TRUE, getCollection().getOutputFields());
}
default MilvusData.BasicPredictData predictByPartition(String partition, String sentence,
Filter filter, List<String> outputFields) {
ImmutableList<List<Float>> vectors = getEmbeddings(Lists.newArrayList(sentence));
if (vectors.isEmpty()) {
return MilvusData.BasicPredictData.EMPTY;
}
List<String> mergedOutputFields = Sets.union(
ImmutableSet.copyOf(outputFields),
ImmutableSet.copyOf(getCollection().getOutputFields()))
.immutableCopy().asList();
return searchByPartition(partition, vectors.get(0), filter, mergedOutputFields);
}
}
/**
* Q & A系统接口。
*/
public interface IQuestionAnswering extends INlpSystem {
}
/**
* 推荐系统接口。
*/
public interface IRecommender extends INlpSystem, ISyncer {
}
插入数据集定义
以列式格式构建插入Milvus的数据集，需要注意的是，Milvus JAVA SDK 2.3.1版本并不支持列式导致dynamic fields，因此我对源码进行了改造，以支持列式插入动态字段。
这个问题，已经反馈给了社区，并且已经在v2.3.2版本中支持。
public interface MilvusData {
interface BasicData {
/**
* Return a list view of the splitted data, to avoid copy.
*/
BasicData[] split(int splitSize);
/**
* Return a view of the range [start, end) data, to avoid copy.
*/
BasicData subData(int groupId, int start, int end);
int size();
}
interface EmbeddingsProducer extends Function<List<String>, ImmutableList<List<Float>>> {
}
@Getter
@Setter
@AllArgsConstructor
@NoArgsConstructor
class Dataset {
private List<BasicInsertData> inserts;
}
abstract class GroupedBasicData implements BasicData {
@Getter
private final int groupId;
@Getter
@Setter
@Accessors(chain = true)
private GroupedBasicData parent;
protected GroupedBasicData(int groupId) {
this.groupId = groupId;
}
/**
* Split the data into more more sub-dataset.
*
* @param groups the number of expected groups
* @return an array of sub-dataset views from the original dataset
*/
public abstract BasicData[] grouped(int groups);
/**
* 每一个切分或是extract的子数据集，都应该拥有一个可以唯一标识它的ID
*/
public String fullGroupId() {
if (parent == null) {
return String.valueOf(groupId);
}
return parent.fullGroupId() + "-" + groupId;
}
}
@Getter
abstract class PartitionedBasicData<T> extends GroupedBasicData {
// 每一个系统都需要指定一个分区键，因此为了能够最小化存储，这里使用一个变量
// 保存整个数据集应该插入
private final T partition;
protected PartitionedBasicData(T partition, int groupId) {
super(groupId);
this.partition = partition;
}
public abstract List<T> getPartitions();
}
/**
* 以列式的形式构建插入数据集，并完成数据导入到Milvus。
* Milvus
*/
class BasicInsertData extends PartitionedBasicData<String> {
private final String collection;
private final ImmutableList<String> ids;
private final ImmutableList<String> embeddingsInput;
private final Supplier<ImmutableList<List<Float>>> vectorsSupplier;
private final EmbeddingsProducer embeddingsProducer;
private ImmutableMap<String, List<?>> dynamicFields;
private final AtomicBoolean vectorsInitialized = new AtomicBoolean(false);
}
结果集定义
public interface MilvusData {
/**
* 一个通用的数据集，可以保存search/query的结果，行式数据结构。
*/
@Getter
class BasicPredictData extends GroupedBasicData {
@Getter
@Builder
@AllArgsConstructor
public static class Row {
@JsonProperty
private String id;
private String embeddingsInput;
@JsonProperty
private Map<String, Object> extensions;
@JsonProperty
private float distance;
@JsonIgnore
private List<Float> vector;
public <T> T getAs(String key, Class<T> clazz) {
return getAs(key, clazz, null);
}
public <T> T getAs(String key, Class<T> clazz, T defaultValue) {
return clazz.cast(extensions.getOrDefault(key, defaultValue));
}
}
}
}
数据插入实例：列式插入
这个代码示例展示了如何构建列式数据集，并将其插入Milvus的流程。
注意到这里特别演示了使用了多线程并行 插入的功能，其原因有二：
- 一个批次的数据集过大，Milvus无法一次快速且稳定地完成插入动作，因此需要将原始数据集进行分组，例如这里分成3个组；
- 通常LLM（Large language Model）的一次API调用，只能支持生成16个向量数组，因此这里又对每一个分组后的子数据集进行横向切分，产生多个Batch，每个Batch包含一条记录。
@Test
void testSyncCollections() throws ExecutionException, InterruptedException {
ExecutorService executorService = Executors.newFixedThreadPool(2);
// 一组唯一值，用于区别每一条数据记录
ImmutableList<String> ids = ImmutableList.of("1", "2", "3", "4", "5");
// 一组指标名，这些指标就是待检索的合法指标集。
ImmutableList<String> metrics = ImmutableList.of("m1", "m2", "m3", "m4", "m5");
// 生成一组包含5个向量的列表，对应于每一个指标名
ImmutableList<List<Float>> vectors = generateVectors(5, TEST_COLLECTION_DIMENSION);
// 构建插入数据集
MilvusData.BasicInsertData data = new MilvusData.BasicInsertData(
config.getStorages().getMilvusConf().getActiveRecommenderCollection().getName(), ids, metrics,
//这里使用Java中的Provider接口，提供执行插入数据任务时，对指标名列表向量化，由于这里事先生成了向量数组，因此直接从索引构建数据
messages -> messages.stream().map(metrics::indexOf).map(vectors::get).collect(toImmutableList()));
ImmutableList<Double> randoms = ImmutableList.of(1.0d, 2.0d, 3.0d, 4.0d, 5.0d);
// 添加动态字段及相应的数据
data.updateDynamicFields(ImmutableMap.of("random", randoms));
// 构建并行插入数据任务
// 3 groups: ([1, 2]), ([3, 4]), ([5])
// 1 batche: ([1],[2]),([3],[4]),([5])
CompletableFuture<Integer>[] futures = milvusService.syncMetrics(data, 3, 1, executorService);
assertEquals(3, futures.length);
CompletableFuture.allOf(futures).join();
assertEquals(2, futures[0].get());
assertEquals(2, futures[1].get());
assertEquals(1, futures[2].get());
}
相似性检索实例：指标推荐
用户输入一个指标（Metric）名，或是包含指标名的语句，可以通过Milvus的Search接口，找到最相近的TOP K指标，前提是需要对输入指标名进行向量化，然后以此向量来r从Milvus库中既存的指标集中计算找到最相似的。
@Test
void testLookingForMetric() {
int topK = config..getMilvusConf().getActiveRecommenderCollection().getTopk();
Optional<MilvusData.BasicPredictData> metrics = milvusService.getActiveRecommenderSys().map(system -> system.predict("销售总额"));
assertTrue(metrics.isPresent());
assertEquals(metrics.get().getRows().size(), topK);
metrics = milvusService.getActiveRecommenderSys().map(system -> system.predict("销售总额", lt("random", 0f)));
assertFalse(metrics.isPresent());
}
相似性检索实例：问答
用户输入一段描述，可以通过Milvus提供的Search接口，找到历史相关的问题，并返回与此问题相关的上下文，并辅助回答AI模型回答用户的当前问题。
@Test
void testSearchWithSimilarityOfMultiVectors() {
ImmutableList<String> testQuestions = ImmutableList.of("用柱形图展示2019年12月的总销售额", "用拆线图展示2020年12月的总净利润");
ImmutableList<String> testIds = testQuestions.stream()
.map(DefaultQuestionAnswering::encodeQuestion)
.collect(ImmutableList.toImmutableList());
IEmbedding embeddingSvc = aiService.getMilvusService().get().getActiveQuestionAnswering().get();
DefaultQuestionAnswering.QAInsertData insertData = system.getInsertDataBuilder()
.ids(testIds)
.questions(testQuestions)
.answers(ImmutableList.of("很好", "不错"))
// 用户对于此问题返回结果的评价
.scores(ImmutableList.of(1.0f, 1.0f))
// 定义embeddings生成器，在插入时才会计算embeddings
.embeddingsProducer(
questions -> {
ImmutableList<List<Float>> qvectors = embeddingSvc.getEmbeddings(questions);
ImmutableList<List<Float>> mvectors = embeddingSvc.getEmbeddings(ImmutableList.of("总销售额", "总净利润"));
return ImmutableList.of(merge(qvectors.get(0), mvectors.get(0)), merge(qvectors.get(1), mvectors.get(1)));
})
.build();
system.insert(new MilvusData.Dataset(ImmutableList.of(insertData)));
// Case 1:
// 用柱形图展示2019年12月的总销售额: 52.0181
// 用拆线图展示2020年12月的总净利润: 147.0664
verifySearch(system, "用拆线图展示2020年12月的总销售额", "总销售额", 0, "销售额", this::merge);
// Case 2:
// 用柱形图展示2019年12月的总销售额: 313.70105
// 用拆线图展示2020年12月的总净利润: 181.3783
verifySearch(system, "今年5月的净利润详情", "净利润", 0, "利润", this::merge);
// Case 3:
// 用柱形图展示2019年12月的总销售额: 160.30568
// 用拆线图展示2020年12月的总净利润: 357.7008
verifySearch(system, "今年5月的销售额详情", "销售额", 0, "销售额", this::merge);
}
总结
Milvus对上层提供了与传统数据库相似的接口，以管理Milvus数据，同时提供了带有过滤功能的数据检索接口，使得上层应用能够很方便地利用传统数据库思维，来设计 和实现自己的系统。
但在使用中也感受到一些局限性或可能提升的点：
- 库中的一行记录只能对应一个embedding vector：只能使用相同模型生成的vector才能更好地检索向量，如果想一处持编码的文本对应多个vectors是不可能的，用户不得不创建新的Collection存储相同文本的不同向量。
- 用户显示Flush/Load Collection：每一次更新数据集，客户端必须要显示地load collection的操作，才能将新的数据加载到Server结点的内存中，同时第一次加载Collection必须是全量。
- 粗糙的表达式字符串：对于API接口的使用，缺少便利的表达式类定义，只能传递字符串，很容易出错，只能在运行时才知道哪些出错了。
- 缓存特性的支持：通常Milvus被用作Cache角色被引入系统中，但Milvus缺少一些缓存特性，如过期自动清理、partial dataset的load/unload功能等。