Realizar a pesquisa de vetores de similaridade no Spanner encontrando os vizinhos mais próximos

Esta página descreve como realizar uma pesquisa de vetor de similaridade em Spanner usando a distância do cosseno, a distância euclidiana e o ponto de vetores de produto para encontrar os vizinhos K mais próximos. Antes de ler esta página, é importante que você entenda os seguintes conceitos:

  • Distância euclidiana: mede a menor distância entre dois vetores.
  • Distância do cosseno: mede o cosseno do ângulo entre dois vetores.
  • Produto escalar: calcula o cosseno do ângulo multiplicado pelo produto do vetor correspondente grandes. Se você sabe que todos os embeddings de vetores no conjunto de dados são normalizado, é possível usar DOT_PRODUCT() como uma função de distância.
  • Vizinhos mais próximos (KNN, na sigla em inglês): um algoritmo de machine learning supervisionado usado para resolver problemas de problemas de regressão.

É possível usar funções de distância vetorial para executar vizinhos mais próximos (KNN, na sigla em inglês) pesquisa vetorial para casos de uso como pesquisa por similaridade ou aumento da recuperação geração de imagens. O Spanner dá suporte aos COSINE_DISTANCE(), as funções EUCLIDEAN_DISTANCE() e DOT_PRODUCT(), que operam em embeddings, permitindo que você encontre o KNN do embedding de entrada.

Por exemplo, depois de gerar e salvar o Spanner operacional dados como embeddings vetoriais, é possível forneça esses embeddings de vetor como um parâmetro de entrada em sua consulta para encontrar o vetores mais próximos em um espaço N-dimensional para procurar semanticamente semelhantes itens relacionados.

Todas as três funções de distância usam os argumentos vector1 e vector2, que são do tipo array<> e precisam ter as mesmas dimensões e ter o mesmo tamanho. Para mais detalhes sobre essas funções, consulte:

Exemplos

Os exemplos a seguir mostram a pesquisa KNN, a pesquisa KNN em dados particionados e usando um índice secundário com KNN.

Todos os exemplos usam EUCLIDEAN_DISTANCE(). Você também pode usar COSINE_DISTANCE(): Além disso, se todos os embeddings de vetores no conjunto de dados são normalizados, é possível usar DOT_PRODUCT() como uma função de distância.

Considere uma tabela Documents que tem uma coluna (DocEmbedding) de variáveis pré-computadas embeddings de texto da coluna DocContents bytes.

GoogleSQL 

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

PostgreSQL 

CREATE TABLE Documents (
UserId       bigint primary key,
DocId        bigint primary key,
Author       varchar(1024),
DocContents  bytea,
DocEmbedding float4[]
);

Supondo que um embedding de entrada para "beisebol, mas não beisebol profissional" for a matriz [0.3, 0.3, 0.7, 0.7], você poderá encontrar os cinco principais documentos mais próximos correspondentes, com a seguinte consulta:

GoogleSQL 

SELECT DocId, DocEmbedding FROM Documents
ORDER BY EUCLIDEAN_DISTANCE(DocEmbedding,
ARRAY<FLOAT32>[0.3, 0.3, 0.7, 0.8])
LIMIT 5;

PostgreSQL 

SELECT DocId, DocEmbedding FROM Documents
ORDER BY spanner.euclidean_distance(DocEmbedding,
'{0.3, 0.3, 0.7, 0.8}'::float4[])
LIMIT 5;

Os resultados esperados deste exemplo:

Documents
+---------------------------+-----------------+
| DocId                     | DocEmbedding    |
+---------------------------+-----------------+
| 24                        | [8, ...]        |
+---------------------------+-----------------+
| 25                        | [6, ...]        |
+---------------------------+-----------------+
| 26                        | [3.2, ...]      |
+---------------------------+-----------------+
| 27                        | [38, ...]       |
+---------------------------+-----------------+
| 14229                     | [1.6, ...]      |
+---------------------------+-----------------+

Exemplo 2: pesquisa de KNN em dados particionados

A consulta no exemplo anterior pode ser modificada adicionando condições ao cláusula WHERE para limitar a pesquisa vetorial a um subconjunto dos seus dados. Um erro comum aplicação é pesquisar dados particionados, como as linhas que pertencem a um UserId específico.

GoogleSQL 

SELECT UserId, DocId, DocEmbedding FROM Documents
WHERE UserId=18
ORDER BY EUCLIDEAN_DISTANCE(DocEmbedding,
ARRAY<FLOAT32>[0.3, 0.3, 0.7, 0.8])
LIMIT 5;

PostgreSQL 

SELECT UserId, DocId, DocEmbedding FROM Documents
WHERE UserId=18
ORDER BY spanner.euclidean_distance(DocEmbedding,
'{0.3, 0.3, 0.7, 0.8}'::float4[])
LIMIT 5;

Os resultados esperados deste exemplo:

Documents
+-----------+-----------------+-----------------+
| UserId    | DocId           | DocEmbedding    |
+-----------+-----------------+-----------------+
| 18        | 234             | [12, ...]       |
+-----------+-----------------+-----------------+
| 18        | 12              | [1.6, ...]      |
+-----------+-----------------+-----------------+
| 18        | 321             | [22, ...]       |
+-----------+-----------------+-----------------+
| 18        | 432             | [3, ...]        |
+-----------+-----------------+-----------------+

Exemplo 3: pesquisa de KNN em intervalos de índices secundários

Se o filtro de cláusula WHERE que você está usando não fizer parte da chave primária da tabela, você pode criar um índice secundário para acelerar a operação com um index-only scan.

GoogleSQL 

CREATE INDEX DocsByAuthor
ON Documents(Author)
STORING (DocEmbedding);

SELECT Author, DocId, DocEmbedding FROM Documents
WHERE Author="Mark Twain"
ORDER BY EUCLIDEAN_DISTANCE(DocEmbedding,
   <embeddings for "book about the time traveling American">)
LIMIT 5;

PostgreSQL 

CREATE INDEX DocsByAuthor
ON Documents(Author)
INCLUDE (DocEmbedding);

SELECT Author, DocId, DocEmbedding FROM Documents
WHERE Author="Mark Twain"
ORDER BY spanner.euclidean_distance(DocEmbedding,
   <embeddings for "that book about the time traveling American">)
LIMIT 5;

Os resultados esperados deste exemplo:

Documents
+------------+-----------------+-----------------+
| Author     | DocId           | DocEmbedding    |
+------------+-----------------+-----------------+
| Mark Twain | 234             | [12, ...]       |
+------------+-----------------+-----------------+
| Mark Twain | 12              | [1.6, ...]      |
+------------+-----------------+-----------------+
| Mark Twain | 321             | [22, ...]       |
+------------+-----------------+-----------------+
| Mark Twain | 432             | [3, ...]        |
+------------+-----------------+-----------------+
| Mark Twain | 375             | [9, ...]        |
+------------+-----------------+-----------------+

A seguir