AstraPy

A pythonic client for DataStax Astra DB.

This README targets AstraPy version 2.0+. Click here for v1 and here for the v0 API (which you should not really be using by now).

Quickstart

Install with pip install astrapy.

Get the API Endpoint and the Token to your Astra DB instance at astra.datastax.com.

Try the following code after replacing the connection parameters:

from astrapy import DataAPIClient
from astrapy.constants import VectorMetric
from astrapy.ids import UUID
from astrapy.info import CollectionDefinition


ASTRA_DB_APPLICATION_TOKEN = "AstraCS:..."
ASTRA_DB_API_ENDPOINT = "https://01234567-....apps.astra.datastax.com"

# Connect and create the Database object
my_client = DataAPIClient()
my_database = my_client.get_database(
    ASTRA_DB_API_ENDPOINT,
    token=ASTRA_DB_APPLICATION_TOKEN,
)

# Create a vector collection
my_collection = my_database.create_collection(
    "dreams_collection",
    definition=(
        CollectionDefinition.builder()
        .with_vector_dimension(3)
        .with_vector_metric(VectorMetric.COSINE)
        .build()
    )
)

# Populate the collection with some documents
my_collection.insert_many(
    [
        {
            "_id": UUID("018e65c9-e33d-749b-9386-e848739582f0"),
            "summary": "Riding the waves",
            "tags": ["sport"],
            "$vector": [0, 0.2, 1],
        },
        {
            "summary": "Friendly aliens in town",
            "tags": ["scifi"],
            "$vector": [-0.3, 0, 0.8],
        },
        {
            "summary": "Meeting Beethoven at the dentist",
            "$vector": [0.2, 0.6, 0],
        },
    ],
)

my_collection.update_one(
    {"tags": "sport"},
    {"$set": {"summary": "Surfers' paradise"}},
)

# Run a vector search
cursor = my_collection.find(
    {},
    sort={"$vector": [0, 0.2, 0.4]},
    limit=2,
    include_similarity=True,
)

for result in cursor:
    print(f"{result['summary']}: {result['$similarity']}")

# This would print:
#   Surfers' paradise: 0.98238194
#   Friendly aliens in town: 0.91873914

# Resource cleanup
my_collection.drop()

Next steps:

• More info and usage patterns are given in the docstrings of classes and methods
• Data API reference
• AstraPy reference
• Package on PyPI
• For interactive sessions, check the astrapy-repl command here

Server-side embeddings

AstraPy works with the "vectorize" feature of the Data API. This means that one can define server-side computation for vector embeddings and use text strings in place of a document vector, both in writing and in reading. The transformation of said text into an embedding is handled by the Data API, using a provider and model you specify.

my_collection = database.create_collection(
    "my_vectorize_collection",
    definition=(
        CollectionDefinition.builder()
        .with_vector_service(
            provider="example_vendor",
            model_name="embedding_model_name",
            authentication={"providerKey": "<STORED_API_KEY_NAME>"}  # if needed
        )
        .build()
    )
)

my_collection.insert_one({"$vectorize": "text to make into embedding"})

documents = my_collection.find(sort={"$vectorize": "vector search query text"})

See the Data API reference for more on this topic.

Hybrid search

AstraPy supports the "findAndRerank" Data API command, which performs a hybrid search by combining results from a lexical search and a vector-based search in a single operation.

r_results = my_collection.find_and_rerank(
    sort={"$hybrid": "query text"},
    limit=10,
    include_scores=True,
)

for r_result in r_results:
    print(r_result.document, r_results.scores)

This command will execute if the collection was created with the required settings (they are enabled by default).

See the Data API reference, and the docstring for the find_and_rerank method, for more on this topic.

Using Tables

The example above uses a collection, where schemaless "documents" can be stored and retrieved. Here is an equivalent code that uses Tables, i.e. uniform, structured data where each row has the same columns, which are of a specific type:

from astrapy import DataAPIClient
from astrapy.constants import VectorMetric
from astrapy.data_types import DataAPIVector
from astrapy.info import (
    CreateTableDefinition,
    ColumnType,
    TableVectorIndexDefinition,
    TableVectorIndexOptions,
)


ASTRA_DB_APPLICATION_TOKEN = "AstraCS:..."
ASTRA_DB_API_ENDPOINT = "https://01234567-....apps.astra.datastax.com"

# Connect and create the Database object
my_client = DataAPIClient()
my_database = my_client.get_database(
    ASTRA_DB_API_ENDPOINT,
    token=ASTRA_DB_APPLICATION_TOKEN,
)

# Create a table and a vector index on it
table_definition = (
    CreateTableDefinition.builder()
    .add_column("dream_id", ColumnType.INT)
    .add_column("summary", ColumnType.TEXT)
    .add_set_column("tags", ColumnType.TEXT)
    .add_vector_column("dream_vector", dimension=3)
    .add_partition_by(["dream_id"])
    .build()
)
index_options=TableVectorIndexOptions(
    metric=VectorMetric.COSINE,
)
my_table = my_database.create_table("dreams_table", definition=table_definition, if_not_exists=True)
my_table.create_vector_index("dreams_table_vec_idx", column="dream_vector", options=index_options, if_not_exists=True)

# Populate the table with some rows
my_table.insert_many(
    [
        {
            "dream_id": 103,
            "summary": "Riding the waves",
            "tags": ["sport"],
            "dream_vector": DataAPIVector([0, 0.2, 1]),
        },
        {
            "dream_id": 119,
            "summary": "Friendly aliens in town",
            "tags": ["scifi"],
            "dream_vector": DataAPIVector([-0.3, 0, 0.8]),
        },
        {
            "dream_id": 37,
            "summary": "Meeting Beethoven at the dentist",
            "dream_vector": DataAPIVector([0.2, 0.6, 0]),
        },
    ],
)

my_table.update_one(
    {"dream_id": 103},
    {"$set": {"summary": "Surfers' paradise"}},
)

# Run a vector search
cursor = my_table.find(
    {},
    sort={"dream_vector": DataAPIVector([0, 0.2, 0.4])},
    limit=2,
    include_similarity=True,
)

for result in cursor:
    print(f"{result['summary']}: {result['$similarity']}")

# This would print:
#   Surfers' paradise: 0.98238194
#   Friendly aliens in town: 0.91873914

# Resource cleanup
my_table.drop()

For more on Tables, consult the Data API documentation about Tables. Note that most features of Collections, with due modifications, hold for Tables as well (e.g. "vectorize", i.e. server-side embeddings).

Maps as association lists

When working with Tables, map columns with key of a type other than text must be expressed as association lists, i.e. nested lists of lists: [[key1, value1], [key2, value2], ...], in payloads.

By default, Table objects in AstraPy will use the association-list representation only for maps wrapped by a DataAPIMap, and leave regular Python dict objects unchanged in payloads: in practice the following row would be sent to the API as shown:

from astrapy.data_types import DataAPIMap

row_to_insert = {
    "map_1": {"k": "v"},
    "map_2": DataAPIMap({"k1": "v1", "k2": "v2"}),
}

# Will be serialized as JSON like this:
#    {
#        "map_1": {"k": "v"},
#        "map_2": [["k1", "v1"], ["k2", "v2"]]
#    }

This behaviour is controlled by the API Option serdes_options.encode_maps_as_lists_in_tables, which defaults to MapEncodingMode.DATAAPIMAPS. Other options are:

• MapEncodingMode.ALWAYS to convert also regular Python dict values into lists of pairs;
• MapEncodingMode.NEVER to never apply such conversion.

The following code demonstrates getting a Table object with the desired setting:

from astrapy.api_options import APIOptions, SerdesOptions
from astrapy.constants import MapEncodingMode

my_options = APIOptions(serdes_options=SerdesOptions(
    encode_maps_as_lists_in_tables=MapEncodingMode.ALWAYS  # or even just "ALWAYS"
))

my_table = my_database.get_table("my_table", spawn_api_options=my_options)

See the section about API Options, and the docstring, for more details.

Caution: if one plans to use regular Python dict objects to express user-defined types (UDTs), setting this option to ALWAYS would interfere with the format expected by the API for UDTs. See the section on UDTs for more details.

User-defined types (UDTs)

The Data API allows creation of "user-defined types" (UDTs), composite data types that can be subsequently used as table column types.

from astrapy.info import (
    ColumnType,
    CreateTableDefinition,
    CreateTypeDefinition,
    TableScalarColumnTypeDescriptor,
)

player_udt_def = CreateTypeDefinition(
    fields={
        "name": TableScalarColumnTypeDescriptor(ColumnType.TEXT),
        "age": TableScalarColumnTypeDescriptor(ColumnType.INT),
    },
)
my_database.create_type("player_udt", definition=player_udt_def)

table_definition = (
    CreateTableDefinition.builder()
    .add_column("match_id", ColumnType.TEXT)
    .add_userdefinedtype_column("player1", "player_udt")
    .add_userdefinedtype_column("player2", "player_udt")
    .add_partition_by(["match_id"])
    .build()
)
udt_table = my_database.create_table("matches", definition=table_definition)

To write data to UDT columns, the default settings (in part. the SerdesOptions settings) admit using plain Python dictionaries. When reading from a Table, by default the UDT will be returned in the form of an astrapy.data_types.DataAPIDictUDT, a subclass of dict.

The default write behaviour can be changed through the serdes_options.encode_maps_as_lists_in_tables API Option, described in the "Maps as association lists" section.

The default read behaviour can be changed by setting serdes_options.custom_datatypes_in_reading to False, in which case UDT data will be returned as a regular dict.

The following code demonstrates the default read and write behaviour, using the udt_table just created:

from astrapy.data_types import DataAPIDictUDT

# both are valid ways to express a 'player_udt':
udt_table.insert_one({
    "match_id": "x001",
    "player1": {"name": "Anita", "age": 49},
    "player2": DataAPIDictUDT({"name": "Pedro", "age": 50}),
})
# Conversely, use of DataAPIMap for UDTs is an error under the default serdes settings.

Note that if the map list-encoding is set to "ALWAYS", the above insertion will fail because of "player1": usage of DataAPIDictUDT is guaranteed to work in all circumstances instead.

Using models for UDTs

In many cases, one wants to 'bind' a certain UDT to a model class. AstraPy provides a way to "register" model classes both for the write and read paths, through the following serdes API Options:

• serdes_options.serializer_by_class: a map from classes (not class names; not class instances) to serializer functions. A serializer function accepts an instance of the model class and returns a corresponding dict representation;
• serdes_options.deserializer_by_udt: a map from UDT names to deserializer functions. A deserializer function accepts a dict representation, along with the UDT definition, and returns an instance of the model class.

The following example demonstrates this procedure, including ser/deserializers. Suppose you have a dataclass corresponding to player_udt:

from dataclasses import dataclass

@dataclass
class Player:
    name: str
    age: int

The two required ser/des functions can be defined and configured for the API options in order to spawn a Table object able to use Player:

from typing import Any

from astrapy.api_options import APIOptions, SerdesOptions
from astrapy.info import CreateTypeDefinition

def player_serializer(pl: Player) -> dict[str, Any]:
    # the logic in this function will depend on the model class being used:
    return pl.__dict__

def player_deserializer(
    pl_dict: dict[str, Any],
    udt_def: CreateTypeDefinition | None,
) -> Player:
    # the logic in this function will depend on the model class being used:
    return Player(**pl_dict)

my_options = APIOptions(serdes_options=SerdesOptions(
    serializer_by_class={Player: player_serializer},
    deserializer_by_udt={"player_udt": player_deserializer},
))

# This statement does not create the table on DB, that is assumed to exist already:
my_model_capable_table = database.get_table("matches", spawn_api_options=my_options)

At this point, writes and reads can seamlessly use the Player class:

my_model_capable_table.insert_one({
    "match_id": "x001",
    "player1": Player(name="Anita", age=49),
    "player2": Player(name="Pedro", age=50),
})

the_match = my_model_capable_table.find_one({"match_id": "x001"})
# the_match["player1"] and the_match["player2"] are Player objects:

print(the_match["player1"])
# prints: Player(name='Anita', age=49)
print(the_match["player2"])
# prints: Player(name='Pedro', age=50)

Summary for UDT usage and maps in Tables

Under the default serdes settings, the following rules apply and constitute the suggested approach:

• use DataAPIMap to write map columns;
• use DataAPIDictUDTs, or plain dict, to write UDTs;
• ex