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

Nesta página, descrevemos como realizar a pesquisa de vetores de similaridade no Spanner usando as funções de distância do cosseno, distância euclidiana e vetor de produto escalar para encontrar vizinhos 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 das magnitudes vetoriais correspondentes. Se você sabe que todos os embeddings de vetores no conjunto de dados são normalizados, pode usar DOT_PRODUCT() como uma função de distância.
  • Vizinhos K-mais próximos (KNN, na sigla em inglês): um algoritmo de machine learning supervisionado usado para resolver problemas de classificação ou regressão.

É possível usar funções de distância vetorial para realizar a pesquisa de vetores de vizinhos mais próximos (KNN, na sigla em inglês) para casos de uso como pesquisa por similaridade ou geração aumentada de recuperação. O Spanner é compatível com as funções COSINE_DISTANCE(), EUCLIDEAN_DISTANCE() e DOT_PRODUCT(), que operam em embeddings de vetores, 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 vetores, é possível fornecer esses embeddings de vetores 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:

Examples

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(). Você também pode usar COSINE_DISTANCE(). Além disso, se todos os embeddings de vetores no conjunto de dados forem normalizados, será possível usar DOT_PRODUCT() como uma função de distância.

Considere uma tabela Documents que tenha 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,
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 "baseball, mas não beisebol profissional" seja a matriz [0.3, 0.3, 0.7, 0.7], será possível 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 à cláusula WHERE para limitar a pesquisa vetorial a um subconjunto dos seus dados. Uma aplicação comum disso é 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 de KNN em intervalos de índices secundários

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