Laion-400M dataset

The Laion-400M dataset contains 400 million images with English image captions. Laion nowadays provides an even larger dataset but working with it will be similar.

The dataset contains the image URL, embeddings for both the image and the image caption, a similarity score between the image and the image caption, as well as metadata, e.g. the image width/height, the licence and a NSFW flag. We can use the dataset to demonstrate approximate nearest neighbor search in ClickHouse.

Data preparation

The embeddings and the metadata are stored in separate files in the raw data. A data preparation step downloads the data, merges the files, converts them to CSV and imports them into ClickHouse. You can use the following download.sh script for that:

Script process.py is defined as follows:

To start the data preparation pipeline, run:

The dataset is split into 410 files, each file contains ca. 1 million rows. If you like to work with a smaller subset of the data, simply adjust the limits, e.g. seq 0 9 | ....

(The python script above is very slow (~2-10 minutes per file), takes a lot of memory (41 GB per file), and the resulting csv files are big (10 GB each), so be careful. If you have enough RAM, increase the -P1 number for more parallelism. If this is still too slow, consider coming up with a better ingestion procedure - maybe converting the .npy files to parquet, then doing all the other processing with clickhouse.)

Create table

To create a table without indexes, run:

To import the CSV files into ClickHouse:

Run a brute-force ANN search (without ANN index)

To run a brute-force approximate nearest neighbor search, run:

target is an array of 512 elements and a client parameter. A convenient way to obtain such arrays will be presented at the end of the article. For now, we can run the embedding of a random cat picture as target.

Result

Run a ANN with an ANN index

Create a new table with an ANN index and insert the data from the existing table:

By default, Annoy indexes use the L2 distance as metric. Further tuning knobs for index creation and search are described in the Annoy index documentation. Let's check now again with the same query:

Result

The speed increased significantly at the cost of less accurate results. This is because the ANN index only provide approximate search results. Note the example searched for similar image embeddings, yet it is also possible to search for positive image caption embeddings.

Creating embeddings with UDFs

One usually wants to create embeddings for new images or new image captions and search for similar image / image caption pairs in the data. We can use UDF to create the target vector without leaving the client. It is important to use the same model to create the data and new embeddings for searches. The following scripts utilize the ViT-B/32 model which also underlies the dataset.

Text embeddings

First, store the following Python script in the user_scripts/ directory of your ClickHouse data path and make it executable (chmod +x encode_text.py).

encode_text.py:

Then create encode_text_function.xml in a location referenced by <user_defined_executable_functions_config>/path/to/*_function.xml</user_defined_executable_functions_config> in your ClickHouse server configuration file.

You can now simply use:

The first run will be slow because it loads the model, but repeated runs will be fast. We can then copy the output to SET param_target=... and can easily write queries.

Image embeddings

Image embeddings can be created similarly but we will provide the Python script the path to a local image instead of the image caption text.

encode_image.py

encode_image_function.xml

Then run this query: