Python SDK for Vector Search

Applies to VAST Cluster 5.5. Examples use the VAST DataBase Python SDK 2.1.0.

The VAST DataBase Python SDK, vastdb, can create vector-indexed tables, insert and import rows, and search with table.vector_search(). The search runs as SQL through the VAST Query Engine, which is reachable through the Python SDK and the ADBC library (see supported query engines), so it needs the ADBC driver. This page covers the vector parts of the SDK, end to end. The signatures and field tables are in the Python SDK reference.

Run the code blocks in order in one Python session. Below each block is the output you should see.

Connect with the ADBC driver

vector_search() runs through the Query Engine's ADBC driver, so install the SDK with its adbc extra and connect to the Query Engine endpoint:

pip install "vastdb[adbc]==2.1.0"

import os
import adbc_driver_vastdb
import pyarrow as pa
import vastdb
from vastdb import VectorIndexSpec

BUCKET = os.environ["VASTDB_BUCKET"]

session = vastdb.connect(
    endpoint=os.environ["QE_ENDPOINT"],
    adbc_driver_path=adbc_driver_vastdb.get_driver_path(),
)
print(vastdb.version())

2.1.0

vastdb.connect() reads the keys from AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY. If adbc-driver-vastdb is installed, you can leave out adbc_driver_path: the SDK finds the driver itself.

Without the driver, you can still create tables and insert rows over the SDK's own API, but vector_search() can't run. It raises vastdb.transaction.NoAdbcConnectionError: Adbc Driver may not have been supplied. See ADBC in the VAST DataBase docs.

When a transaction opens its ADBC connection, the driver manager may print Warning: Cannot disable autocommit; conn will not be DB-API 2.0 compliant on stderr. You can ignore it.

Create an indexed table

Pass a VectorIndexSpec to create_table(). The vector column must be a fixed-size Arrow list of float32 whose items aren't nullable:

dim = 128
vec = pa.list_(pa.field("item", pa.float32(), nullable=False), dim)
arrow_schema = pa.schema([("id", pa.int64()), ("embedding", vec)])

with session.transaction() as tx:
    schema = tx.bucket(BUCKET).create_schema("vs_sdk", fail_if_exists=False)
    table = schema.create_table(
        "t",
        arrow_schema,
        vector_index=VectorIndexSpec(
            column="embedding",
            distance_metric="l2sq",
            cluster_quantization="scalar_8bit",
            search_quantization="scalar_8bit",
        ),
        fail_if_exists=False,
    )
    table.insert(pa.table(schema=arrow_schema, data=[
        [1, 2, 3, 4, 5],
        [[float(i)] * dim for i in range(5)],
    ]))
    print(table.vector_index)

VectorIndex(column='embedding', distance_metric='l2sq', sql_distance_function='array_distance', cluster_quantization='scalar_8bit', search_quantization='scalar_8bit')

VectorIndexSpec takes:

Argument	Values on VAST Cluster 5.5	Default
`column`	The name of the vector column	Required
`distance_metric`	`"l2sq"`, `"cosine"` or `"ip"`	Required
`cluster_quantization`	`"scalar_8bit"` or `"none"`	The cluster's default, `"scalar_8bit"`
`search_quantization`	`"scalar_8bit"` or `"none"`	The cluster's default, `"scalar_8bit"`

The SDK 2.1.0 docstring lists "l2sq" and "ip"; the cluster also accepts "cosine". If you leave a quantization out, the cluster's default applies.

The SDK needs cluster 5.5 or later for vector tables. On an older cluster it raises Vector tables require 5.5+ VAST release.

Invalid quantization pairs

The cluster rejects some combinations when you create the table. The SDK raises vastdb.errors.BadRequest:

try:
    with session.transaction() as tx:
        tx.bucket(BUCKET).schema("vs_sdk").create_table(
            "bad_pair",
            arrow_schema,
            vector_index=VectorIndexSpec(
                column="embedding",
                distance_metric="l2sq",
                cluster_quantization="scalar_8bit",
                search_quantization="none",
            ),
        )
except vastdb.errors.BadRequest as e:
    print(type(e).__name__, e.status, e.code)
    print(e.message.split("]: ", 1)[-1])

BadRequest 400 InvalidArgument
search_quantization='none' requires cluster_quantization='none'

Some failed create_table() calls still leave a table behind. A vector index with a vastdb_rowid column fails with failed to enable vector index eres=14, but creates the table without an index. A vector column over the maximum dimension fails, but leaves a table that schema.tables() lists and that answers drop() with 404 NoSuchKey. After a failed create, list the schema's tables before you try again.

The SDK also logs the failed request and the rollback on stderr. See quantization.

Read the index definition

table.vector_index holds the index definition. The SDK loads it when you open the table, and again when you call reload_schema(). The table's statistics carry the same definition, in table.stats.vector_index, which reload_stats() refreshes:

with session.transaction() as tx:
    table = tx.bucket(BUCKET).schema("vs_sdk").table("t")
    table.reload_schema()
    table.reload_stats()
    print(table.stats.vector_index == table.vector_index)
    print(table.vector_index.sql_distance_function)

True
array_distance

sql_distance_function is the SQL function that vector_search() uses for this index.

Search

with session.transaction() as tx:
    table = tx.bucket(BUCKET).schema("vs_sdk").table("t")
    table.reload_schema()
    reader = table.vector_search(
        vec=[0.1] * dim,
        columns=["id", "embedding"],
        limit=3,
        predicate=table["id"] > 1,
    )
    result = reader.read_all()

print(result.column_names)
print(result["id"].to_pylist())
print([round(d, 3) for d in result["vast_pysdk_vector_dist"].to_pylist()])

['id', 'embedding', 'vast_pysdk_vector_dist']
[2, 3, 4]
[10.182, 21.496, 32.81]

vector_search() takes:

vec: the query vector, as a list of floats with the column's dimension.
columns: the columns to return. The distance is added as vast_pysdk_vector_dist.
limit: how many rows to return.
predicate: an optional filter, written like a select() predicate. See predicates.

The full signature is in the SDK reference.

It returns a pyarrow.RecordBatchReader. Read it before the transaction ends.

Under the hood, vector_search() builds one SQL query and runs it on the transaction's ADBC connection: the table's sql_distance_function on the vector column and the query vector, cast to FLOAT[<dimension>], then your predicate as a WHERE clause, ORDER BY the distance, and LIMIT. That's the query shape the index uses; see query design.

Inner-product indexes: search with SQL

For an index with distance_metric="ip", don't use vector_search() on VAST Cluster 5.5 with SDK 2.1.0. If you search from the SDK, create the index with l2sq or cosine. The cluster reports array_inner_product as the index's SQL function, and vector_search() sorts it in ascending order. So it returns the least similar rows first, and the query doesn't match the index, which ranks inner product with array_negative_inner_product.

For an existing ip table, run the search as SQL on the transaction's ADBC connection instead, with array_negative_inner_product and ORDER BY ascending. See the SQL reference.

vec2 = pa.list_(pa.field("item", pa.float32(), nullable=False), 2)
ip_schema = pa.schema([("id", pa.int64()), ("embedding", vec2)])

with session.transaction() as tx:
    ip_table = tx.bucket(BUCKET).schema("vs_sdk").create_table(
        "ip_t",
        ip_schema,
        vector_index=VectorIndexSpec(column="embedding", distance_metric="ip"),
        fail_if_exists=False,
    )
    ip_table.insert(pa.table(schema=ip_schema, data=[
        [1, 2, 3],
        [[1.0, 0.0], [0.0, 1.0], [3.0, 0.0]],
    ]))

with session.transaction() as tx:
    ip_table = tx.bucket(BUCKET).schema("vs_sdk").table("ip_t")
    print(ip_table.vector_index.sql_distance_function)

    sdk_result = ip_table.vector_search(vec=[1.0, 0.0], columns=["id"], limit=2).read_all()
    print("vector_search():", sdk_result["id"].to_pylist())

    cur = tx.adbc_conn.cursor
    cur.execute(f"""
        SELECT id, array_negative_inner_product(embedding::FLOAT[2], [1.0, 0.0]::FLOAT[2]) AS dist
        FROM {ip_table.ref.query_engine_full_path}
        ORDER BY dist
        LIMIT 2""")
    print("SQL:", [row[0] for row in cur.fetchall()])

array_inner_product
vector_search(): [2, 1]
SQL: [3, 1]

Row 3 has the largest inner product with [1.0, 0.0], so it's the best match. table.ref.query_engine_full_path gives the table's full name for SQL.

Clean up

with session.transaction() as tx:
    schema = tx.bucket(BUCKET).schema("vs_sdk")
    for name in ("t", "ip_t"):
        schema.table(name).drop()
    schema.drop()
print("Dropped vs_sdk")

Dropped vs_sdk