2024-10-23

Was ist neu in der Vector-Similarity-Search-Erweiterung?

Max Gabrielsson

Im vorherigen Blogbeitrag haben wir die DuckDB-Vector-Similarity-Search-(VSS-)Erweiterung vorgestellt. Die Erweiterung ist noch recht experimentell, aber wir finden es interessant, ins Detail einiger neuer Features und Verbesserungen zu gehen, an denen wir seit dem ersten Release gearbeitet haben.

Schnelleres Indexieren

Wie zuvor dokumentiert, ist das Anlegen eines HNSW-(Hierarchical Navigable Small Worlds-)Index über eine bereits gefüllte Tabelle viel effizienter, als zuerst den Index anzulegen und dann in die Tabelle einzufügen. Es ist viel leichter vorherzusagen, wie groß der Index wird, wenn die Gesamtzahl der Zeilen von vornherein bekannt ist – dann lässt sich die Arbeit in Chunks aufteilen, die groß genug sind, um sie über mehrere Threads zu verteilen. Im ersten Release war diese Arbeitsverteilung etwas zu grob: Wir haben nur für jede Row Group (standardmäßig etwa 120.000 Zeilen) in der Tabelle einen zusätzlichen Worker-Thread eingeplant.

Wir haben jetzt einen Extra-Buffer-Schritt in der Index-Erzeugungspipeline eingeführt, der feinere Arbeitsverteilung, klügere Speicherallokation und weniger Contention zwischen Worker-Threads ermöglicht. Das führt zu deutlich höherer CPU-Sättigung und einem signifikanten Speedup beim Bau von HNSW-Indizes in Umgebungen mit vielen Threads – unabhängig davon, wie groß oder klein die zugrunde liegende Tabelle ist.

Ein Bonus dieser Änderung: Wir können jetzt eine Fortschrittsanzeige beim Indexbau ausgeben – ein netter Touch, wenn Sie trotzdem noch eine Weile warten müssen (trotz der jetzt viel besseren Nutzung der Systemressourcen!).

Neue Distanzfunktionen

Im ersten VSS-Release haben wir drei Distanzfunktionen unterstützt: array_distance, array_cosine_similarity und array_inner_product. Allerdings ist nur array_distance wirklich eine Distanzfunktion, weil sie Ergebnisse näher an 0 liefert, wenn die Vektoren ähnlich sind, und näher an 1, wenn sie unähnlich sind – im Gegensatz zu z. B. array_cosine_similarity, das 1 liefert, wenn die Vektoren identisch sind. Hoppla!

Zur Abhilfe haben wir zwei neue Distanzen eingeführt:

array_cosine_distance, gleichwertig zu 1 - array_cosine_simililarity
array_negative_inner_product, gleichwertig zu -array_inner_product

Diese werden jetzt durch den HNSW-Index beschleunigt, sodass Query-Muster und Sortierung für alle unterstützten Metriken konsistent sind, unabhängig davon, ob Sie den HNSW-Index nutzen oder nicht. Außerdem: Haben Sie einen HNSW mit z. B. der Metrik cosine und schreiben eine Top-k-Query mit 1 - array_cosine_similarity als Ranking-Kriterium, sollte der Optimizer den Ausdruck zu array_cosine_distance normalisieren und den Index auch für diese Funktion nutzen.

Der Vollständigkeit halber haben wir die entsprechenden Distanzfunktionen für den dynamisch großen Datentyp LIST ergänzt (Präfix list_ statt array_) und den binären Operator <=> zum Alias von array_cosine_distance gemacht, passend zur Semantik der pgvector-Erweiterung für PostgreSQL.

Indexbeschleunigte „Top-K“-Aggregate

Eine weitere coole Sache in Core-DuckDB seit letztem Mal: DuckDB hat jetzt Extra-Overloads für die Aggregatfunktionen min_by und max_by (und ihre Aliase arg_min und arg_max). Diese neuen Overloads nehmen ein optionales drittes Argument n, das die Zahl der Top-k- (bzw. Top-n-)Elemente angibt, und geben sie als sortierten LIST-Wert aus. Ein Beispiel:

-- Create a table with some example data
CREATE OR REPLACE TABLE vecs AS
    SELECT
        row_number() OVER () AS id,
        [a, b, c]::FLOAT[3] AS vec
    FROM
        range(1,4) AS x(a), range(1,4) AS y(b), range(1,4) AS z(c);

-- Find the top 3 rows with the vector closest to [2, 2, 2]
SELECT
    arg_min(vecs, array_distance(vec, [2, 2, 2]::FLOAT[3]), 3)
FROM
    vecs;

[{'id': 14, 'vec': [2.0, 2.0, 2.0]}, {'id': 13, 'vec': [2.0, 1.0, 2.0]}, {'id': 11, 'vec': [1.0, 2.0, 2.0]}]

Natürlich enthält die VSS-Erweiterung jetzt Optimizer-Regeln, um den HNSW-Index für diese Top-k-Aggregate zu nutzen, wenn die Sortiereingabe eine Distanzfunktion ist, die auf eine indexierte Vektorspalte verweist – analog zum Query-Muster SELECT a FROM b ORDER BY array_distance(a.vec, query_vec) LIMIT k aus dem vorherigen Blogbeitrag. Die neuen Overloads erlauben, dieselbe Query knapper und lesbarer auszudrücken, und vermeiden trotzdem einen vollständigen Scan und Sort der zugrunde liegenden Tabelle (solange die Tabelle einen passenden HNSW-Index hat).

Indexbeschleunigte `LATERAL`-Joins

Nach Benchmarks der ersten VSS-Version haben wir merkt: Obwohl Index-Lookups auf unserem HNSW-Index wirklich schnell sind (Dank an die Bibliothek USearch!), hat die Suche nach einzelnen Vektoren mit DuckDB viel Latenz im Vergleich zu anderen Lösungen. Die Gründe sind vielfältig und nuanciert, aber klar ist: Unsere Wahl der HNSW-Implementierung, USearch, ist hier nicht der Engpass – das Profiling zeigte nur etwa 2 % der Laufzeit innerhalb von USearch.

Der Großteil des Overheads pro Query kommt daher, dass DuckDB nicht für Point Queries optimiert ist, also Queries, die im Wesentlichen eine einzelne Zeile holen und verarbeiten. Weil DuckDB auf einer vektorisierten Ausführungsengine basiert, ist die kleinste Arbeitseinheit nicht 1 Zeile, sondern 2.048, und weil wir erwarten, eine Menge Daten zu knacken, investieren wir oft viel Zeit im Voraus, um den Query-Plan zu optimieren und große Buffer und Caches vorab zu allokieren, damit alles so effizient wie möglich ist, sobald die Ausführung startet. Viel davon wird überflüssig, wenn die tatsächliche Working Set so klein ist. Lohnt es sich wirklich, jedes Element eines konstanten 768-langen Query-Vektors zu inspizieren und zu hashen, um gemeinsame Teilausdrücke zu suchen, wenn Sie wissen, dass im Ergebnis nur eine Handvoll Zeilen stehen?

Wir haben Ideen, wie wir dieses Szenario künftig verbessern, haben uns aber vorerst entschieden, nicht auf unsere Schwächen, sondern auf unsere Stärken zu setzen: eine Menge Daten knacken! Statt die Query „1:N“ zu optimieren („zu diesem einen Embedding die nächsten N Embeddings“), konzentrieren wir uns auf „N:M“: „zu all diesen N Embeddings jeweils die nächsten M Embeddings zuordnen“. Wie sähe das aus? Natürlich als LATERAL-Join!

Im Wesentlichen können wir HNSW-Indizes jetzt nutzen, um LATERAL-Joins zu beschleunigen, bei denen die „innere“ Query wie die Top-k-Queries aussieht, die wir normalerweise anvisieren, zum Beispiel:

SELECT a
FROM b
ORDER BY array_distance(a.vec, query_vec)
LIMIT k;

Aber das Array query_vec ist jetzt eine Referenz auf eine „äußere“ Join-Tabelle. Die einzige Voraussetzung: Die innere Tabelle braucht einen HNSW-Index auf der Vektorspalte, der zur Distanzfunktion passt. Ein Beispiel:

-- Set the random seed for reproducibility
SELECT setseed(0.42);

-- Create some example tables
CREATE TABLE queries AS
    SELECT
        i AS id,
        [random(), random(), random()]::FLOAT[3] AS embedding
    FROM generate_series(1, 10_000) r(i);

CREATE TABLE items AS
    SELECT
        i AS id,
        [random(), random(), random()]::FLOAT[3] AS embedding
FROM generate_series(1, 10_000) r(i);

-- Collect the 5 closest items to each query embedding
SELECT queries.id AS id, list(inner_id) AS matches
    FROM queries, LATERAL (
        SELECT
            items.id AS inner_id,
            array_distance(queries.embedding, items.embedding) AS dist
        FROM items
        ORDER BY dist
        LIMIT 5
    )
GROUP BY queries.id;

Die Ausführung auf meinem MacBook mit Apple M3 Pro und 36 GB Speicher dauert etwa 10 Sekunden.

Wenn wir diesen Query-Plan mit EXPLAIN ansehen, sehen wir viele fortgeschrittene Operatoren:

PRAGMA explain_output = 'optimized_only';
EXPLAIN ...

Vanilla-Query-Plan (Operatoren und erwartete Kardinalitäten)

┌───────────────────────────┐
│         PROJECTION        │
│    ────────────────────   │
│       ~5000000 Rows       │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│       HASH_GROUP_BY       │
│    ────────────────────   │
│       ~5000000 Rows       │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         PROJECTION        │
│    ────────────────────   │
│       ~10000000 Rows      │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         PROJECTION        │
│    ────────────────────   │
│       ~10000000 Rows      │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│      RIGHT_DELIM_JOIN     │
│    ────────────────────   │
│       ~10000000 Rows      ├──────────────┐
└─────────────┬─────────────┘              │
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│         SEQ_SCAN          ││         HASH_JOIN         │
│    ────────────────────   ││    ────────────────────   │
│        ~10000 Rows        ││       ~10000000 Rows      ├──────────────┐
└───────────────────────────┘└─────────────┬─────────────┘              │
                             ┌─────────────┴─────────────┐┌─────────────┴─────────────┐
                             │         PROJECTION        ││         DUMMY_SCAN        │
                             │    ────────────────────   ││                           │
                             │       ~10000000 Rows      ││                           │
                             └─────────────┬─────────────┘└───────────────────────────┘
                             ┌─────────────┴─────────────┐
                             │           FILTER          │
                             │    ────────────────────   │
                             │       ~10000000 Rows      │
                             └─────────────┬─────────────┘
                             ┌─────────────┴─────────────┐
                             │         PROJECTION        │
                             │    ────────────────────   │
                             │       ~50000000 Rows      │
                             └─────────────┬─────────────┘
                             ┌─────────────┴─────────────┐
                             │           WINDOW          │
                             └─────────────┬─────────────┘
                             ┌─────────────┴─────────────┐
                             │         PROJECTION        │
                             │    ────────────────────   │
                             │       ~50000000 Rows      │
                             └─────────────┬─────────────┘
                             ┌─────────────┴─────────────┐
                             │       CROSS_PRODUCT       ├──────────────┐
                             └─────────────┬─────────────┘              │
                             ┌─────────────┴─────────────┐┌─────────────┴─────────────┐
                             │         SEQ_SCAN          ││         DELIM_SCAN        │
                             │    ────────────────────   ││    ────────────────────   │
                             │        ~10000 Rows        ││         ~5000 Rows        │
                             └───────────────────────────┘└───────────────────────────┘

Der Plan sieht sehr kompliziert aus; am besorgniserregendsten ist der CROSS_PRODUCT weiter unten, der die erwartete Kardinalität aufbläht und zeigt, dass wir viel Arbeit tun, die wir wahrscheinlich nicht tun wollen. Legen wir aber einen HNSW-Index auf der Tabelle items an mit

CREATE INDEX my_hnsw_idx ON items USING HNSW(embedding);

und führen EXPLAIN erneut aus, bekommen wir diesen Plan:

Query-Plan mit HNSW-Index (Operatoren und erwartete Kardinalitäten)

┌───────────────────────────┐
│         PROJECTION        │
│    ────────────────────   │
│        ~50000 Rows        │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│       HASH_GROUP_BY       │
│    ────────────────────   │
│        ~50000 Rows        │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         PROJECTION        │
│    ────────────────────   │
│        ~50000 Rows        │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         PROJECTION        │
│    ────────────────────   │
│        ~50000 Rows        │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         PROJECTION        │
│    ────────────────────   │
│        ~50000 Rows        │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│      HNSW_INDEX_JOIN      │
│    ────────────────────   │
│        ~50000 Rows        │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         SEQ_SCAN          │
│    ────────────────────   │
│        ~10000 Rows        │
└───────────────────────────┘

Dieser Plan ist drastisch vereinfacht; am wichtigsten: Der neue Operator HNSW_INDEX_JOIN ersetzt den CROSS_PRODUCT-Knoten, und die geschätzte Kardinalität sank von 5.000.000 auf 50.000! Die Ausführung dauert jetzt etwa 0,15 Sekunden. Das ist fast ein 66-facher Speedup!

Diese Optimierung ist erst kürzlich in die VSS-Erweiterung gekommen. Wenn Sie vss bereits für DuckDB v1.1.2 installiert haben, führen Sie folgenden Befehl aus, um die aktuelle Version zu holen:

UPDATE EXTENSIONS (vss);

Fazit

Das war’s für diesmal! Wir hoffen, dieses Update zur DuckDB-Vector-Similarity-Search-Erweiterung hat Ihnen gefallen. Der Fokus lag auf neuen Features und Verbesserungen wie schnellerem Indexieren, zusätzlichen Distanzfunktionen und mehr Optimizer-Regeln; wir arbeiten weiter an den Einschränkungen aus dem vorherigen Blogbeitrag. Bald hoffen wir, mehr zu Custom Indexes und indexbasierten Optimierungen zu teilen! Fragen oder Feedback gern im duckdb-vss-GitHub-Repository oder im DuckDB Discord. Bis bald!