Adding a GIST index actually seems to make K-Nearest Neighbor (KNN) ORDER BY queries on cube columns worse in PostgreSQL. Why would that be, and what can be done about it?
Here's what I mean. In a PostgreSQL database I have a table whose DDL is create sample (id serial primary key, title text, embedding cube) where the embedding column is an embedding vector of the title obtained with a Google language model. The cube data type is provided by the cube extension, which I have installed. Incidentally, these are titles of Wikipedia articles. In any case, there are 1 million records. I then perform a KNN query with the following query. This query defines distance using the Euclidean distance operator <->, though results are similar for the other two metrics. It does an ORDER BY and applies a LIMIT in order to find 10 Wikipedia articles with "similar" titles (the most similar being the target title itself). That all works fine.
select sample.title, sample.embedding <-> cube('(0.18936706, -0.12455666, -0.31581765, 0.0192692, -0.07364611, 0.07851536, 0.0290586, -0.02582532, -0.03378124, -0.10564457, -0.03903799, 0.08668878, -0.15357816, -0.17793414, -0.01826405, 0.01969068, 0.11386908, 0.1555583, 0.09368557, 0.13697313, -0.05610929, -0.06536788, -0.12212707, 0.26356605, -0.06004387, -0.01966437, -0.1250324, -0.16645767, -0.13525756, 0.22482251, -0.1709727, 0.28966117, -0.07927769, -0.02498624, -0.10018375, -0.10923951, 0.04770213, 0.11573371, 0.04619929, 0.05216618, 0.19176421, 0.12948817, 0.08719034, -0.16109011, -0.02411379, -0.05638905, -0.37334979, 0.31225419, 0.0744801, 0.27044332)') distance from sample order by distance limit 10;
What's puzzling to me, however, is that, if I put a GIST index on the embedding column, the query performance actually is worse. Adding the index, the query plan changes as expected, in the way expected, insofar as it uses the index. But...it gets slower!
This seems to run contrary to the documentation for cube which states:
In addition, a cube GiST index can be used to find nearest neighbors using the metric operators <->, <#>, and <=> in ORDER BY clauses
They even provide an example query, which is very similar to mine.
SELECT c FROM test ORDER BY c <-> cube(array[0.5,0.5,0.5]) LIMIT 1
Here's the query plan and timing info before dropping the index.
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Limit (cost=0.41..6.30 rows=10 width=29)
-> Index Scan using sample_embedding_idx on sample (cost=0.41..589360.33 rows=999996 width=29)
Order By: (embedding <-> '(0.18936706, -0.12455666, -0.31581765, 0.0192692, -0.07364611, 0.07851536, 0.0290586, -0.02582532, -0.03378124, -0.10564457, -0.03903799, 0.08668878, -0.15357816, -0.17793414, -0.01826405, 0.01969068, 0.11386908, 0.1555583, 0.09368557, 0.13697313, -0.05610929, -0.06536788, -0.12212707, 0.26356605, -0.06004387, -0.01966437, -0.1250324, -0.16645767, -0.13525756, 0.22482251, -0.1709727, 0.28966117, -0.07927769, -0.02498624, -0.10018375, -0.10923951, 0.04770213, 0.11573371, 0.04619929, 0.05216618, 0.19176421, 0.12948817, 0.08719034, -0.16109011, -0.02411379, -0.05638905, -0.37334979, 0.31225419, 0.0744801, 0.27044332)'::cube)
(3 rows)
title | distance
----------------------+--------------------
david petrarca | 0.5866321762629475
david adamski | 0.5866321762629475
richard ansdell | 0.6239883862603475
linda darke | 0.6392124797481789
ilias tsiliggiris | 0.6996660649119893
watson, jim | 0.7059481479504834
sk radni%c4%8dki | 0.71718948226995
burnham, pa | 0.7384858030758069
arthur (europa-park) | 0.7468462897336924
ivan kecojevic | 0.7488206082281348
(10 rows)
Time: 1226.457 ms (00:01.226)
And, here's the query plan and timing info after dropping the index.
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Limit (cost=74036.32..74037.48 rows=10 width=29)
-> Gather Merge (cost=74036.32..171264.94 rows=833330 width=29)
Workers Planned: 2
-> Sort (cost=73036.29..74077.96 rows=416665 width=29)
Sort Key: ((embedding <-> '(0.18936706, -0.12455666, -0.31581765, 0.0192692, -0.07364611, 0.07851536, 0.0290586, -0.02582532, -0.03378124, -0.10564457, -0.03903799, 0.08668878, -0.15357816, -0.17793414, -0.01826405, 0.01969068, 0.11386908, 0.1555583, 0.09368557, 0.13697313, -0.05610929, -0.06536788, -0.12212707, 0.26356605, -0.06004387, -0.01966437, -0.1250324, -0.16645767, -0.13525756, 0.22482251, -0.1709727, 0.28966117, -0.07927769, -0.02498624, -0.10018375, -0.10923951, 0.04770213, 0.11573371, 0.04619929, 0.05216618, 0.19176421, 0.12948817, 0.08719034, -0.16109011, -0.02411379, -0.05638905, -0.37334979, 0.31225419, 0.0744801, 0.27044332)'::cube))
-> Parallel Seq Scan on sample (cost=0.00..64032.31 rows=416665 width=29)
(6 rows)
title | distance
----------------------+--------------------
david petrarca | 0.5866321762629475
david adamski | 0.5866321762629475
richard ansdell | 0.6239883862603475
linda darke | 0.6392124797481789
ilias tsiliggiris | 0.6996660649119893
watson, jim | 0.7059481479504834
sk radni%c4%8dki | 0.71718948226995
burnham, pa | 0.7384858030758069
arthur (europa-park) | 0.7468462897336924
ivan kecojevic | 0.7488206082281348
(10 rows)
Time: 381.419 ms
Notice:
- With Index: 1226.457 ms
- Without Index: 381.419 ms
This very puzzling behavior! All of it is documented in a GitHub repo so that others can try it. I'll add documentation about how to generate the embedding vectors, but that shouldn't be needed, as in the Quick-Start I show that pre-computed embedding vectors can be downloaded from my Google Drive folder.
Addendum
It was asked in the comments below to provide the output of explain (analyze, buffers). Here that is, where
- I re-create the (covering) index
- I run the query with
explain (analyze, buffers) - I drop the index
- I run the query with
explain (analyze, buffers)again
pgbench=# create index on sample using gist (embedding) include (title);
CREATE INDEX
Time: 51966.315 ms (00:51.966)
pgbench=#
QUERY PLAN
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Limit (cost=0.41..4.15 rows=10 width=29) (actual time=3215.956..3216.667 rows=10 loops=1)
Buffers: shared hit=1439 read=87004 written=7789
-> Index Only Scan using sample_embedding_title_idx on sample (cost=0.41..373768.39 rows=999999 width=29) (actual time=3215.932..3216.441 rows=10 loops=1)
Order By: (embedding <-> '(0.18936706, -0.12455666, -0.31581765, 0.0192692, -0.07364611, 0.07851536, 0.0290586, -0.02582532, -0.03378124, -0.10564457, -0.03903799, 0.08668878, -0.15357816, -0.17793414, -0.01826405, 0.01969068, 0.11386908, 0.1555583, 0.09368557, 0.13697313, -0.05610929, -0.06536788, -0.12212707, 0.26356605, -0.06004387, -0.01966437, -0.1250324, -0.16645767, -0.13525756, 0.22482251, -0.1709727, 0.28966117, -0.07927769, -0.02498624, -0.10018375, -0.10923951, 0.04770213, 0.11573371, 0.04619929, 0.05216618, 0.19176421, 0.12948817, 0.08719034, -0.16109011, -0.02411379, -0.05638905, -0.37334979, 0.31225419, 0.0744801, 0.27044332)'::cube)
Heap Fetches: 0
Buffers: shared hit=1439 read=87004 written=7789
Planning:
Buffers: shared hit=14 read=6 dirtied=2
Planning Time: 0.432 ms
Execution Time: 3316.266 ms
(10 rows)
Time: 3318.333 ms (00:03.318)
pgbench=# drop index sample_embedding_title_idx;
DROP INDEX
Time: 182.324 ms
pgbench=#
QUERY PLAN
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Limit (cost=74036.35..74037.52 rows=10 width=29) (actual time=6052.845..6057.210 rows=10 loops=1)
Buffers: shared hit=70 read=58830
-> Gather Merge (cost=74036.35..171265.21 rows=833332 width=29) (actual time=6052.825..6057.021 rows=10 loops=1)
Workers Planned: 2
Workers Launched: 2
Buffers: shared hit=70 read=58830
-> Sort (cost=73036.33..74077.99 rows=416666 width=29) (actual time=6002.928..6003.019 rows=8 loops=3)
Sort Key: ((embedding <-> '(0.18936706, -0.12455666, -0.31581765, 0.0192692, -0.07364611, 0.07851536, 0.0290586, -0.02582532, -0.03378124, -0.10564457, -0.03903799, 0.08668878, -0.15357816, -0.17793414, -0.01826405, 0.01969068, 0.11386908, 0.1555583, 0.09368557, 0.13697313, -0.05610929, -0.06536788, -0.12212707, 0.26356605, -0.06004387, -0.01966437, -0.1250324, -0.16645767, -0.13525756, 0.22482251, -0.1709727, 0.28966117, -0.07927769, -0.02498624, -0.10018375, -0.10923951, 0.04770213, 0.11573371, 0.04619929, 0.05216618, 0.19176421, 0.12948817, 0.08719034, -0.16109011, -0.02411379, -0.05638905, -0.37334979, 0.31225419, 0.0744801, 0.27044332)'::cube))
Sort Method: top-N heapsort Memory: 26kB
Buffers: shared hit=70 read=58830
Worker 0: Sort Method: top-N heapsort Memory: 26kB
Worker 1: Sort Method: top-N heapsort Memory: 26kB
-> Parallel Seq Scan on sample (cost=0.00..64032.33 rows=416666 width=29) (actual time=0.024..3090.103 rows=333333 loops=3)
Buffers: shared read=58824
Planning:
Buffers: shared hit=3 read=3 dirtied=1
Planning Time: 0.129 ms
Execution Time: 6057.388 ms
(18 rows)
Time: 6053.284 ms (00:06.053)
