Realiza una búsqueda de similitud de vectores en Spanner encontrando los K vecinos más cercanos

En esta página, se describe cómo realizar una búsqueda de similitud de vectores en Spanner con las funciones de vectores de distancia coseno, distancia euclidiana y producto punto para encontrar los K vecinos más cercanos. Esta información se aplica a las bases de datos con dialecto de GoogleSQL y a las bases de datos con dialecto de PostgreSQL. Antes de leer esta página, es importante que comprendas los siguientes conceptos:

  • Distancia euclidiana: Mide la distancia más corta entre dos vectores.
  • Distancia de coseno: Mide el coseno del ángulo entre dos vectores.
  • Producto escalar: Calcula el coseno del ángulo multiplicado por el producto de las magnitudes del vector correspondiente. Si sabes que todas las embedding de vector de tu conjunto de datos están normalizadas, puedes usar DOT_PRODUCT() como función de distancia.
  • K-vecinos más cercanos (KNN): Es un algoritmo de aprendizaje automático supervisado que se usa para resolver problemas de clasificación o regresión.

Puedes usar funciones de distancia vectorial para realizar búsquedas de vectores de K-vecinos más cercanos (KNN) para casos de uso como la búsqueda de similitud o la generación aumentada por recuperación. Spanner admite las funciones COSINE_DISTANCE(), EUCLIDEAN_DISTANCE() y DOT_PRODUCT(), que operan en embeddings de vectores, lo que te permite encontrar los KNN del embedding de entrada.

Por ejemplo, después de generar y guardar tus datos operacionales de Spanner como embeddings de vectores, puedes proporcionar estos embeddings de vectores como un parámetro de entrada en tu consulta para encontrar los vectores más cercanos en el espacio N-dimensional y buscar elementos relacionados o similares desde el punto de vista semántico.

Las tres funciones de distancia toman los argumentos vector1 y vector2, que son del tipo array<>, y deben constar de las mismas dimensiones y tener la misma longitud. Para obtener más detalles sobre estas funciones, consulta los siguientes recursos:

Ejemplos

En los siguientes ejemplos, se muestran la búsqueda de KNN, la búsqueda de KNN en datos particionados y el uso de un índice secundario con KNN.

Todos los ejemplos usan EUCLIDEAN_DISTANCE(). También puedes usar COSINE_DISTANCE(). Además, si todas las embeddings de vectores de tu conjunto de datos están normalizadas, puedes usar DOT_PRODUCT() como función de distancia.

Considera una tabla Documents que tiene una columna (DocEmbedding) de incorporaciones de texto precalculadas de la columna 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)
);

Suponiendo que una incorporación de entrada para "béisbol, pero no béisbol profesional" es el array [0.3, 0.3, 0.7, 0.7], puedes encontrar los cinco documentos más cercanos que coinciden con la siguiente búsqueda:

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;

Los resultados esperados de este ejemplo son los siguientes:

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

Ejemplo 2: Búsqueda de KNN en datos particionados

La consulta del ejemplo anterior se puede modificar agregando condiciones a la cláusula WHERE para limitar la búsqueda de vectores a un subconjunto de tus datos. Una aplicación común de esto es la búsqueda en datos particionados, como las filas que pertenecen a un 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;

Los resultados esperados de este ejemplo son los siguientes:

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

Ejemplo 3: Búsqueda de KNN en rangos de índices secundarios

Si el filtro de la cláusula WHERE que usas no forma parte de la clave primaria de la tabla, puedes crear un índice secundario para acelerar la operación con un análisis solo 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;

Los resultados esperados de este ejemplo son los siguientes:

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, ...]        |
+------------+-----------------+-----------------+

¿Qué sigue?