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

Esta página descreve como realizar uma pesquisa de similaridade de vetor no Spanner usando as funções de vetor de distância de cosseno, distância euclidiana e produto escalar para encontrar os K vizinhos mais próximos. Essas informações se aplicam a bancos de dados com dialeto do GoogleSQL e do PostgreSQL. Antes de ler esta página, é importante entender os seguintes conceitos:

  • Distância euclidiana: mede a distância mais curta 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 de vetor correspondentes. Se você souber que todos os embeddings de vetor no conjunto de dados estão normalizados, use DOT_PRODUCT() como uma função de distância.
  • K vizinhos 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 K vizinhos mais próximos (KNN) para 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 o KNN do embedding de entrada.

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

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

Exemplos

Os exemplos a seguir mostram a pesquisa de KNN, a pesquisa de 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 conjunto de dados estiverem normalizados, use DOT_PRODUCT() como uma função de distância.

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

GoogleSQL 

CREATE TABLE Documents (
UserId       INT64 NOT NULL,
DocId        INT64 NOT NULL,
Author       STRING(1024),
DocContents  BYTES(MAX),
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 um 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 à 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;

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 pela adição de condições à cláusula WHERE para limitar a pesquisa de vetor a um subconjunto dos dados. Uma aplicação comum disso é pesquisar dados particionados, como linhas pertencentes 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;

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 índice secundário

Se o filtro de 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;

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