Realizar pesquisa de similaridade vetorial no Spanner encontrando os vizinhos mais próximos K

Esta página descreve como realizar a pesquisa de similaridade de vetor no Spanner usando a distância de cosseno, a distância euclidiana e as funções de vetor de produto vetorial para encontrar os vizinhos mais próximos. Antes de ler esta página, é importante entender os seguintes conceitos:

  • Distância euclidiana: mede a menor distância entre dois vetores.
  • Distância de cosseno: mede o cosseno do ângulo entre dois vetores.
  • Produto escalar: calcula o cosseno do ângulo multiplicado pelo produto das magnitudes vetoriais correspondentes. Se você souber que todos os embeddings de vetor no seu conjunto de dados estão normalizados, use DOT_PRODUCT() como uma função de distância.
  • Vizinhos k-mais próximos (KNN): um algoritmo de aprendizado de máquina supervisionado usado para resolver problemas de classificação ou regressão.

É possível usar funções de distância de vetor para realizar a pesquisa de vetor de vizinho mais próximo (KNN) em casos de uso como pesquisa de similaridade ou geração aumentada de recuperação. O Spanner oferece suporte às funções COSINE_DISTANCE(), EUCLIDEAN_DISTANCE() e DOT_PRODUCT(), que operam em embeddings de vetor, permitindo que você encontre a KNN do embedding de entrada.

Por exemplo, depois de gerar e salvar seus dados operacionais do Spanner como embeddings vetoriais, você pode fornecer esses embeddings vetoriais como um parâmetro de entrada na consulta para encontrar os vetores mais próximos no espaço N-dimensional e pesquisar itens semanticamente semelhantes ou relacionados.

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 o mesmo comprimento. Para mais detalhes sobre essas funções, consulte:

Exemplos

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

Todos os exemplos usam EUCLIDEAN_DISTANCE(). Também é possível usar COSINE_DISTANCE(). Além disso, se todos os embeddings de vetor no seu conjunto de dados forem normalizados, você poderá usar DOT_PRODUCT() como uma função de distância.

Considere uma tabela Documents que tenha uma coluna (DocEmbedding) de representações de texto pré-calculadas da coluna de bytes DocContents.

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 NOT NULL,
DocId        bigint NOT NULL,
Author       varchar(1024),
DocContents  bytea,
DocEmbedding float4[],
PRIMARY KEY  (UserId, DocId)
);

Supondo que uma embedding de entrada para "baseball, but not professional baseball" seja a matriz [0.3, 0.3, 0.7, 0.7], você pode encontrar os cinco documentos mais próximos que correspondem, 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 KNN em dados particionados

A consulta no exemplo anterior pode ser modificada adicionando condições à cláusula WHERE para limitar a pesquisa de vetores a um subconjunto dos seus dados. Uma aplicação comum é pesquisar dados particionados, como 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 KNN em intervalos de índice secundário

Se o filtro da cláusula WHERE que você está usando não fizer parte da chave primária da tabela, crie um índice secundário para acelerar a operação com uma verificação somente de índice.

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