Quick tip: Using Approximate Nearest Neighbor (ANN) Search with SingleStoreDB
Akmal Chaudhri
Posted on January 19, 2024
Abstract
The new SingleStoreDB release v8.5 provides several new vector features. In this short article, we'll evaluate ANN Index Search with the new VECTOR
data type using the Fashion MNIST dataset from Zalando.
The notebook file and SQL code are available on GitHub.
Create a SingleStoreDB Cloud account
A previous article showed the steps to create a free SingleStoreDB Cloud account. We'll use the following settings:
- Workspace Group Name: ANN Demo Group
- Cloud Provider: AWS
- Region: US East 1 (N. Virginia)
- Workspace Name: ann-demo
- Size: S-00
Create a Database and Tables
In our SingleStore Cloud account, let's use the SQL Editor to create a new database. Call this fmnist_db
, as follows:
CREATE DATABASE IF NOT EXISTS fmnist_db;
We'll also create several tables using the BLOB
data type and new VECTOR
data type, as follows:
USE fmnist_db;
CREATE TABLE IF NOT EXISTS train_data_blob (
idx INT(10) UNSIGNED NOT NULL,
label VARCHAR(20),
vector BLOB,
KEY(idx)
);
CREATE TABLE IF NOT EXISTS test_data_blob (
idx INT(10) UNSIGNED NOT NULL,
label VARCHAR(20),
vector BLOB,
KEY(idx)
);
CREATE TABLE IF NOT EXISTS train_data_vec (
idx INT(10) UNSIGNED NOT NULL,
label VARCHAR(20),
vector VECTOR(784) NOT NULL,
KEY(idx)
);
CREATE TABLE IF NOT EXISTS test_data_vec (
idx INT(10) UNSIGNED NOT NULL,
label VARCHAR(20),
vector VECTOR(784) NOT NULL,
KEY(idx)
);
We have train
and test
tables using both formats. We'll load data into the two different sets of tables.
New notebook
We'll follow the instructions to create a new notebook as described in a previous article. We'll call the notebook ann_demo.
Fill out the Notebook
First, we'll install some libraries:
!pip install matplotlib --quiet
!pip install tensorflow --quiet
Next, let's set up our environment:
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import struct
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
from tensorflow import keras
from keras.datasets import fashion_mnist
Load the Dataset
We'll use the Fashion MNIST dataset from Zalando.
First, we'll get the train
and test
data:
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
Let's take a look at the shape of the data:
print("train_images: " + str(train_images.shape))
print("train_labels: " + str(train_labels.shape))
print("test_images: " + str(test_images.shape))
print("test_labels: " + str(test_labels.shape))
The result should be as follows:
train_images: (60000, 28, 28)
train_labels: (60000,)
test_images: (10000, 28, 28)
test_labels: (10000,)
We have 60,000 images for training and 10,000 images for testing. The images are greyscaled, 28 pixels by 28 pixels, and we can take a look at one of these:
print(train_images[0])
The result should be (28 columns by 28 rows):
[[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 13 73 0 0 1 4 0 0 0 0 1 1 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 3 0 36 136 127 62 54 0 0 0 1 3 4 0 0 3]
[ 0 0 0 0 0 0 0 0 0 0 0 0 6 0 102 204 176 134 144 123 23 0 0 0 0 12 10 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 155 236 207 178 107 156 161 109 64 23 77 130 72 15]
[ 0 0 0 0 0 0 0 0 0 0 0 1 0 69 207 223 218 216 216 163 127 121 122 146 141 88 172 66]
[ 0 0 0 0 0 0 0 0 0 1 1 1 0 200 232 232 233 229 223 223 215 213 164 127 123 196 229 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 183 225 216 223 228 235 227 224 222 224 221 223 245 173 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 193 228 218 213 198 180 212 210 211 213 223 220 243 202 0]
[ 0 0 0 0 0 0 0 0 0 1 3 0 12 219 220 212 218 192 169 227 208 218 224 212 226 197 209 52]
[ 0 0 0 0 0 0 0 0 0 0 6 0 99 244 222 220 218 203 198 221 215 213 222 220 245 119 167 56]
[ 0 0 0 0 0 0 0 0 0 4 0 0 55 236 228 230 228 240 232 213 218 223 234 217 217 209 92 0]
[ 0 0 1 4 6 7 2 0 0 0 0 0 237 226 217 223 222 219 222 221 216 223 229 215 218 255 77 0]
[ 0 3 0 0 0 0 0 0 0 62 145 204 228 207 213 221 218 208 211 218 224 223 219 215 224 244 159 0]
[ 0 0 0 0 18 44 82 107 189 228 220 222 217 226 200 205 211 230 224 234 176 188 250 248 233 238 215 0]
[ 0 57 187 208 224 221 224 208 204 214 208 209 200 159 245 193 206 223 255 255 221 234 221 211 220 232 246 0]
[ 3 202 228 224 221 211 211 214 205 205 205 220 240 80 150 255 229 221 188 154 191 210 204 209 222 228 225 0]
[ 98 233 198 210 222 229 229 234 249 220 194 215 217 241 65 73 106 117 168 219 221 215 217 223 223 224 229 29]
[ 75 204 212 204 193 205 211 225 216 185 197 206 198 213 240 195 227 245 239 223 218 212 209 222 220 221 230 67]
[ 48 203 183 194 213 197 185 190 194 192 202 214 219 221 220 236 225 216 199 206 186 181 177 172 181 205 206 115]
[ 0 122 219 193 179 171 183 196 204 210 213 207 211 210 200 196 194 191 195 191 198 192 176 156 167 177 210 92]
[ 0 0 74 189 212 191 175 172 175 181 185 188 189 188 193 198 204 209 210 210 211 188 188 194 192 216 170 0]
[ 2 0 0 0 66 200 222 237 239 242 246 243 244 221 220 193 191 179 182 182 181 176 166 168 99 58 0 0]
[ 0 0 0 0 0 0 0 40 61 44 72 41 35 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]]
We can check the label associated with this image:
print(train_labels[0])
The result should be:
9
This value represents an Ankle Boot.
We can do a quick plot, as follows:
classes = [
"t_shirt_top",
"trouser",
"pullover",
"dress",
"coat",
"sandal",
"shirt",
"sneaker",
"bag",
"ankle_boot"
]
num_classes = len(classes)
for i in range(num_classes):
ax = plt.subplot(2, 5, i + 1)
plt.imshow(
np.column_stack(train_images[i].reshape(1, 28, 28)),
cmap = plt.cm.binary
)
plt.axis("off")
ax.set_title(classes[train_labels[i]])
The result is shown in Figure 1.
Prepare Pandas Dataframe
We need to reshape our dataset so that we can store it correctly later:
train_images = train_images.reshape((train_images.shape[0], -1))
test_images = test_images.reshape((test_images.shape[0], -1))
and we can check the shapes:
print("train_images: " + str(train_images.shape))
print("test_images: " + str(test_images.shape))
The result should be:
train_images: (60000, 784)
test_images: (10000, 784)
So, we have flattened the image structure.
Now we'll create two Pandas Dataframes, as follows:
train_data_df = pd.DataFrame([
(i,
image.astype(int).tolist(),
classes[int(label)],
) for i, (image, label) in enumerate(zip(train_images, train_labels))
], columns = ["idx", "img", "label"])
test_data_df = pd.DataFrame([
(i,
image.astype(int).tolist(),
classes[int(label)],
) for i, (image, label) in enumerate(zip(test_images, test_labels))
], columns = ["idx", "img", "label"])
We need to convert the values in the img
column to a suitable format for SingleStoreDB. We can do this using the following code:
def data_to_binary(data: list[float]):
format_string = "f" * len(data)
return struct.pack(format_string, *data)
train_data_df["vector"] = train_data_df["img"].apply(data_to_binary)
test_data_df["vector"] = test_data_df["img"].apply(data_to_binary)
We can now drop the img
column:
train_data_df.drop("img", axis = 1, inplace = True)
test_data_df.drop("img", axis = 1, inplace = True)
Write Pandas Dataframes to SingleStoreDB
We are now ready to write the Dataframes train_data_df
and test_data_df
to the tables train_data_blob
and test_data_blob
, respectively.
First, we'll set up the connection to SingleStoreDB:
from sqlalchemy import *
db_connection = create_engine(connection_url)
Finally, we are ready to write the Dataframes to SingleStoreDB. First, train_data_df
:
train_data_df.to_sql(
"train_data_blob",
con = db_connection,
if_exists = "append",
index = False,
chunksize = 1000
)
and then test_data_df
:
test_data_df.to_sql(
"test_data_blob",
con = db_connection,
if_exists = "append",
index = False,
chunksize = 1000
)
Example Queries
Now that we have built our system, we can run some queries using the SQL Editor.
Using the BLOB type
First, let's create two variables:
SET @qv_train_blob = (
SELECT vector
FROM train_data_blob
WHERE idx = 30000
);
SET @qv_test_blob = (
SELECT vector
FROM test_data_blob
WHERE idx = 500
);
In the first case, we are selecting an image vector 50% through the train
data. In the second case, we are selecting an image vector 5% through the test
data.
Now, let's use EUCLIDEAN_DISTANCE
with the train
data:
SELECT label, EUCLIDEAN_DISTANCE(vector, @qv_train_blob) AS score
FROM train_data_blob
ORDER BY score
LIMIT 5;
The result should be:
+-------+-------------------+
| label | score |
+-------+-------------------+
| dress | 0 |
| dress | 570.5322076798119 |
| dress | 612.5422434412177 |
| dress | 653.6390441214478 |
| dress | 665.1052548281363 |
+-------+-------------------+
Next, let's try the same query but use the test
data:
SELECT label, EUCLIDEAN_DISTANCE(vector, @qv_test_blob) AS score
FROM train_data_blob
ORDER BY score
LIMIT 5;
The result should be:
+----------+--------------------+
| label | score |
+----------+--------------------+
| pullover | 1211.59399140141 |
| pullover | 1295.9332544541019 |
| pullover | 1316.508640305866 |
| pullover | 1320.24278070361 |
| pullover | 1346.3539653449236 |
+----------+--------------------+
Using the VECTOR type
First, we'll copy the data from the tables using the BLOB
type to the tables using the VECTOR
type, as follows:
INSERT INTO train_data_vec (idx, label, vector) (
SELECT idx, label, vector
FROM train_data_blob
);
INSERT INTO test_data_vec (idx, label, vector) (
SELECT idx, label, vector
FROM test_data_blob
);
Next, we'll define an index, as follows:
ALTER TABLE train_data_vec ADD VECTOR INDEX (vector)
INDEX_OPTIONS '{
"index_type":"IVF_FLAT",
"nlist":1000,
"metric_type":"EUCLIDEAN_DISTANCE"
}';
Many vector indexing options are available. Please see the Vector Indexing documentation.
First, let's create two variables:
SET @qv_train_vec = (
SELECT vector
FROM train_data_vec
WHERE idx = 30000
);
SET @qv_test_vec = (
SELECT vector
FROM test_data_vec
WHERE idx = 500
);
In the first case, we are selecting an image vector 50% through the train
data. In the second case, we are selecting an image vector 5% through the test
data.
Now, let's use the Infix Operator <->
with the train
data:
SELECT label, vector <-> @qv_train_vec AS score
FROM train_data_vec
ORDER BY score
LIMIT 5;
The result could be:
+-------+-------------------+
| label | score |
+-------+-------------------+
| dress | 0 |
| dress | 570.5322076798119 |
| dress | 612.5422434412177 |
| dress | 653.6390441214478 |
| dress | 665.1052548281363 |
+-------+-------------------+
Next, let's try the same query but use the test
data:
SELECT label, vector <-> @qv_test_vec AS score
FROM train_data_vec
ORDER BY score
LIMIT 5;
The result could be:
+----------+--------------------+
| label | score |
+----------+--------------------+
| pullover | 1211.59399140141 |
| pullover | 1295.9332544541019 |
| pullover | 1316.508640305866 |
| pullover | 1320.24278070361 |
| pullover | 1346.3539653449236 |
+----------+--------------------+
Comparing the results, we can see that both approaches work well. However, the new ANN Index Search provides many benefits, as discussed in the Vector Indexing documentation.
Summary
In this short article, we've seen how to create an ANN Index using the new VECTOR
data type with a well-known dataset. We've seen that the existing approach to storing vectors in SingleStoreDB using the BLOB
type works well, but using the new vector features offers greater flexibility and choices.
Check out the blog post on the new VECTOR
data type for additional hints and tips.
License
The MIT License (MIT) Copyright © [2017] Zalando SE, https://tech.zalando.com
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
Posted on January 19, 2024
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.
Related
January 19, 2024