Procedimiento para optimizar el rendimiento al usar pgvector en Azure Cosmos DB for PostgreSQL

Artículo
08/15/2024

SE APLICA A: Azure Cosmos DB for PostgreSQL (con tecnología de la extensión de base de datos de Citus en PostgreSQL)

La extensión pgvector agrega una búsqueda de similitud vectorial de código abierto a PostgreSQL.

En este artículo se exploran las limitaciones y desventajas de pgvector y se muestra cómo usar la creación de particiones, la indexación y la configuración de búsqueda para mejorar el rendimiento.

Para más información sobre la propia extensión, vea Conceptos básicos de pgvector. También puede consultar el archivo LÉAME oficial del proyecto.

Rendimiento

Siempre debe empezar con la investigación del plan de consulta. Si la consulta finaliza razonablemente rápido, ejecute EXPLAIN (ANALYZE,VERBOSE, BUFFERS).

EXPLAIN (ANALYZE, VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

Para las consultas que tardan demasiado tiempo en ejecutarse, considere la posibilidad de quitar la palabra clave ANALYZE. El resultado contiene menos detalles, pero se proporciona al instante.

EXPLAIN (VERBOSE, BUFFERS) SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

Los sitios de terceros, como explain.depesz.com, pueden ser útiles para comprender los planes de consulta. Algunas preguntas que debe intentar responder son las siguientes:

¿Se ha paralelizado la consulta?
¿Se ha usado un índice?
¿He usado la misma condición en la cláusula WHERE que en una definición de índice parcial?
Si uso la creación de particiones, ¿se eliminan las particiones que no son necesarias?

Si los vectores se normalizan hasta la longitud 1, como las inserciones de OpenAI. Considere la posibilidad de usar el producto interno (<#>) para obtener el mejor rendimiento.

Ejecución en paralelo

En la salida del plan de explicación, busque Workers Planned y Workers Launched (esto último solo cuando se ha usado la palabra clave ANALYZE). El parámetro max_parallel_workers_per_gather de PostgreSQL define cuántos trabajos en segundo plano puede iniciar la base de datos para cada nodo de plan Gather y Gather Merge. El aumento de este valor podría acelerar las consultas de búsqueda exactas sin tener que crear índices. Pero tenga en cuenta que es posible que la base de datos no decida ejecutar el plan en paralelo incluso cuando este valor sea alto.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 3;

                                        QUERY PLAN
------------------------------------------------------------------------------------------
 Limit  (cost=4214.82..4215.16 rows=3 width=33)
   ->  Gather Merge  (cost=4214.82..13961.30 rows=84752 width=33)
         Workers Planned: 1
         ->  Sort  (cost=3214.81..3426.69 rows=84752 width=33)
               Sort Key: ((embedding <-> '[1,2,3]'::vector))
               ->  Parallel Seq Scan on t_test  (cost=0.00..2119.40 rows=84752 width=33)
(6 rows)

Indización

Sin índices presentes, la extensión realiza una búsqueda exacta, lo que proporciona una coincidencia perfecta a costa del rendimiento.

Para realizar la búsqueda de vecinos más próximos aproximada, puede crear índices en los datos, lo que cambia la coincidencia por el rendimiento de la ejecución.

Siempre que sea posible, cargue los datos antes de indexarlos. Es más rápido crear el índice de esta manera y el diseño resultante es más óptimo.

Hay dos tipos de índice admitidos:

Archivo invertido con compresión plana (IVVFlat)
Mundo pequeño navegable jerárquico (HNSW)

El índice IVFFlat tiene tiempos de compilación más rápidos y usa menos memoria que HNSW, pero tiene un rendimiento de consulta menor (en términos de equilibrio de velocidad y recuperación).

Límites

Para indexar una columna, debe tener dimensiones definidas. Al intentar indexar una columna definida como col vector resultado, se produce el error: ERROR: column does not have dimensions.
Solo puede indexar una columna que tenga hasta 2000 dimensiones. Al intentar indexar una columna con más dimensiones se produce el error: ERROR: column cannot have more than 2000 dimensions for INDEX_TYPE index donde INDEX_TYPE es ivfflat o hnsw.

Aunque puede almacenar vectores con más de 2000 dimensiones, no se pueden indexar. Puede usar la reducción de dimensionalidad para ajustarse a los límites. Como alternativa, puede recurrir a la creación de particiones con Azure Cosmos DB for PostgreSQL para lograr un rendimiento aceptable sin indexación.

Archivo invertido con compresión plana (IVVFlat)

ivfflat es un índice para la búsqueda de vecino más próximo (ANN) aproximada. Este método usa un índice de archivo invertido para dividir el conjunto de datos en varias listas. El parámetro probes controla cuántas listas se buscan, lo que puede mejorar la precisión de los resultados de la búsqueda a costa de una velocidad de búsqueda más lenta.

Si el parámetro probes se establece en el número de listas del índice, se buscan todas las listas y la búsqueda se convierte en una búsqueda de vecino más próximo exacta. En este caso, el planificador no usa el índice porque la búsqueda de todas las listas equivale a realizar una búsqueda por fuerza bruta en todo el conjunto de datos.

El método de indexación divide el conjunto de datos en varias listas mediante el algoritmo de agrupación en clústeres k-means. Cada lista contiene vectores más cercanos a un centro de clúster determinado. Durante una búsqueda, el vector de consulta se compara con los centros de clúster para determinar qué listas son más probables que contengan los vecinos más próximo. Si el parámetro probes está establecido en 1, solo se buscará en la lista correspondiente al centro de clúster más cercano.

Opciones de índice

La selección del valor correcto para el número de sondeos que se van a realizar y los tamaños de las listas podrían afectar al rendimiento de la búsqueda. Un buen punto de partida es el siguiente:

Usar lists igual a rows / 1000 para las tablas con hasta un millón de filas y sqrt(rows) para conjuntos de datos de mayor tamaño.
Para probes, empezar con lists / 10 para las tablas con hasta un millón de filas y sqrt(lists) para conjuntos de datos de mayor tamaño.

La cantidad de lists se define al crear el índice con la opción lists:

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 5000);

Los sondeos se pueden establecer para toda la conexión o para cada transacción (mediante SET LOCAL dentro de un bloque de transacción):

SET ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes

BEGIN;

SET LOCAL ivfflat.probes = 10;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 10 probes

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, one probe

Progreso de indexación

Con PostgreSQL 12 y versiones posteriores, puede usar pg_stat_progress_create_index para comprobar el progreso de la indexación.

SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

Las fases para la creación de índices IVFFlat son:

initializing
performing k-means
assigning tuples
loading tuples

Nota:

El porcentaje de progreso (%) solo se rellena durante la fase loading tuples.

Mundos pequeños jerárquicos (HNSW)

El hnsw es un índice para la búsqueda de vecinos más próximos (ANN) aproximadas mediante el algoritmo de mundos pequeños jerárquicos. Funciona creando un gráfico alrededor de los puntos de entrada seleccionados aleatoriamente que encuentran sus vecinos más próximos, el gráfico se extiende a continuación con varias capas, donde cada capa inferior contiene más puntos. Cuando se busca en este grafo multicapa, se empieza por arriba y se va reduciendo hasta llegar a la capa más baja, que contiene los vecinos más próximos de la consulta.

La creación de este índice toma más tiempo y memoria que IVFFlat, pero tiene un mejor equilibrio de velocidad y recuperación. Además, no hay ningún paso de entrenamiento como con IVFFlat, por lo que el índice se puede crear en una tabla vacía.

Opciones de índice

Al crear el índice, puede ajustar dos parámetros:

m: número máximo de conexiones por capa (el valor predeterminado es 16)
ef_construction: tamaño de la lista de candidatos dinámicos que se usa para la construcción del grafo (el valor predeterminado es 64)

CREATE INDEX t_test_hnsw_l2_idx ON t_test USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);

Durante las consultas, puede especificar la lista de candidatos dinámicos para la búsqueda (el valor predeterminado es 40).

La lista de candidatos dinámicos para la búsqueda se puede establecer para toda la conexión o por transacción (mediante SET LOCAL dentro de un bloque de transacciones):

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates

BEGIN;

SET hnsw.ef_search = 100;
SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses 100 candidates

COMMIT;

SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5; -- uses default, 40 candidates

Progreso de indexación

Con PostgreSQL 12 y versiones posteriores, puede usar pg_stat_progress_create_index para comprobar el progreso de la indexación.

SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;

Las fases para crear índices HNSW son:

initializing
loading tuples

Selección de la función de acceso al índice

El tipo vector permite realizar tres tipos de búsquedas en los vectores almacenados. Debe seleccionar la función de acceso correcta para el índice para que la base de datos tenga en cuenta el índice al ejecutar las consultas. Los ejemplos se muestran en los tipos de índice ivfflat, pero se puede hacer lo mismo para los índices hnsw. La opción lists solo se aplica a los índices ivfflat.

Distancia de coseno

Para la búsqueda de similitud de coseno, use el mecanismo de acceso vector_cosine_ops.

CREATE INDEX t_test_embedding_cosine_idx ON t_test USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

Para usar el índice anterior, la consulta debe realizar una búsqueda de similitud coseno, que se realiza con el operador <=>.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <=> '[1,2,3]' LIMIT 5;

                                              QUERY PLAN
------------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_cosine_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <=> '[1,2,3]'::vector)
(3 rows)

Distancia L2

Para la distancia L2 (también conocida como distancia euclidiana), use el mecanismo de acceso vector_l2_ops.

CREATE INDEX t_test_embedding_l2_idx ON t_test USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);

Para usar el índice anterior, la consulta debe realizar una búsqueda de distancia L2, que se realiza con el operador <->.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_l2_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <-> '[1,2,3]'::vector)
(3 rows)

Producto interno

Para la similitud de productos interna, use el mecanismo de acceso vector_ip_ops.

CREATE INDEX t_test_embedding_ip_idx ON t_test USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);

Para usar el índice anterior, la consulta debe realizar una búsqueda de similitud de productos interna, que se realiza con el operador <#>.

EXPLAIN SELECT * FROM t_test ORDER BY embedding <#> '[1,2,3]' LIMIT 5;

                                            QUERY PLAN
--------------------------------------------------------------------------------------------------
 Limit  (cost=5.02..5.23 rows=5 width=33)
   ->  Index Scan using t_test_embedding_ip_idx on t_test  (cost=5.02..175.06 rows=4003 width=33)
         Order By: (embedding <#> '[1,2,3]'::vector)
(3 rows)

Índices parciales

En algunos escenarios, es beneficioso tener un índice que abarque solo un conjunto parcial de los datos. Por ejemplo, se podría crear un índice solo para los usuarios premium:

CREATE INDEX t_premium ON t_test USING ivfflat (vec vector_ip_ops) WITH (lists = 100) WHERE tier = 'premium';

Ahora se puede ver que el nivel premium usa el índice:

explain select * from t_test where tier = 'premium' order by vec <#> '[2,2,2]';

                                     QUERY PLAN
------------------------------------------------------------------------------------
 Index Scan using t_premium on t_test  (cost=65.57..25638.05 rows=245478 width=39)
   Order By: (vec <#> '[2,2,2]'::vector)
(2 rows)

Aunque los usuarios del nivel gratis carecen de las ventajas:

explain select * from t_test where tier = 'free' order by vec <#> '[2,2,2]';

                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44019.01..44631.37 rows=244941 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15395.25 rows=244941 width=39)
         Filter: (tier = 'free'::text)
(4 rows)

Tener solo un subconjunto de datos indexado significa que el índice ocupa menos espacio en el disco y es más rápido al realizar búsquedas.

Es posible que PostgreSQL no reconozca que el índice se puede utilizar con seguridad si el formulario usado en la cláusula WHERE de la definición de índice parcial no coincide con el usado en las consultas. En nuestro conjunto de datos de ejemplo, solo tenemos los valores exactos 'free', 'test' y 'premium' como los valores distintos de la columna de nivel. Incluso con una consulta en la que se usa tier LIKE 'premium', PostgreSQL no usa el índice.

explain select * from t_test where tier like 'premium' order by vec <#> '[2,2,2]';

                              QUERY PLAN
-----------------------------------------------------------------------
 Sort  (cost=44086.30..44700.00 rows=245478 width=39)
   Sort Key: ((vec <#> '[2,2,2]'::vector))
   ->  Seq Scan on t_test  (cost=0.00..15396.59 rows=245478 width=39)
         Filter: (tier ~~ 'premium'::text)
(4 rows)

Creación de particiones

Una manera de mejorar el rendimiento consiste en dividir el conjunto de datos en varias particiones. Imagine un sistema en el que es natural hacer referencia a datos del año actual o posiblemente de hace dos años. En este tipo de sistema, puede crear particiones de los datos por un intervalo de fechas y, después, capitalizar el rendimiento mejorado cuando el sistema pueda leer solo las particiones pertinentes según lo definido por el año consultado.

Ahora se definirá una tabla con particiones:

CREATE TABLE t_test_partitioned(vec vector(3), vec_date date default now()) partition by range (vec_date);

Se pueden crear particiones para cada año manualmente, o bien usar la función de utilidad Citus (disponible en Cosmos DB for PostgreSQL).

    select create_time_partitions(
      table_name         := 't_test_partitioned',
      partition_interval := '1 year',
      start_from         := '2020-01-01'::timestamptz,
      end_at             := '2024-01-01'::timestamptz
    );

Compruebe las particiones creadas:

\d+ t_test_partitioned

                                Partitioned table "public.t_test_partitioned"
  Column  |   Type    | Collation | Nullable | Default | Storage  | Compression | Stats target | Description
----------+-----------+-----------+----------+---------+----------+-------------+--------------+-------------
 vec      | vector(3) |           |          |         | extended |             |              |
 vec_date | date      |           |          | now()   | plain    |             |              |
Partition key: RANGE (vec_date)
Partitions: t_test_partitioned_p2020 FOR VALUES FROM ('2020-01-01') TO ('2021-01-01'),
            t_test_partitioned_p2021 FOR VALUES FROM ('2021-01-01') TO ('2022-01-01'),
            t_test_partitioned_p2022 FOR VALUES FROM ('2022-01-01') TO ('2023-01-01'),
            t_test_partitioned_p2023 FOR VALUES FROM ('2023-01-01') TO ('2024-01-01')

Para crear manualmente una partición:

CREATE TABLE t_test_partitioned_p2019 PARTITION OF t_test_partitioned FOR VALUES FROM ('2019-01-01') TO ('2020-01-01');

Después, asegúrese de que las consultas se filtran realmente a un subconjunto de las particiones disponibles. Por ejemplo, en la consulta siguiente hemos filtrado por dos particiones:

explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01';

                                                                  QUERY PLAN
-----------------------------------------------------------------------------------------------------------------------------------------------
 Append  (cost=0.00..58.16 rows=12 width=36) (actual time=0.014..0.018 rows=3 loops=1)
   ->  Seq Scan on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.00..29.05 rows=6 width=36) (actual time=0.013..0.014 rows=1 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
   ->  Seq Scan on t_test_partitioned_p2023 t_test_partitioned_2  (cost=0.00..29.05 rows=6 width=36) (actual time=0.002..0.003 rows=2 loops=1)
         Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.125 ms
 Execution Time: 0.036 ms

Puede indexar una tabla con particiones.

CREATE INDEX ON t_test_partitioned USING ivfflat (vec vector_cosine_ops) WITH (lists = 100);

explain analyze select * from t_test_partitioned where vec_date between '2022-01-01' and '2024-01-01' order by vec <=> '[1,2,3]' limit 5;

                                                                                         QUERY PLAN
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Limit  (cost=4.13..12.20 rows=2 width=44) (actual time=0.040..0.042 rows=1 loops=1)
   ->  Merge Append  (cost=4.13..12.20 rows=2 width=44) (actual time=0.039..0.040 rows=1 loops=1)
         Sort Key: ((t_test_partitioned.vec <=> '[1,2,3]'::vector))
         ->  Index Scan using t_test_partitioned_p2022_vec_idx on t_test_partitioned_p2022 t_test_partitioned_1  (cost=0.04..4.06 rows=1 width=44) (actual time=0.022..0.023 rows=0 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
         ->  Index Scan using t_test_partitioned_p2023_vec_idx on t_test_partitioned_p2023 t_test_partitioned_2  (cost=4.08..8.11 rows=1 width=44) (actual time=0.015..0.016 rows=1 loops=1)
               Order By: (vec <=> '[1,2,3]'::vector)
               Filter: ((vec_date >= '2022-01-01'::date) AND (vec_date <= '2024-01-01'::date))
 Planning Time: 0.167 ms
 Execution Time: 0.139 ms
(11 rows)

Conclusión

Enhorabuena, acaba de aprender los inconvenientes, las limitaciones y los procedimientos recomendados para lograr el mejor rendimiento con pgvector.

Compartir vía

Procedimiento para optimizar el rendimiento al usar pgvector en Azure Cosmos DB for PostgreSQL

Rendimiento

Ejecución en paralelo

Indización

Límites

Archivo invertido con compresión plana (IVVFlat)

Opciones de índice

Progreso de indexación

Mundos pequeños jerárquicos (HNSW)

Opciones de índice

Progreso de indexación

Selección de la función de acceso al índice

Distancia de coseno

Distancia L2

Producto interno

Índices parciales

Creación de particiones

Conclusión

Comentarios

Recursos adicionales