Skip to content

Commit

Permalink
[vector] Add support for nprobe session variable and query level hint
Browse files Browse the repository at this point in the history 
Summary:
Nearest neighbor search is most accurate with "flat" indexes. With large # of vector embeddings in the dataset, this search can get slow. To speed up the search,  approximate nearest neighbor search can be done via indexes like the IVF (family of) indexes. Search accuracy with IVF can be controlled using the "nprobe" value. A higher value of "nprobe" will (possibly) increase the accuracy and a lower value will speed up the query, but may possibly lower the accuracy. The results can be very dataset and usecase-dependent.

This diff adds support for defining a default "nprobe" value using a new session variable "rocksdb_nprobe", as well as a new query hint. The query hint can be used to essentially override the default "nprobe" session value. Valid values for both are between 1 and 10000.

  set session fb_vector_search_nprobe = 10;

Query overrides the session default:
  select /*+ SET_VAR(fb_vector_search_nprobe = 1) */ *, fb_vector_l2(vector1, '[0.3, 0.3]')  as dis from t1 force index(key1) order by dis limit 10;

Differential Revision: D54886377

fbshipit-source-id: 6acfa2e
  • Loading branch information
@saumitrp @inikep
saumitrp authored and inikep committed Jul 30, 2024
1 parent 8a42e14 commit 85fd21c
Show file tree
Hide file tree
Showing 10 changed files with 584 additions and 6 deletions.
12 changes: 12 additions & 0 deletions mysql-test/r/mysqld--help-notwin.result
View file
Original file line number Diff line number Diff line change
Expand Up @@ -593,6 +593,17 @@ The following options may be given as the first argument:
maximum vector dimension
--fb-vector-min-dimension[=#]
minimum vector dimension
--fb-vector-search-nprobe[=#]
This parameter controls the vector search radius for a
query vector when doing nearest neighbour search or
similarity search. The nprobe value corresponds to the
number of closest centroids that become part of the
vector search space. Only vectors in the nprobe-closest
voronoi cells or partitions are searched. This session
default can be superceded by a query level nprobe value
using a query hint like this: 'SELECT /*+
SET_VAR(fb_vector_search_nprobe = 3) */ ... '. Default:
16
--filesort-max-file-size=#
The max size of a file to use for filesort. Raise an
error when this is exceeded. 0 means no limit.
Expand Down Expand Up @@ -3391,6 +3402,7 @@ failure-injection-points
fast-integer-to-string FALSE
fb-vector-max-dimension 4096
fb-vector-min-dimension 3
fb-vector-search-nprobe 16
filesort-max-file-size 0
flush FALSE
flush-only-old-table-cache-entries FALSE
Expand Down
234 changes: 234 additions & 0 deletions mysql-test/suite/fb_vectordb/r/ivf_vector_index_nprobe.result
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,234 @@
insert into VECTORDB_DATA values ('id1', 'metadata', 0, JSON_OBJECT('version', 1)), ('id1', 'quantizer', 0, '[0, 0, 0, 0]'), ('id1', 'quantizer', 1, '[1, 1, 0, 0]'),
('id1', 'quantizer', 2, '[2, 2, 0, 0]'), ('id1', 'quantizer', 3, '[3, 3, 0, 0]'), ('id1', 'quantizer', 4, '[4, 4, 0, 0]'), ('id1', 'quantizer', 5, '[5, 5, 0, 0]'),
('id1', 'quantizer', 6, '[6, 6, 0, 0]'), ('id1', 'quantizer', 7, '[7, 7, 0, 0]'), ('id1', 'quantizer', 8, '[8, 8, 0, 0]'), ('id1', 'quantizer', 9, '[9, 9, 0, 0]');
CREATE TABLE t1 (
id BIGINT NOT NULL PRIMARY KEY,
vector1 JSON,
name varchar(64),
INDEX key1(vector1) FB_VECTOR_INDEX_TYPE 'ivfflat' fb_vector_dimension 4 FB_VECTOR_TRAINED_INDEX_TABLE 'VECTORDB_DATA' FB_VECTOR_TRAINED_INDEX_ID 'id1'
);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i, json_array(i*0.0001, i*0.0001), concat('val', i) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+1000, json_array(1+i*0.0001, 1-i*0.0001), concat('val', i+1000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+2000, json_array(2+i*0.0001, 2-i*0.0001), concat('val', i+2000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+3000, json_array(3+i*0.0001, 3-i*0.0001), concat('val', i+3000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+4000, json_array(4+i*0.0001, 4-i*0.0001), concat('val', i+4000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+5000, json_array(5+i*0.0001, 5-i*0.0001), concat('val', i+5000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+6000, json_array(6+i*0.0001, 6-i*0.0001), concat('val', i+6000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+7000, json_array(7+i*0.0001, 7-i*0.0001), concat('val', i+7000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+8000, json_array(8+i*0.0001, 8-i*0.0001), concat('val', i+8000) from a);
insert into t1 (WITH RECURSIVE a(i) AS (SELECT 0 union all select i+1 from a where i < 4) SELECT i+9000, json_array(9+i*0.0001, 9-i*0.0001), concat('val', i+9000) from a);
Check the default value of nprobe
show variables like 'fb_vector_search_nprobe';
Variable_name Value
fb_vector_search_nprobe 16

the following select will end up including all available centroids due to the default nprobe value
with search as (select *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100) select COUNT(*) from search;
COUNT(*)
50
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
50 1 5 5 5

Set a lower session default value for nprobe
set session fb_vector_search_nprobe = 3;
show variables like 'fb_vector_search_nprobe';
Variable_name Value
fb_vector_search_nprobe 3

In spite of LIMIT 100, the following select should only have limited results due to the session nprobe value
select *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
1000 [1, 1] val1000 0.9799999594688416
1001 [1.0001, 0.9999] val1001 0.9800000190734863
1002 [1.0002, 0.9998] val1002 0.9800001382827759
1004 [1.0004, 0.9996] val1004 0.9800001978874207
1003 [1.0003, 0.9997] val1003 0.9800002574920654
2001 [2.0001, 1.9999] val2001 5.779999732971191
2003 [2.0003, 1.9997] val2003 5.779999732971191
2000 [2, 2] val2000 5.78000020980835
2002 [2.0002, 1.9998] val2002 5.780000686645508
2004 [2.0004, 1.9996] val2004 5.780000686645508
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
50 2 5 5 5

Query hint can be used to override the default nprobe
The following SELECT should have only 5 result rows
select /*+ SET_VAR(fb_vector_search_nprobe = 1) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
50 3 5 5 5

Set a new, lower, session default value for nprobe
set session fb_vector_search_nprobe = 1;
show variables like 'fb_vector_search_nprobe';
Variable_name Value
fb_vector_search_nprobe 1

Less results expected now
select *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
50 4 5 5 5

Query hint can be used to override the default nprobe and expand search space
The following SELECT should have 15 result rows
select /*+ SET_VAR(fb_vector_search_nprobe = 3) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
1000 [1, 1] val1000 0.9799999594688416
1001 [1.0001, 0.9999] val1001 0.9800000190734863
1002 [1.0002, 0.9998] val1002 0.9800001382827759
1004 [1.0004, 0.9996] val1004 0.9800001978874207
1003 [1.0003, 0.9997] val1003 0.9800002574920654
2001 [2.0001, 1.9999] val2001 5.779999732971191
2003 [2.0003, 1.9997] val2003 5.779999732971191
2000 [2, 2] val2000 5.78000020980835
2002 [2.0002, 1.9998] val2002 5.780000686645508
2004 [2.0004, 1.9996] val2004 5.780000686645508
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
50 5 5 5 5

Add new vectors into the vector database to verify that they become part of the search
insert into t1 values (1010, json_array(0.32, 0.32), 'val1010');
insert into t1 values (1011, json_array(0.28, 0.28), 'val1011');

Newly added vectors should show up in the previous search
select /*+ SET_VAR(fb_vector_search_nprobe = 3) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
1010 [0.32, 0.32] val1010 0.0007999984663911164
1011 [0.28, 0.28] val1011 0.0008000008529052138
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
1000 [1, 1] val1000 0.9799999594688416
1001 [1.0001, 0.9999] val1001 0.9800000190734863
1002 [1.0002, 0.9998] val1002 0.9800001382827759
1004 [1.0004, 0.9996] val1004 0.9800001978874207
1003 [1.0003, 0.9997] val1003 0.9800002574920654
2001 [2.0001, 1.9999] val2001 5.779999732971191
2003 [2.0003, 1.9997] val2003 5.779999732971191
2000 [2, 2] val2000 5.78000020980835
2002 [2.0002, 1.9998] val2002 5.780000686645508
2004 [2.0004, 1.9996] val2004 5.780000686645508
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 6 5 7 5

Valid values for nprobe are between 1 and 10000
By providing nprobe as < 1 or > 10000, the nprobe hint will get truncated to the min/max limits (1 and 10000). Examples below:
select /*+ SET_VAR(fb_vector_search_nprobe = 0) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
1010 [0.32, 0.32] val1010 0.0007999984663911164
1011 [0.28, 0.28] val1011 0.0008000008529052138
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
Warnings:
Warning 1292 Truncated incorrect fb_vector_search_nprobe value: '0'
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 7 5 7 5
with search as (select /*+ SET_VAR(fb_vector_search_nprobe = 10001) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100) select COUNT(*) from search;
COUNT(*)
52
Warnings:
Warning 1292 Truncated incorrect fb_vector_search_nprobe value: '10001'
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 8 5 7 5
with search as (select /*+ SET_VAR(fb_vector_search_nprobe = 1G) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100) select COUNT(*) from search;
COUNT(*)
52
Warnings:
Warning 1292 Truncated incorrect fb_vector_search_nprobe value: '1073741824'
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 9 5 7 5

On providing value other than unsigned int, the nprobe hint will become ineffective, and the session variable value will be used
Example of bad syntax for nprobe query hint
select /*+ SET_VAR(fb_vector_search_nprobe = -1) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
1010 [0.32, 0.32] val1010 0.0007999984663911164
1011 [0.28, 0.28] val1011 0.0008000008529052138
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
Warnings:
Warning 1064 Optimizer hint syntax error near '-1) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) ' at line 1
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 10 5 7 5
select /*+ SET_VAR(fb_vector_search_nprobe = 0.1) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
1010 [0.32, 0.32] val1010 0.0007999984663911164
1011 [0.28, 0.28] val1011 0.0008000008529052138
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
Warnings:
Warning 1232 Incorrect argument type to variable 'fb_vector_search_nprobe'
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 11 5 7 5
select /*+ SET_VAR(fb_vector_search_nprobe = 1.0) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
1010 [0.32, 0.32] val1010 0.0007999984663911164
1011 [0.28, 0.28] val1011 0.0008000008529052138
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
Warnings:
Warning 1232 Incorrect argument type to variable 'fb_vector_search_nprobe'
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 12 5 7 5
select /*+ SET_VAR(fb_vector_search_nprobe =A) */ *, fb_vector_l2(vector1, '[0.3, 0.3]') as dis from t1 force index(key1) order by dis limit 100;
id vector1 name dis
1010 [0.32, 0.32] val1010 0.0007999984663911164
1011 [0.28, 0.28] val1011 0.0008000008529052138
4 [0.0004, 0.0004] val4 0.17952032387256622
3 [0.0003, 0.0003] val3 0.17964020371437073
2 [0.0002, 0.0002] val2 0.17976008355617523
1 [0.0001, 0.0001] val1 0.1798800528049469
0 [0, 0] val0 0.18000000715255737
Warnings:
Warning 1232 Incorrect argument type to variable 'fb_vector_search_nprobe'
SELECT NTOTAL, HIT, MIN_LIST_SIZE, MAX_LIST_SIZE, AVG_LIST_SIZE FROM INFORMATION_SCHEMA.ROCKSDB_VECTOR_INDEX WHERE TABLE_NAME = 't1';
NTOTAL HIT MIN_LIST_SIZE MAX_LIST_SIZE AVG_LIST_SIZE
52 13 5 7 5
drop table t1;
Loading

0 comments on commit 85fd21c

Please sign in to comment.