2024-05-03
Vektorsuche in DuckDB
Max Gabrielsson
In DuckDB v0.10.0 haben wir den ARRAY-Datentyp eingeführt, der Listen fester Länge speichert und den bestehenden variablen LIST-Datentyp ergänzt.
Die ursprüngliche Motivation für diesen Datentyp war, optimierte Operationen für Listen zu bieten, die die Positionssemantik ihrer Kindelemente nutzen und Verzweigungen vermeiden können, weil alle Listen dieselbe Länge haben. Denken Sie z. B. an die Art von Array-Manipulationen, die Sie in NumPy machen: Stapeln, Verschieben, Multiplizieren – Sie nennen es. Außerdem wollten wir die Interoperabilität mit Apache Arrow verbessern, weil Arrows Listen fester Länge bisher beim Einlesen in DuckDB in gewöhnliche Listen variabler Länge umgewandelt wurden und dabei Typinformationen verloren gingen.
Während der Hype um Vektorembeddings und semantische Ähnlichkeitssuche wuchs, haben wir aber auch ein paar Distanzmetrik-Funktionen für diesen neuen ARRAY-Typ eingeschmuggelt:
array_distance,
array_negative_inner_product und
array_cosine_distance
Falls Sie zu den heutigen glücklichen 10.000 gehören und noch nichts von Word Embeddings oder Vektorsuche gehört haben: Kurz gesagt ist das eine Technik, um Dokumente, Bilder, Entitäten – Daten als hochdimensionale Vektoren darzustellen und dann in einem Vektorraum nach ähnlichen Vektoren zu suchen, mit einem mathematischen „Distanz“-Ausdruck als Ähnlichkeitsmaß. Das wird in vielen Anwendungen genutzt, von der natürlichen Sprachverarbeitung über Empfehlungssysteme bis zur Bilderkennung, und ist durch generative KI und vortrainierte Modelle zuletzt stark in Mode gekommen.
Das hat die Community wirklich begeistert! Während wir (DuckLabs) zunächst öffentlich gesagt haben, dass wir keinen Index für Vektorsuche in DuckDB einbauen würden, weil uns das zu weit außerhalb des Scopes schien, waren wir sehr daran interessiert, benutzerdefinierte Indizes über Extensions generell zu unterstützen. Ich persönlich nörgele seit dem Start der Spatial-Extension von DuckDB darüber, einen „R-Tree“-Index einstöpseln zu wollen! Als eines unserer Kundenprojekte sich zu einem Proof-of-Concept für eine eigene „HNSW“-Index-Extension entwickelte, haben wir gesagt, wir versuchen es. Und … nun ja, eins kam zum anderen.
Schnellvorlauf bis jetzt: Wir freuen uns, die Verfügbarkeit der vss-Extension für Vektorsuche in DuckDB anzukündigen! Manche sagen vielleicht, wir kommen spät zur Vektorsuche-Party, wir finden, die Party fängt gerade erst an!
Also, was steckt in vss?
Die Vector-Similarity-Search-(VSS-)Extension
Oberflächlich wirkt vss wie eine vergleichsweise kleine DuckDB-Extension. Sie liefert keine neuen Datentypen, Skalarfunktionen oder Copy-Funktionen, sondern einen einzigen neuen Indextyp: HNSW (Hierarchical Navigable Small Worlds), eine graphbasierte Indexstruktur, die besonders gut für hochdimensionale Vektorsuche geeignet ist.
-- Create a table with an array columnCREATE TABLE embeddings (vec FLOAT[3]);
-- Create an HNSW index on the columnCREATE INDEX idx ON embeddings USING HNSW (vec);Dieser Indextyp kann nicht wie der eingebaute ART-Index Constraints oder Eindeutigkeit erzwingen und auch nicht Joins beschleunigen oder reguläre Spalten indizieren. Der HNSW-Index gilt nur für Spalten vom Typ ARRAY mit FLOAT-Elementen und wird nur genutzt, um Abfragen zu beschleunigen, die die „Distanz“ zwischen einem konstanten FLOAT-ARRAY und den FLOAT-ARRAYs in der indizierten Spalte berechnen, nach dieser Distanz sortieren und die Top-n-Ergebnisse zurückgeben. Also Abfragen der Form:
SELECT *FROM embeddingsORDER BY array_distance(vec, [1, 2, 3]::FLOAT[3])LIMIT 3;werden im logischen Plan zu einer Projektion über einen neuen HNSW-Index-Scan-Operator optimiert, der Limit und Sort komplett entfernt. Das können wir am EXPLAIN-Output prüfen:
EXPLAINSELECT *FROM embeddingsORDER BY array_distance(vec, [1, 2, 3]::FLOAT[3])LIMIT 3;┌───────────────────────────┐│ PROJECTION ││ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ││ #0 │└─────────────┬─────────────┘┌─────────────┴─────────────┐│ PROJECTION ││ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ││ vec ││array_distance(vec, [1.0, 2││ .0, 3.0]) │└─────────────┬─────────────┘┌─────────────┴─────────────┐│ HNSW_INDEX_SCAN ││ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ││ t1 (HNSW INDEX SCAN : ││ idx) ││ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ││ vec ││ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ││ EC: 3 │└───────────────────────────┘Beim Erzeugen des HNSW-Index können Sie einen Parameter metric übergeben, um die Distanzmetrik festzulegen. Unterstützt werden l2sq, cosine und inner_product, passend zu den drei eingebauten Distanzfunktionen: array_distance, array_cosine_distance und array_negative_inner_product.
Standard ist l2sq, also euklidische Distanz (array_distance):
CREATE INDEX l2sq_idx ON embeddings USING HNSW (vec)WITH (metric = 'l2sq');Für Kosinusdistanz (array_cosine_distance):
CREATE INDEX cos_idx ON embeddings USING HNSW (vec)WITH (metric = 'cosine');Für Inner Product (array_negative_inner_product):
CREATE INDEX ip_idx ON embeddings USING HNSW (vec)WITH (metric = 'ip');Implementierung
Die vss-Extension basiert auf der Bibliothek usearch, die eine flexible C++-Implementierung der HNSW-Indexdatenstruktur mit sehr beeindruckenden Performance-Benchmarks liefert. Wir nutzen derzeit nur eine Teilmenge der von usearch gebotenen Funktionen und Tuning-Optionen, freuen uns aber darauf zu erkunden, wie wir künftig mehr davon nutzen können. Bisher sind wir vor allem froh, dass es so gut zu DuckDBs Entwicklungsethos passt. Wie DuckDB selbst ist usearch in portablem C++11 ohne externe Abhängigkeiten geschrieben und unter einer permissiven Lizenz veröffentlicht, sodass die Integration in unsere Extension-Build- und -Verteilungspipeline sehr glatt war.
Einschränkungen
Die große Einschränkung derzeit: Der HNSW-Index kann nur in In-Memory-Datenbanken erzeugt werden, sofern nicht der Konfigurationsparameter SET hnsw_enable_experimental_persistence = ⟨bool⟩{:.language-sql .highlight} auf true gesetzt ist. Ist dieser Parameter nicht gesetzt, führt jeder Versuch, einen HNSW-Index in einer festplattenbasierten Datenbank zu erzeugen, zu einer Fehlermeldung. Ist er gesetzt, wird der Index nicht nur im Speicher erzeugt, sondern beim Checkpointing auch als Teil der DuckDB-Datenbankdatei auf die Festplatte persistiert. Nach einem Neustart oder dem Laden einer Datenbankdatei mit persistiertem HNSW-Index wird der Index beim ersten Zugriff auf die zugehörige Tabelle lazy wieder in den Speicher geladen, was deutlich schneller ist als den Index von Grund auf neu zu erzeugen.
Der Grund, dieses Feature hinter einer experimentellen Flag zu sperren, ist, dass wir noch bekannte Probleme bei der Persistenz benutzerdefinierter Indizes haben, die wir vor dem Standard-Einschalten lösen wollen. Insbesondere ist WAL-Recovery für benutzerdefinierte Indizes noch nicht richtig implementiert: Stürzt das System ab oder wird die Datenbank unerwartet beendet, während es uncommittete Änderungen an einer HNSW-indizierten Tabelle gibt, kann es zu Datenverlust oder Beschädigung des Index kommen. Technisch ist eine manuelle Wiederherstellung nach einem unerwarteten Shutdown möglich, indem man DuckDB zuerst separat startet, die vss-Extension lädt und dann die Datenbankdatei mit ATTACH anhängt, sodass die HNSW-Index-Funktionalität während der WAL-Wiedergabe verfügbar ist – darauf sollten Sie sich für Produktionslasten aber nicht verlassen.
Wir arbeiten aktiv daran, das und andere Probleme rund um Index-Persistenz zu beheben; das soll hoffentlich in DuckDB v0.10.3 landen. Vorerst empfehlen wir, den HNSW-Index nur in In-Memory-Datenbanken zu nutzen.
Zur Laufzeit muss der HNSW-Index wie der ART vollständig in den RAM passen, und der vom HNSW zur Laufzeit allokierte Speicher liegt „außerhalb“ des DuckDB-Speicherverwaltungssystems, respektiert also DuckDBs Konfigurationsparameter memory_limit nicht.
Eine weitere aktuelle Einschränkung des HNSW-Index ist, dass er für die Array-Elemente nur den Typ FLOAT (32-Bit-Gleitkommazahl einfacher Genauigkeit) und nur Distanzmetriken zu den drei eingebauten Distanzfunktionen array_distance, array_negative_inner_product und array_cosine_distance unterstützt. Auch das wollen wir in naher Zukunft erweitern; das ist weniger eine technische Grenze als ein „wir sind noch nicht dazu gekommen“.
Fazit
Die vss-Extension für DuckDB ist eine neue Extension, die HNSW-Indizes auf Listen fester Länge in DuckDB ermöglicht und Abfragen zur Vektorsuche beschleunigt. Die Extension lässt sich derzeit auf DuckDB v0.10.2 auf allen unterstützten Plattformen (einschließlich Wasm!) mit INSTALL vss; LOAD vss installieren. Die vss-Extension betritt für DuckDB-Extensions Neuland, indem sie einen eigenen Indextyp bereitstellt, und wir freuen uns, diese Funktionalität weiter zu verfeinern und auszubauen.
Wir arbeiten noch an einigen der obigen Einschränkungen, besonders rund um Persistenz (und Performance), wollen diese frühe Version der vss-Extension aber trotzdem teilen, weil wir glauben, dass sie der Community viele spannende Möglichkeiten eröffnet. Schauen Sie sich die vss-Extension-Dokumentation an, um mehr über die Arbeit mit dieser Extension zu erfahren!
Diese Arbeit wurde durch das Sponsoring eines DuckLabs-Kunden möglich! Wenn Sie an ähnlicher Arbeit für bestimmte Fähigkeiten interessiert sind, wenden Sie sich bitte an DuckLabs. Alternativ heißen wir Beitragende gerne willkommen! Melden Sie sich beim DuckLabs-Team auf Discord oder im vss-Extension-GitHub-Repository, um die neuesten Entwicklungen zu verfolgen.