Introducing sqlite-vss: A SQLite Extension for Vector Search

sqlite-vss is a new SQLite extension that adds vector search capabilities to SQLite, based on Faiss. It's several steps above storing embeddings as JSON or pickle, but doesn't scale as well as heavier vector databases like Pinecone/Qdrant/Weaviate. Great for local-first applications and no-nonsense deployments!

Vector search is an integral part of new-age AI applications. You make "embeddings" of your data, compare those embeddings using standard vector distance algorithms, and with some elbow grease, you can build entire semantic search engines, recommendation systems, or a Q&A platform. And now with a plethora of embedding generations tools out there like OpenAI's Embeddings API, HuggingFace's Inference API, sentence-transformers, or literally dozens of other open source models out there, it's never been easier to make embeddings of your own data!

Storing and querying these embeddings can be tricky, however. Sure, you can just dump your data into a JSON file or pickle it, but with enough data this can easily blow up in size and rack up storage costs. For querying, you could just manually loop through all your vectors and calculate cosine distances, but an exhaustive search across all your vectors can be time-consuming. Typically this is when people recommend large-scale vector databases like Pinecone or Qdrant, but sometimes they feel overkill or expensive.

sqlite-vss is a new vector search and storage solution that aims to be a "happy medium" between those two approaches. Your embeddings are stored inside a SQLite database and can be queried using SQL, similar to the fts5 full-text search SQLite extension. It's based on Faiss which means querying can be extremely fast and efficient, and still flexible enough to meet your needs.

Once you have your embeddings in a vss0 virtual table, running KNN search is as simple as an extra WHERE clause:

-- "get the 20 most similar headline embeddings to row 123"
select
  rowid,
  distance
from vss_articles
where vss_search(
  headline_embedding,
  (
    select headline_embedding
    from articles
    where rowid = 123
  )
)
limit 20;

Demo

We'll use sqlite-vss on a real-world example: Let's take the News Category Dataset from Kaggle and make a semantic search engine. This dataset has around 210,000 news articles with headlines, short descriptions, links, and some other metadata. We will build a semantic search engine that will search these headlines and descriptions! We will use sentence-transformers for a free "good enough" embeddings generator and host the database with Datasette on a single fly.io app.

The Datasette instance is hosted at this URL, and we'll query that database directly from this Observable notebook with Datasette Client.

Importing the JSON dataset

First we'll import the 200k news article headlines + descriptions into a single SQLite table. The data source is a 84MB newline-delimited JSON file, which we can easily import into a table with sqlite-lines and SQLite's builtin JSON support:

.load ./lines0

create table articles(
  headline text,
  headline_embedding blob,
  description text,
  description_embedding blob,
  link text,
  category text,
  authors text,
  date
);

insert into articles(headline, description, link, category, authors, date)
  select
    line ->> '$.headline'           as headline,
    line ->> '$.short_description'  as description,
    line ->> '$.link'               as link,
    line ->> '$.category'           as category,
    line ->> '$.authors'            as authors,
    line ->> '$.date'               as date
  from lines_read('News_Category_Dataset_v3.json');

Note that we're leaving headline_embedding and description_embedding as NULL, that's coming next!

Now we have a nice articles table with ~200k rows of our raw data:

Generating embeddings with sentence-transformers

We'll use the sentence-transformers library to generate embeddings (vectors) of the 'headline' and 'description' columns. We'll later use these to generate fast indexes with the vss0 virtual table.

import sqlite3
import json

from sentence_transformers import SentenceTransformer
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

db = sqlite3.connect("test.db")

# go through all rows until all headline_embedding values are not null
while True:
  result = db.execute(
    """
      select rowid, headline
      from articles
      where headline_embeddings is null
      limit 1
    """
  ).fetchone()
  if result == None:
    break
  rowid, headline = result
  headline_embedding = model.encode([headline])[0]
  db.execute(
    """
      update articles
      set description_embedding = ?
      where rowid = ?
    """,
    [embedding.tobytes(), rowid]
  )
  db.commit()

In this script, we are looping through all headline values in each row of the articles table, and embedding them with the sentence-transformers/all-MiniLM-L6-v2 model.

Do note that there are many ways to make this more efficient. The script I used to build this instead batched 256 rows together to batch into model.encode(), but there are likely other improvements that could be made. Here's the full Python script for reference.

And here's the result, a full table with BLOB headline_embedding and description_embedding columns!

Building the vss0 Virtual Table

Now that we have the columns, let's build a vector search index for faster queries! We'll create a virtual table named vss_articles using the vss0 module, where every row in vss_articles corresponds to a row in articles, sharing the same rowid.

create virtual table vss_articles using vss0(
  headline_embedding(384),
  description_embedding(384),
);

We specify separate columns for headline_embedding and description_embedding, declaring that the vectors have 384 dimensions.

Now we can insert data into the vss_articles table with rows from the articles table!

insert into vss_articles(rowid, headline_embedding, description_embedding)
  select
    rowid,
    headline_embedding,
    description_embedding
  from articles;

And that's it! We now have ready-to-be-queried indexes of our embeddings.

Querying the Embeddings

You can query the vss_articles table with normal SELECT statements like SELECT * FROM vss_articles LIMIT 5. But to take advantage of the vector indexes that sqlite-vss offers, use the vss_search() function in the WHERE clause and a LIMIT statement to perform KNN-style similarity searches.

In this query, we are getting the headline_embedding value for the row in the articles table with a rowid of 590. We pass it into the vss_search function inside of the WHERE clause, and specify that we want to search on the headline_embedding column of vss_articles.

"Semantic Search" Querying

You can query your own search terms in real time! This Datasette instances exposes a st_encode() function that creates an embedding of whatever text we give it, with the same sentence-transformers model we used to make the article embeddings. To search with our own search terms, we can construct a query similar to the ones above.

Comparisons with...

sqlite-vss is very similar to many vector search tools, libraries, and databases, so here's a quick rundown of how it's different than some of the popular ones.

Pinecone/ Milvus/ Weaviate/ Qdrant

These databases can store millions or billions of vectors, while sqlite-vss has a limit of ~1GB per column. They also scale across multiple machines with proper distributed systems, while sqlite-vss only works on one machine.

On the other hand, sqlite-vss is completely open source and can be self-hosted, while some of these databases can not. sqlite-vss is also fairly simple to setup - it only requires SQLite, a few apt-get installable packages, and the sqlite-vss loadable extension. These other solutions may require several Docker containers, entire dedicated clusters, or worse - a credit card.

Most of these databases also offer additional filtering on top of similarity search. sqlite-vss doesn't support this, but will in the future!

Storing Embeddings as JSON or using pickle

The first low-tech embeddings storage technique people use are typically json.dump()'ing into a JSON file or pickle.dump()'ing into a binary file. If it works for you, then it works for you! But, if it starts to take up too much disk space or the exhaustive cosine queries take too long, then sqlite-vss is an easy no-hassle upgrade that'll likely outperform.

Faiss

sqlite-vss uses Faiss under the hood! But Faiss already lets you serialize an index to a file, so what does sqlite-vss offer on top of that?

For one, a SQL interface. Instead of running index.write_index("file.bin") or index.search() in Python, you might prefer traditional CREATE TABLE/SELECT * FROM xxx SQL syntax instead.

txtai

txtai helps you build end-to-end semantic search tools, but while sqlite-vss can be bundled alongside any SQLite client in any programming language, txtai is a Python library.

datasette-faiss

datasette-faiss is also very similar to sqlite-vss, since they both are based on SQLite-stored vectors and Faiss. However, datasette-faiss only works for in-memory IndexFlatL2 indexes built at startup. sqlite-vss has configurable index factory strings, indexes are stored on disk, and it works with any extension-enabled SQLite client.