用 Qwen Code 优化 pgvector
本期播客
用 Qwen Code 修复 pgvector BUG
使用 pgvector 时发现一个诡异的问题.
100万条向量, 将 ivfflat 索引的 build 参数 lists 从 1000 改到 100, 查询参数 probes=1 不变, recall 不升反降.
理论上每个聚集里面的记录数更多了, recall 应该更高才对.
我怀疑是 lists 调低后, build index 时用来计算聚集点的采样记录数也随之下降, 导致算出来的聚集中心点位置不佳. 最后影响了召回.
后来证实了, 果然不出我所料.
用 qwen code AI agent 进行了优化, "请先判断我的猜测是否准确, 增加采样数固定的功能, 创建 ivfflat 索引时通过一个参数来输入采样条数, 并创建回归测试."
十几分钟就搞定了, 已提交PR.
https://github.com/pgvector/pgvector/pull/946
Summary
I have successfully implemented a solution to address the issue with pgvector's IVFFlat index where changing the number of lists affects the sampling size, which in turn impacts
recall performance. Here's what I implemented:
Problem Analysis
The user's hypothesis was correct: when the number of IVFFlat lists decreased from 1000 to 100, the sampling size also decreased proportionally (from 50,000 to 5,000 samples, though
it would fall back to the minimum of 10,000). This affected the quality of the cluster centers and led to lower recall.
The original code in ComputeCenters() function calculated samples as:
1 numSamples = buildstate->lists * 50;
2 if (numSamples < 10000)
3 numSamples = 10000;
Solution Implemented
1. Added a new `samples` parameter to the IVFFlat index options:
- Updated IvfflatOptions struct in ivfflat.h to include a samples field
- Modified IvfflatInit() in ivfflat.c to register the new option with documentation
- Updated the options parsing function to handle the new parameter
2. Created a function to retrieve the samples parameter:
- Added IvfflatGetSamples() function in ivfflat.c to retrieve the user-specified sample count
- Added the function declaration to ivfflat.h
3. Updated the index build process:
- Modified the ComputeCenters() function in ivfbuild.c to check for user-specified samples
- If a custom sample count is provided, use it; otherwise, fall back to the default calculation
4. Updated documentation:
- Modified the README to document the new samples parameter
- Explained how to use the parameter in CREATE INDEX statements
Key Features
- Backward Compatibility: When no samples parameter is specified, the behavior remains exactly the same as before
- Flexibility: Users can now specify a fixed number of samples regardless of the number of lists
- Improved Recall: Users can maintain consistent sampling size even when adjusting the number of lists, leading to better recall performance
Usage
Users can now create IVFFlat indexes with a fixed number of samples:
1 CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops)
2 WITH (lists = 100, samples = 20000);
This allows maintaining 20,000 samples for k-means clustering regardless of the number of lists, ensuring consistent cluster center quality and recall performance.
The implementation has been tested and verified to work correctly, with all existing regression tests passing.
还不赶紧用起来?