Compartilhar via

Como otimizar o desempenho ao usar pgvector no Banco de Dados do Azure para PostgreSQL – Servidor Flexível

  • Artigo

APLICA-SE A: Banco de Dados do Azure para PostgreSQL - Servidor Flexível

A extensão pgvector adiciona uma pesquisa de similaridade de vetor de código aberto ao servidor flexível do Banco de Dados do Azure para PostgreSQL.

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

Para obter mais informações sobre a extensão, confiranoções básicas de pgvector. Talvez você também queira consultar o LEIAME oficial do projeto.

Desempenho

Você sempre deve começar investigando o plano de consulta. Se a consulta for encerrada de forma razoavelmente rápida, 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 remover a palavra-chave ANALYZE. 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:

Se os vetores forem normalizados para o tamanho 1, como inserções do OpenAI. Você deve considerar o uso do produto interno (<#>) para obter o melhor desempenho.

Execução paralela

Na saída do plano de explicação, procure Workers Planned e Workers Launched (o último somente quando a palavra-chave ANALYZE foi usada). O parâmetro max_parallel_workers_per_gather do PostgreSQL define quantos trabalhos em segundo plano o banco de dados pode iniciar para cada nó de plano Gather e Gather Merge. O aumento nesse valor pode acelerar suas consultas de pesquisa exatas sem precisar criar índices. No entanto, observe que o banco de dados pode decidir não executar o plano em paralelo, mesmo quando esse valor for 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 executa uma pesquisa exata, que fornece um recall perfeito, mas reduz o desempenho.

Para executar uma pesquisa aproximada do vizinho mais próximo, você pode criar índices em seus dados, que trocam o recall pelo desempenho de execução.

Quando possível, sempre carregue seus dados antes de indexá-los. É mais rápido criar o índice dessa maneira e o layout resultante é melhor.

Há dois tipos de índice com suporte:

O índice IVFFlat leva menos tempo para ser criado e usa menos memória do que o HNSW, mas tem um desempenho de consulta inferior (em termos de compensação de recuperação de velocidade).

Limites

  • Para indexar uma coluna, ela precisa ter dimensões definidas. Tentar indexar uma coluna definida como col vector resulta no erro: ERROR: column does not have dimensions.
  • Você só pode indexar uma coluna que tenha até 2.000 dimensões. Tentar 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 é ivfflat ou hnsw.

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

IVFFlat (Arquivo invertido com compactação plana)

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

Se o parâmetro de investigação for definido como o número de listas no índice, todas as listas serão pesquisadas e a pesquisa se tornará uma pesquisa exata do vizinho mais próximo. Nesse caso, o planejador não está usando o índice porque a pesquisa de todas as listas é equivalente 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 clustering K-means. Cada lista contém vetores mais próximos de um centro de cluster específico. 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 de investigação for definido como 1, somente a lista correspondente ao centro de cluster mais próximo será pesquisada.

Opções de índice

A seleção do valor correto para o número de investigações a serem executadas e os tamanhos das listas podem afetar o desempenho da pesquisa. Algumas opções para começar são:

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

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

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

As investigações podem ser definidas 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 ou posterior, 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;

Fases para construir os índices IVFFlat são:

  1. initializing
  2. performing k-means
  3. assigning tuples
  4. loading tuples

Observação

Percentual de progresso (%) somente é preenchido na fase loading tuples.

HNSW (Hierarchical Navigable Small Worlds)

O hnsw é um índice para pesquisa de vizinho mais próximo aproximado (ANN) usando o algoritmo HNSW. Ele funciona criando um grafo em torno de pontos de entrada selecionados aleatoriamente, encontrando seus vizinhos mais próximos. O gráfico é então estendido com múltiplas camadas, cada camada inferior contendo mais pontos. Esse gráfico multicamadas, quando pesquisado, começa no topo, estreitando-se até atingir a camada mais baixa que contém os vizinhos mais próximos da consulta.

Construir esse índice leva mais tempo e memória do que o índice IVFFlat, porém tem uma melhor compensação entre velocidade e recuperação. Além disso, não há etapa de treinamento como no IVFFlat, portanto o índice pode ser criado em uma tabela vazia.

Opções de índice

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

  1. m - número máximo de conexões por camada (o padrão é 16)
  2. ef_construction - tamanho da lista dinâmica de candidatos usada para a construção do grafo (o 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 pra pesquisa (o padrão é 40).

A lista dinâmica de candidatos 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 ou posterior, 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;

Fases para criar os índices HNSW são:

  1. initializing
  2. loading tuples

Seleção da função de acesso ao índice

O tipo vector permite que você execute três tipos de pesquisas nos vetores armazenados. Você precisa selecionar a função de acesso correta para o índice para que o banco de dados considere seu índice ao executar suas consultas. Os exemplo usam os tipos de índice ivfflat, mas o mesmo pode ser feito para índices hnsw. A opção lists aplica-se somente a índices ivfflat.

Distância do cosseno

Para pesquisa de similaridade de cosseno, use o método de acesso vector_cosine_ops.

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 executar uma pesquisa de similaridade de 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 distância L2 (também conhecida como distância euclidiana), use o método de acesso vector_l2_ops.

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 executar 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 obter similaridade interna do produto, use o método de acesso vector_ip_ops.

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 executar uma pesquisa de similaridade de do produto interno, 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 abrange apenas um conjunto parcial dos dados. Podemos, por exemplo, criar um índice apenas para nossos usuários Premium:

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

Podemos ver que a camada 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 da camada gratuita 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 indexado significa que o índice ocupa menos espaço no disco e pode ser pesquisado mais rapidamente.

O PostgreSQL pode não reconhecer que o índice é seguro para uso se o formulário usado na cláusula WHERE 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 exatos 'free', 'test' e 'premium' como os valores distintos da coluna de camada. Mesmo com uma consulta usando tier LIKE 'premium', o 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)

Particionamento

Uma maneira de melhorar o desempenho é dividir o conjunto de dados em várias partições. Podemos imaginar um sistema onde é natural se referir a dados apenas do ano atual ou talvez dos últimos dois anos. Nesse sistema, você pode particionar seus dados por um intervalo de datas e, em seguida, aproveitar melhor o desempenho quando o sistema puder 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 for 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, verifique se as consultas realmente filtram um subconjunto de partições disponíveis. Por exemplo, na consulta abaixo, filtramos 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)

Próximas etapas

Parabéns, você acabou de aprender as compensações, limitações e práticas recomendadas para obter o melhor desempenho com pgvector.

Gerar inserções de vetores com o OpenAI do Azure no Banco de Dados do Azure para PostgreSQL – Servidor Flexível