Criar índices e vetores de consulta

Esta página descreve como usar embeddings armazenados para gerar índices e consultar embeddings usando os índices ScaNN, IVF, IVFFlat e HNSW com o AlloyDB para PostgreSQL. Para mais informações sobre como armazenar embeddings, consulte Armazenar embeddings de vetor.

Antes de começar

Antes de começar a criar índices, conclua os seguintes pré-requisitos.

  • Os vetores de embedding são adicionados a uma tabela no seu banco de dados do AlloyDB.

  • A versão 0.5.0 ou mais recente da extensão vector, que é baseada em pgvector, estendida pelo Google para AlloyDB, é instalada.

    CREATE EXTENSION IF NOT EXISTS vector;

  • Para gerar índices ScaNN, instale a extensão alloydb_scann, além da extensão vector.

    CREATE EXTENSION IF NOT EXISTS alloydb_scann;

Criar um índice

É possível criar um dos seguintes tipos de índice para tabelas no seu banco de dados.

Criar um índice ScaNN

AlloyDB alloydb_scann, uma extensão do PostgreSQL desenvolvida pelo Google que implementa um índice de vizinho mais próximo altamente eficiente, com base no algoritmo ScaNN.

O índice ScaNN é um índice de quantização baseado em árvore para pesquisa aproximada de vizinho mais próximo. Ele fornece um tempo de criação de índice menor e uma pegada de memória menor em comparação com HNSW. Além disso, ele oferece QPS mais rápido em comparação com HNSW com base na carga de trabalho.

Índice ScaNN de árvore de dois níveis

Para aplicar um índice de árvore de dois níveis usando o algoritmo ScaNN a uma coluna que contém embeddings de vetor armazenados, execute a consulta DDL a seguir:

CREATE INDEX INDEX_NAME ON TABLE
  USING scann (EMBEDDING_COLUMN DISTANCE_FUNCTION)
  WITH (num_leaves=NUM_LEAVES_VALUE);

Substitua:

  • INDEX_NAME: o nome do índice que você quer criar, por exemplo, my-scann-index. Os nomes dos índices são compartilhados no seu banco de dados. Verifique se cada nome de índice é exclusivo para cada tabela no seu banco de dados.

  • TABLE: a tabela em que o índice será adicionado.

  • EMBEDDING_COLUMN: uma coluna que armazena dados de vector.

  • DISTANCE_FUNCTION: a função de distância a ser usada com esse índice. Escolha uma destas opções:

    • Distância de L2:l2

    • Produto escalar: dot_product

    • Distância do cosseno:cosine

  • NUM_LEAVES_VALUE: o número de partições a serem aplicadas a esse índice. Defina qualquer valor entre 1 e 1.048.576. Para mais informações sobre como decidir esse valor, consulte Ajustar um índice ScaNN.

Índice ScaNN de árvore de três níveis

Para criar um índice de árvore de três níveis usando o algoritmo ScaNN em uma coluna que contém embeddings de vetor armazenados, execute a seguinte consulta DDL:

CREATE INDEX INDEX_NAME ON TABLE
  USING scann (EMBEDDING_COLUMN DISTANCE_FUNCTION)
  WITH (num_leaves=NUM_LEAVES_VALUE, max_num_levels = MAX_NUM_LEVELS);

Substitua:

  • MAX_NUM_LEVELS: o número máximo de níveis da árvore de agrupamento K-means. Defina como 1(padrão) para quantização baseada em árvore de dois níveis e como 2 para quantização baseada em árvore de três níveis.

Depois de criar o índice, você pode executar consultas de pesquisa de vizinho mais próximo que fazem uso do índice seguindo as instruções em Fazer uma consulta de vizinho mais próximo com um texto fornecido.

Os parâmetros do índice precisam ser definidos para encontrar o equilíbrio certo entre QPS e resposta. Para mais informações sobre como ajustar o índice ScaNN, consulte Ajustar um índice ScaNN.

Para criar esse índice em uma coluna incorporada que usa o tipo de dados real[] em vez de vector, converta a coluna no tipo de dados vector:

CREATE INDEX INDEX_NAME ON TABLE
  USING scann (CAST(EMBEDDING_COLUMN AS vector(DIMENSIONS)) DISTANCE_FUNCTION)
  WITH (num_leaves=NUM_LEAVES_VALUE, max_num_levels = MAX_NUM_LEVELS);

Substitua DIMENSIONS pela largura dimensional da coluna de embedding. Para mais informações sobre como encontrar as dimensões, consulte a função vector_dims em Funções vetoriais.

Para conferir o progresso da indexação, use a visualização pg_stat_progress_create_index:

SELECT * FROM pg_stat_progress_create_index;

A coluna phase mostra o estado atual da criação do índice, e a fase building index: tree training desaparece depois que o índice é criado.

Para ajustar o índice para uma recuperação desejada e um equilíbrio de QPS, consulte Ajustar um índice ScaNN.

Executar uma consulta

Depois de armazenar e indexar embeddings no banco de dados, você pode começar a fazer consultas usando a funcionalidade de consulta pgvector. Não é possível executar consultas de pesquisa em massa usando a extensão alloydb_scann.

Para encontrar os vizinhos semânticos mais próximos de um vetor de incorporação, execute a consulta de exemplo a seguir, em que você define a mesma função de distância usada durante a criação do índice.

  SELECT * FROM TABLE
    ORDER BY EMBEDDING_COLUMN DISTANCE_FUNCTION_QUERY ['EMBEDDING']
    LIMIT ROW_COUNT

Substitua:

  • TABLE: a tabela que contém o embedding para comparar o texto.

  • INDEX_NAME: o nome do índice que você quer usar. Por exemplo, my-scann-index.

  • EMBEDDING_COLUMN: a coluna que contém os embeddings armazenados.

  • DISTANCE_FUNCTION_QUERY: a função de distância a ser usada com essa consulta. Escolha uma das seguintes opções com base na função de distância usada ao criar o índice:

    • Distância de L2:<->

    • Produto interno:<#>

    • Distância do cosseno:<=>

  • EMBEDDING: o vetor de embedding para encontrar os vizinhos semânticos armazenados mais próximos.

  • ROW_COUNT: o número de linhas que serão retornadas.

    Especifique 1 se quiser apenas a melhor correspondência.

Para mais informações sobre outros exemplos de consulta, consulte Como consultar.

Você também pode usar a função embedding() para converter o texto em um vetor. Você aplica o vetor a um dos operadores de vizinho mais próximo pgvector, <-> para distância L2, para encontrar as linhas do banco de dados com os embeddings mais semanticamente semelhantes.

Como embedding() retorna uma matriz real, é necessário transmitir a chamada embedding() para vector para usar esses valores com operadores pgvector.

  CREATE EXTENSION IF NOT EXISTS google_ml_integration VERSION '1.2';
  CREATE EXTENSION IF NOT EXISTS vector;

  SELECT * FROM TABLE
    ORDER BY EMBEDDING_COLUMN::vector
    <-> embedding('MODEL_IDVERSION_TAG', 'TEXT')
    LIMIT ROW_COUNT

Substitua:

  • MODEL_ID: o ID do modelo a ser consultado.

    Se você estiver usando o Model Garden da Vertex AI, especifique text-embedding-005 como o ID do modelo. Esses são os modelos baseados na nuvem que o AlloyDB pode usar para embeddings de texto. Para mais informações, consulte Embeddings de texto.

  • Opcional: VERSION_TAG: a tag da versão do modelo a ser consultada. Adicione @ ao início da tag.

    Se você estiver usando um dos modelos em inglês text-embedding-005 com a Vertex AI, especifique uma das tags de versão, por exemplo, text-embedding-005, listada em Versões de modelo.

    O Google recomenda que você sempre especifique a tag de versão. Se você não especificar a tag de versão, o AlloyDB sempre usará a versão mais recente do modelo, o que pode levar a resultados inesperados.

  • TEXT: o texto a ser traduzido em um embedding vetorial.

A seguir