Realizar búsquedas de similitud de vectores en Spanner buscando los vecinos más cercanos

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

  • Distancia euclídea: mide la distancia más corta entre dos vectores.
  • Distancia del coseno: mide el coseno del ángulo entre dos vectores.
  • Producto escalar: calcula el coseno del ángulo multiplicado por el producto de las magnitudes de los vectores correspondientes. Si sabes que todas las inserciones vectoriales de tu conjunto de datos están normalizadas, puedes usar DOT_PRODUCT() como función de distancia.
  • K vecinos más cercanos (KNN): 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 vectoriales de los k vecinos más cercanos (KNN) en casos prácticos como la búsqueda de similitudes o la generación aumentada de recuperación. Spanner admite las funciones COSINE_DISTANCE(), EUCLIDEAN_DISTANCE() y DOT_PRODUCT(), que operan en inserciones de vectores, lo que te permite encontrar los vecinos más cercanos (KNN) de la inserción de entrada.

Por ejemplo, después de generar y guardar tus datos operativos de Spanner como incrustaciones de vectores, puedes proporcionar estas incrustaciones de vectores como parámetro de entrada en tu consulta para encontrar los vectores más cercanos en el espacio N-dimensional y buscar elementos similares o relacionados semánticamente.

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

Ejemplos

En los siguientes ejemplos se muestra la búsqueda de vecinos más cercanos (KNN), la búsqueda de KNN en datos particionados y el uso de un índice secundario con KNN.

En todos los ejemplos se usa EUCLIDEAN_DISTANCE(). También puedes usar COSINE_DISTANCE(). Además, si todas las inserciones vectoriales de tu conjunto de datos están normalizadas, puedes usar DOT_PRODUCT() como función de distancia.

Supongamos que tenemos una tabla Documents que tiene una columna (DocEmbedding) de inserciones de texto precalculadas a partir de la columna de DocContents bytes.

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)
);

Si la inserción de entrada de "béisbol, pero no béisbol profesional" es la matriz [0.3, 0.3, 0.7, 0.7], puedes encontrar los cinco documentos más cercanos que coincidan con la siguiente 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;

Estos son los resultados esperados de este ejemplo:

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

Ejemplo 2: Búsqueda de k vecinos más cercanos en datos con particiones

La consulta del ejemplo anterior se puede modificar añadiendo condiciones a la cláusula WHERE para limitar la búsqueda de vectores a un subconjunto de sus datos. Una aplicación habitual de esta función es buscar 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;

Estos son los resultados esperados de este ejemplo:

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

Ejemplo 3: Búsqueda de los k vecinos más cercanos en intervalos de índices secundarios

Si el filtro de cláusula WHERE que estás usando no forma parte de la clave principal 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;

Estos son los resultados esperados de este ejemplo:

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

Siguientes pasos