Run TileDB-Vector-Search Locally

ai/ml
vector search
tutorials
local access
How to install TileDB-Vector-Search and interact with TileDB Cloud from your local machine.

This tutorial shows how you can install TileDB-Vector-Search and the TileDB Cloud Python client so that you can use the software from your local machine.

Installation

To install the TileDB-Vector-Search and TileDB Cloud Python client, run the following in bash:

# Install the TileDB Cloud Python client
pip install tiledb-cloud

# Install the TileDB-Vector-Search library
# - Using Pip
pip install tiledb-vector-search
# - Or, using conda
conda install -c conda-forge -c tiledb tiledb-vector-search

Local storage

To create and query vector indexes in local storage, read the Quickstart or Tutorials: Ingestion and Querying tutorials.

Storage on S3

To create and query vector indexes on S3, read the Tutorials: Basic S3 example tutorial.

Storage on TileDB Cloud

Take note of the following when interacting with TileDB Cloud outside of the notebook environments of the TileDB Cloud service:

  1. You need to create a REST API token from the TileDB Cloud console.
  2. You need to set the REST API token in an environment variable and login using the token value.
  3. When creating a new index, you need to use a URI in the form tiledb://<your_username>/<S3_path>/<index_name>, where S3_path is the location on S3 where you wish to physically store the index.
  4. When referring to the index after creating it (for example, when submitting queries), use a URI in the form tiledb://<your_username>/<index_name> (that is, no need to specify the S3 physical path anymore).

Start by importing the necessary libraries, loading the appropriate token and username, logging in, setting up the URIs and cleaning up any previously created index with the same name.

# Import necessary libraries
import os

import numpy as np
import tiledb
import tiledb.cloud
import tiledb.vector_search as vs

# You should set the appropriate environment variables with
# your TileDB Cloud token and username.
token = os.environ["TILEDB_REST_TOKEN"]
username = os.environ["TILEDB_REST_USERNAME"]

# Get the bucket from an environment variable
s3_bucket = os.environ["S3_BUCKET"]

# Login to TileDB Cloud
tiledb.cloud.login(token=token)

# Set index URI
index_name = "running_locally"
index_uri = "tiledb://" + username + "/" + index_name
index_reg_uri = "tiledb://" + username + "/" + s3_bucket + "/" + index_name

# The TileDB Cloud context
ctx = tiledb.cloud.Ctx()

# Clean up index if it already exists
if tiledb.object_type(index_uri, ctx=ctx) == "group":
    tiledb.cloud.asset.delete(index_uri, recursive=True)

Create the index, using the URI that incorporates the S3 physical path as explained above:

# Create an index, where the dimensionality of each vector is 3,
# the type of the vector values is float32, and the index will
# use 3 partitions.
index = vs.ivf_flat_index.create(
    ctx=ctx,
    uri=index_reg_uri,
    dimensions=3,
    partitions=3,
    vector_type=np.dtype(np.float32),
)

Ingest some vectors.

# Apply a set of appends to the index, adding one vector at a time
update_vectors = np.empty([5], dtype=object)
update_vectors[0] = np.array([0, 0, 0], dtype=np.dtype(np.float32))
update_vectors[1] = np.array([1, 1, 1], dtype=np.dtype(np.float32))
update_vectors[2] = np.array([2, 2, 2], dtype=np.dtype(np.float32))
update_vectors[3] = np.array([3, 3, 3], dtype=np.dtype(np.float32))
update_vectors[4] = np.array([4, 4, 4], dtype=np.dtype(np.float32))
index.update_batch(vectors=update_vectors, external_ids=np.array([0, 1, 2, 3, 4]))

Submit a query.

# Create a query
query_vector = np.array([[2, 2, 2]], dtype=np.float32)

# Perform the query
result_d, result_i = index.query(query_vector, k=3, nprobe=3)
print("Result vector ids:\n")
print(result_i)
print("\nResult vector distances:\n")
print(result_d)
Result vector ids:

[[2 3 1]]

Result vector distances:

[[0. 3. 3.]]

Clean up.

# Clean up
if tiledb.object_type(index_uri, ctx=ctx) == "group":
    tiledb.cloud.asset.delete(index_uri, recursive=True)