Encontrar vizinhos aproximados mais próximos (ANN) e embeddings de vetor de consulta

Esta página descreve como encontrar vizinhos mais próximos aproximados (ANN, na sigla em inglês) e consultar embeddings de vetor usando as funções de distância ANN.

Quando um conjunto de dados é pequeno, é possível usar o vizinho mais próximo (KNN, na sigla em inglês) para encontrar os vetores k mais próximos exatos. No entanto, à medida que o conjunto de dados cresce, a latência e o custo de uma pesquisa de KNN também aumentam. É possível usar o ANN para encontrar os k vizinhos mais próximos aproximados com latência e custo significativamente reduzidos.

Em uma pesquisa de ANN, os vetores k retornados não são os k vizinhos mais próximos verdadeiros, porque a pesquisa de ANN calcula distâncias aproximadas e pode não analisar todos os vetores no conjunto de dados. Ocasionalmente, alguns vetores que não estão entre os k vizinhos mais próximos são retornados. Isso é conhecido como perda de recall. A quantidade de perda de recall aceitável depende do caso de uso, mas, na maioria dos casos, perder um pouco de recall em troca de um desempenho aprimorado do banco de dados é uma compensação aceitável.

Para mais detalhes sobre as funções de distância aproximada com suporte no Spanner, consulte as seguintes páginas de referência para o dialeto do banco de dados:

Consultar embeddings de vetor

O Spanner acelera as pesquisas de vetor do vizinho mais próximo aproximado (ANN) usando um índice vetorial. É possível usar um índice vetorial para consultar embeddings de vetor. Para consultar embeddings de vetor, você deve primeiro criar um índice vetorial. Em seguida, use uma das três funções de distância aproximada para encontrar o ANN.

As restrições ao usar as funções de distância aproximada incluem o seguinte:

  • A função de distância aproximada precisa calcular a distância entre uma coluna de embedding e uma expressão constante (por exemplo, um parâmetro ou um literal).
  • A saída da função de distância aproximada precisa ser usada em uma cláusula ORDER BY como a única chave de classificação, e um LIMIT precisa ser especificado após o ORDER BY.
  • A consulta precisa filtrar explicitamente as linhas que não estão indexadas. Na maioria dos casos, isso significa que a consulta precisa incluir uma WHERE <column_name> IS NOT NULL cláusula que corresponda à definição do índice vetorial, a menos que a coluna já esteja marcada como NOT NULL na definição da tabela.

Para uma lista detalhada de limitações, consulte a página de referência da função de distância aproximada.

Exemplos

Considere uma tabela Documents que tenha uma coluna DocEmbedding de embeddings de texto pré-computados da coluna de bytes DocContents e uma NullableDocEmbedding coluna preenchida de outras fontes que podem ser nulas.

GoogleSQL 

CREATE TABLE Documents (
  UserId       INT64 NOT NULL,
  DocId        INT64 NOT NULL,
  Author       STRING(1024),
  DocContents  BYTES(MAX),
  DocEmbedding ARRAY<FLOAT32> NOT NULL,
  NullableDocEmbedding ARRAY<FLOAT32>,
  WordCount    INT64
) PRIMARY KEY (UserId, DocId);

PostgreSQL 

CREATE TABLE documents (
  user_id      bigint not null,
  doc_id       bigint not null,
  author       varchar(1024),
  doc_contents bytea,
  doc_embedding float4[] not null,
  nullable_doc_embedding float4[],
  word_count   bigint,
  PRIMARY KEY (user_id, doc_id)
);

Para pesquisar os 100 vetores mais próximos de [1.0, 2.0, 3.0]:

GoogleSQL 

SELECT DocId
FROM Documents
WHERE WordCount > 1000
ORDER BY APPROX_EUCLIDEAN_DISTANCE(
  ARRAY<FLOAT32>[1.0, 2.0, 3.0], DocEmbedding,
  options => JSON '{"num_leaves_to_search": 10}')
LIMIT 100

PostgreSQL 

SELECT doc_id
FROM documents
WHERE word_count > 1000
ORDER BY spanner.approx_euclidean_distance(
  ARRAY[1.0, 2.0, 3.0]::float4[], doc_embedding,
  options=>jsonb'{"num_leaves_to_search": 10}'
)
LIMIT 100;

Para pesquisar os 100 vetores mais próximos de um embedding gerado por uma expressão SQL, use o padrão a seguir. Neste exemplo, a consulta está encontrando os documentos mais semelhantes ao embedding de UserId = 1 e DocId = 1:

GoogleSQL 

WITH emb AS (
  SELECT DocEmbedding AS value
  FROM Documents
  WHERE UserId = 1 AND DocId = 1
  LIMIT 1
)
SELECT DocId
FROM Documents, emb
ORDER BY APPROX_EUCLIDEAN_DISTANCE(
  emb.value, DocEmbedding,
  options => JSON '{"num_leaves_to_search": 10}')
LIMIT 100

PostgreSQL 

SELECT documents.doc_id
FROM
  documents,
  (
    SELECT doc_embedding AS value
    FROM documents
    WHERE user_id = 1 AND doc_id = 1
    LIMIT 1
  ) vector
WHERE documents.doc_embedding IS NOT NULL
ORDER BY spanner.APPROX_EUCLIDEAN_DISTANCE(documents.doc_embedding,
         vector.value, options=>'{"num_leaves_to_search": 10}'::jsonb)
LIMIT 100

Se a coluna de embedding for anulável:

GoogleSQL 

SELECT DocId
FROM Documents
WHERE NullableDocEmbedding IS NOT NULL AND WordCount > 1000
ORDER BY APPROX_EUCLIDEAN_DISTANCE(
  ARRAY<FLOAT32>[1.0, 2.0, 3.0], NullableDocEmbedding,
  options => JSON '{"num_leaves_to_search": 10}')
LIMIT 100

PostgreSQL 

SELECT doc_id
FROM documents
WHERE nullable_doc_embedding IS NOT NULL AND word_count > 1000
ORDER BY spanner.approx_euclidean_distance(
  ARRAY[1.0, 2.0, 3.0]::float4[], nullable_doc_embedding,
  options=>jsonb'{"num_leaves_to_search": 10}'
)
LIMIT 100;

Adicionar a distância aos resultados da consulta

Nos exemplos anteriores, a distância do ANN é usada para ordenar os resultados da consulta, mas não é mostrada nesses resultados. Para adicionar a distância aos resultados, coloque a função de distância na cláusula SELECT e crie um alias para ela. Em seguida, use esse alias como a única chave de ordenação na cláusula ORDER BY.

Exemplos

GoogleSQL 

SELECT DocId,
  APPROX_EUCLIDEAN_DISTANCE(
    ARRAY<FLOAT32>[1.0, 2.0, 3.0], NullableDocEmbedding,
    options => JSON '{"num_leaves_to_search": 10}') AS distance
FROM Documents
WHERE NullableDocEmbedding IS NOT NULL AND WordCount > 1000
ORDER BY distance
LIMIT 100

PostgreSQL 

SELECT doc_id,
  spanner.approx_euclidean_distance(
    ARRAY[1.0, 2.0, 3.0]::float4[], nullable_doc_embedding,
    options=>jsonb'{"num_leaves_to_search": 10}'
  ) AS distance
FROM documents
WHERE nullable_doc_embedding IS NOT NULL AND word_count > 1000
ORDER BY distance
LIMIT 100;

A seguir