Como otimizar o desempenho ao usar pgvector no Azure Cosmos DB para PostgreSQL

Artigo
08/15/2024

APLICA-SE A: Azure Cosmos DB para PostgreSQL (alimentado pela extensão de banco de dados Citus para PostgreSQL)

A pgvector extensão adiciona uma pesquisa de semelhança vetorial de código aberto ao PostgreSQL.

Este artigo explora as limitações e compensações e mostra como usar as configurações de pgvector particionamento, indexação e pesquisa para melhorar o desempenho.

Para obter mais informações sobre a extensão em si, consulte as noções básicas de pgvector. Você também pode consultar o LEIA-ME oficial do projeto.

Desempenho

Você deve sempre começar investigando o plano de consulta. Se a sua consulta terminar razoavelmente rápido, execute EXPLAIN (ANALYZE,VERBOSE, BUFFERS).

EXPLAIN (ANALYZE, VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

Para consultas que levam muito tempo para serem executadas, considere descartar a ANALYZE palavra-chave. O resultado contém menos detalhes, mas é fornecido instantaneamente.

EXPLAIN (VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

Sites de terceiros, como explain.depesz.com podem ser úteis para entender os planos de consulta. Algumas perguntas que você deve tentar responder são:

A consulta foi paralelizada?
Foi utilizado um índice?
Utilizei a mesma condição na cláusula WHERE que numa definição de índice parcial?
Se eu usar particionamento, as partições não necessárias foram removidas?

Se seus vetores são normalizados para o comprimento 1, como incorporações OpenAI. Você deve considerar o uso do produto interno (<#>) para obter o melhor desempenho.

Execução paralela

Na saída do seu plano explicativo, procure Workers Planned e Workers Launched (último apenas quando ANALYZE a palavra-chave foi usada). O max_parallel_workers_per_gather parâmetro PostgreSQL define quantos trabalhadores em segundo plano o banco de dados pode iniciar para cada Gather nó e planejar Gather Merge . Aumentar esse valor pode acelerar suas consultas de pesquisa exatas sem precisar criar índices. No entanto, observe que o banco de dados pode não decidir executar o plano em paralelo, mesmo quando esse valor é alto.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 3;

                                        QUERY PLAN
------------------------------------------------------------------------------------------
 Limit  (cost=4214.82..4215.16 rows=3 width=33)
   ->  Gather Merge  (cost=4214.82..13961.30 rows=84752 width=33)
         Workers Planned: 1
         ->  Sort  (cost=3214.81..3426.69 rows=84752 width=33)
               Sort Key: ((embedding <-> '[1,2,3]'::vector))
               ->  Parallel Seq Scan on t_test  (cost=0.00..2119.40 rows=84752 width=33)
(6 rows)

Indexação

Sem índices presentes, a extensão realiza uma pesquisa exata, que proporciona um recall perfeito em detrimento do desempenho.

Para realizar a pesquisa aproximada do vizinho mais próximo, você pode criar índices em seus dados, que negociam recall para desempenho de execução.

Sempre que possível, carregue sempre os dados antes de indexá-los. É mais rápido criar o índice dessa maneira e o layout resultante é mais ideal.

Há dois tipos de índice suportados:

Arquivo invertido com compactação plana (IVVFlat)
Pequenos Mundos Navegáveis Hierárquicos (HNSW)

O IVFFlat índice tem tempos de compilação mais rápidos e usa menos memória do que HNSWo , mas tem menor desempenho de consulta (em termos de compensação de recuperação de velocidade).

Limites

Para indexar uma coluna, ela tem que ter dimensões definidas. A tentativa de indexar uma coluna definida como col vector resulta no erro: ERROR: column does not have dimensions.
Só é possível indexar uma coluna que tenha até 2000 dimensões. A tentativa de indexar uma coluna com mais dimensões resulta no erro: ERROR: column cannot have more than 2000 dimensions for INDEX_TYPE index onde INDEX_TYPE está ou ivfflat hnsw.

Embora você possa armazenar vetores com mais de 2000 dimensões, não é possível indexá-los. Você pode usar a redução de dimensionalidade para se encaixar dentro dos limites. Como alternativa, confie no particionamento e/ou fragmentação com o Azure Cosmos DB para PostgreSQL para obter um desempenho aceitável sem indexação.

Arquivo invertido com compactação plana (IVVFlat)

O ivfflat é um índice para pesquisa de vizinho mais próximo aproximado (ANN). Esse método usa um índice de arquivo invertido para particionar o conjunto de dados em várias listas. O parâmetro probes controla quantas listas são pesquisadas, o que pode melhorar a precisão dos resultados da pesquisa ao custo de uma velocidade de pesquisa mais lenta.

Se o parâmetro probes for definido como o número de listas no índice, todas as listas serão pesquisadas e a pesquisa se tornará uma pesquisa de vizinho mais próximo exata. Nesse caso, o planejador não está usando o índice porque pesquisar todas as listas equivale a executar uma pesquisa de força bruta em todo o conjunto de dados.

O método de indexação particiona o conjunto de dados em várias listas usando o algoritmo de agrupamento k-means. Cada lista contém vetores mais próximos de um determinado centro de cluster. Durante uma pesquisa, o vetor de consulta é comparado aos centros de cluster para determinar quais listas têm maior probabilidade de conter os vizinhos mais próximos. Se o parâmetro probes for definido como 1, apenas a lista correspondente ao centro de cluster mais próximo será pesquisada.

Opções de índice

Selecionar o valor correto para o número de testes a serem executados e os tamanhos das listas pode afetar o desempenho da pesquisa. Bons lugares para começar são:

Use lists igual a rows / 1000 para tabelas com até 1 milhão de linhas e sqrt(rows) para conjuntos de dados maiores.
Para probes começar lists / 10 , para tabelas de até 1 milhão de linhas e sqrt(lists) para conjuntos de dados maiores.

A quantidade de é definida após a criação do lists índice com a lists opção:

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 5000);

Os testes podem ser definidos para toda a conexão ou por transação (usando SET LOCAL dentro de um bloco de transação):

SET ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes

BEGIN;

SET LOCAL ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, one probe

Progresso da indexação

Com o PostgreSQL 12 e versões mais recentes, você pode usar pg_stat_progress_create_index para verificar o progresso da indexação.

SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

As fases para a construção de índices IVFFlat são:

initializing
performing k-means
assigning tuples
loading tuples

Nota

A percentagem de progresso (%) só é preenchida durante a loading tuples fase.

Pequenos Mundos Navegáveis Hierárquicos (HNSW)

O hnsw é um índice para pesquisa de vizinho mais próximo aproximado (ANN) usando o algoritmo Hierarchical Navigable Small Worlds. Ele funciona criando um gráfico em torno de pontos de entrada selecionados aleatoriamente encontrando seus vizinhos mais próximos, o gráfico é então estendido com várias camadas, cada camada inferior contendo mais pontos. Este gráfico multicamadas, quando pesquisado, começa na parte superior, estreitando-se até atingir a camada mais baixa que contém os vizinhos mais próximos da consulta.

A criação desse índice leva mais tempo e memória do que o IVFFlat, no entanto, ele tem uma compensação de recuperação de velocidade melhor. Além disso, não há nenhuma etapa de treinamento como com o IVFFlat, então o índice pode ser criado em uma tabela vazia.

Opções de índice

Ao criar o índice, você pode ajustar dois parâmetros:

m - número máximo de conexões por camada (padrão para 16)
ef_construction - tamanho da lista dinâmica de candidatos usada para a construção de gráficos (padrão 64)

CREATE INDEX t_test_hnsw_l2_idx ON t_test USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);

Durante as consultas, você pode especificar a lista dinâmica de candidatos para pesquisa (o padrão é 40).

A lista dinâmica de candidatos para pesquisa pode ser definida para toda a conexão ou por transação (usando SET LOCAL dentro de um bloco de transação):

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates

BEGIN;

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, 40 candidates

Progresso da indexação

Com o PostgreSQL 12 e versões mais recentes, você pode usar pg_stat_progress_create_index para verificar o progresso da indexação.

SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

As fases para a construção de índices HNSW são:

initializing
loading tuples

Selecionando a função de acesso ao índice

O vector tipo permite realizar três tipos de pesquisas nos vetores armazenados. Você precisa selecionar a função de acesso correta para seu índice para que o banco de dados considere seu índice ao executar suas consultas. Os exemplos demonstram sobre os tipos de ivfflat índice, no entanto, o mesmo pode ser feito para hnsw os índices. A lists opção só se aplica a ivfflat índices.

Distância cosseno

Para pesquisa de semelhança de cosseno, use o método de vector_cosine_ops acesso.

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

Para usar o índice acima, a consulta precisa realizar uma pesquisa de semelhança cosseno, que é feita com o <=> operador.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5;

                                              QUERY PLAN
------------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_cosine_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <=> '[1,2,3]'::vector)
(3 rows)

Distância L2

Para a distância L2 (também conhecida como distância euclidiana), use o método de vector_l2_ops acesso.

CREATE INDEX t_test_embedding_l2_idx ON t_test USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);

Para usar o índice acima, a consulta precisa realizar uma pesquisa de distância L2, que é feita com o <-> operador.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_l2_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <-> '[1,2,3]'::vector)
(3 rows)

Produto interno

Para semelhança interna do produto, use o método de vector_ip_ops acesso.

CREATE INDEX t_test_embedding_ip_idx ON t_test USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);

Para usar o índice acima, a consulta precisa realizar uma pesquisa interna de similaridade de produto, que é feita com o <#> operador.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <#> '[1,2,3]' LIMIT 5;

                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_ip_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <#> '[1,2,3]'::vector)
(3 rows)

Índices parciais

Em alguns cenários, é benéfico ter um índice que cubra apenas um conjunto parcial de dados. Podemos, por exemplo, criar um índice apenas para os nossos utilizadores premium:

CREATE INDEX t_premium ON t_test USING ivfflat (vec vector_ip_ops) WITH (lists = 100) WHERE tier = 'premium';

Agora podemos ver que o nível premium agora usa o índice:

explain select * from t_test where tier = 'premium' order by vec <#> '[2,2,2]';

                                     QUERY PLAN
------------------------------------------------------------------------------------
 Index Scan using t_premium on t_test  (cost=65.57..25638.05 rows=245478 width=39)
   Order By: (vec <#> '[2,2,2]'::vector)
(2 rows)

Enquanto os usuários do nível gratuito não têm o benefício:

explain select * from t_test where tier = 'free' order by vec <#> '[2,2,2]';

                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44019.01..44631.37 rows=244941 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15395.25 rows=244941 width=39)
         Filter: (tier = 'free'::text)
(4 rows)

Ter apenas um subconjunto de dados indexados significa que o índice ocupa menos espaço no disco e é mais rápido de pesquisar.

O PostgreSQL pode não reconhecer que o índice é seguro de usar se o WHERE formulário usado na cláusula da definição de índice parcial não corresponder ao usado em suas consultas. Em nosso conjunto de dados de exemplo, temos apenas os valores 'free''test' exatos e 'premium' como os valores distintos da coluna de camada. Mesmo com uma consulta usando tier LIKE 'premium' PostgreSQL não está usando o índice.

explain select * from t_test where tier like 'premium' order by vec <#> '[2,2,2]';

                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44086.30..44700.00 rows=245478 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15396.59 rows=245478 width=39)
         Filter: (tier ~~ 'premium'::text)
(4 rows)

Criação de partições

Uma maneira de melhorar o desempenho é dividir o conjunto de dados em várias partições. Podemos imaginar um sistema em que é natural fazer referência a dados apenas do ano atual ou talvez dos últimos dois anos. Em tal sistema, você pode particionar seus dados por um intervalo de datas e, em seguida, capitalizar sobre o desempenho melhorado quando o sistema é capaz de ler apenas as partições relevantes, conforme definido pelo ano consultado.

Vamos definir uma tabela particionada:

CREATE TABLE t_test_partitioned(vec vector(3), vec_date date default now()) partition by range (vec_date);

Podemos criar partições manualmente para cada ano ou usar a função de utilitário Citus (disponível no Cosmos DB para PostgreSQL).

    select create_time_partitions(
      table_name         := 't_test_partitioned',
      partition_interval := '1 year',
      start_from         := '2020-01-01'::timestamptz,
      end_at             := '2024-01-01'::timestamptz
    );

Verifique as partições criadas:

\d+ t_test_partitioned

                                Partitioned table "public.t_test_partitioned"
  Column  |   Type    | Collation | Nullable | Default | Storage  | Compression | Stats target | Description
----------+-----------+-----------+----------+---------+----------+-------------+--------------+-------------
 vec      | vector(3) |           |          |         | extended |             |              |
 vec_date | date      |           |          | now()   | plain    |             |              |
Partition key: RANGE (vec_date)
Partitions: t_test_partitioned_p2020 FOR VALUES FROM ('2020-01-01') TO ('2021-01-01'),
            t_test_partitioned_p2021 FOR VALUES FROM ('2021-01-01') TO ('2022-01-01'),
            t_test_partitioned_p2022 FOR VALUES FROM ('2022-01-01') TO ('2023-01-01'),
            t_test_partitioned_p2023 FOR VALUES FROM ('2023-01-01') TO ('2024-01-01')

Para criar manualmente uma partição:

CREATE TABLE t_test_partitioned_p2019 PARTITION OF t_test_partitioned FOR VALUES FROM ('2019-01-01') TO ('2020-01-01');

Em seguida, certifique-se de que suas consultas realmente filtram para um subconjunto de partições disponíveis. Por exemplo, na consulta abaixo, filtramos para duas partições:

explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01';

                                                                  QUERY PLAN
-----------------------------------------------------------------------------------------------------------------------------------------------
 Append  (cost=0.00..58.16 rows=12 width=36) (actual time=0.014..0.018 rows=3 loops=1)
   ->  Seq Scan on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.00..29.05 rows=6 width=36) (actual time=0.013..0.014 rows=1 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
   ->  Seq Scan on t_test_partitioned_p2023 t_test_partitioned_2  (cost=0.00..29.05 rows=6 width=36) (actual time=0.002..0.003 rows=2 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.125 ms
 Execution Time: 0.036 ms

Você pode indexar uma tabela particionada.

CREATE INDEX ON t_test_partitioned USING ivfflat (vec vector_cosine_ops) WITH (lists = 100);

explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01' order by vec <=> '[1,2,3]' limit 5;

                                                                                         QUERY PLAN
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Limit  (cost=4.13..12.20 rows=2 width=44) (actual time=0.040..0.042 rows=1 loops=1)
   ->  Merge Append  (cost=4.13..12.20 rows=2 width=44) (actual time=0.039..0.040 rows=1 loops=1)
         Sort Key: ((t_test_partitioned.vec <=> '[1,2,3]'::vector))
         ->  Index Scan using t_test_partitioned_p2022_vec_idx on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.04..4.06 rows=1 width=44) (actual time=0.022..0.023 rows=0 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
         ->  Index Scan using t_test_partitioned_p2023_vec_idx on t_test_partitioned_p2023 t_test_partitioned_2  (cost=4.08..8.11 rows=1 width=44) (actual time=0.015..0.016 rows=1 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.167 ms
 Execution Time: 0.139 ms
(11 rows)

Conclusão

Parabéns, você acabou de aprender as compensações, limitações e melhores práticas para alcançar o melhor desempenho com pgvectoro .

Partilhar via

Como otimizar o desempenho ao usar pgvector no Azure Cosmos DB para PostgreSQL

Desempenho

Execução paralela

Indexação

Limites

Arquivo invertido com compactação plana (IVVFlat)

Opções de índice

Progresso da indexação

Pequenos Mundos Navegáveis Hierárquicos (HNSW)

Opções de índice

Progresso da indexação

Selecionando a função de acesso ao índice

Distância cosseno

Distância L2

Produto interno

Índices parciais

Criação de partições

Conclusão

Comentários

Recursos adicionais