mirrors/scylladb
1
0
Fork 0
mirror of https://github.com/scylladb/scylladb.git synced 2026-04-23 18:10:39 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
d1b4fd56833d34f2fc2ca214ce6e618cfabc79fe
scylladb/test/alternator/test_vector.py
Nadav Har'El 58538e18e8 test/alternator: tests for vector index support 
In this patch we add a large collection of basic functional tests for the
vector index support, covering the CreateTable, UpdateTable,  DescribeTable
and Query operations and the various ways in which those are allowed to
work - or expected to fail. These tests were written in parallel with
writing the code so they (hopefully) cover all the corner cases considered
during development, and make sure these corner cases are all handled
correctly and will not regress in the future.

Some of these tests do not involve querying of the index and focus on
the structure of requests and the kind of syntax allowed. But other tests
are end-to-end, requiring the vector store to be running and trying to
index Alternator data and query it. These tests are marked
"needs_vector_store", and are immediately skipped in Scylla is not
configured to connect to a vector store. In a later patch we'll add a
an option to test/alternator/run to be able to run these end-to-end
tests by automatically running both Scylla and the Vector Store.

We'll have additional end-to-end tests in the vector-store repository.

Note that vector search is a new API feature that doesn't exist in DynamoDB,
so we are adding new parameters and outputs to existing operations. The AWS
SDKs don't normally allow doing that, so the test added here begins by
teaching the Python SDK to use the new APIs we added. This piece of code
can also be used by end-users to use vector search (at least in Python...)
before we officially add this support to ScyllaDB's SDK wrappers.
2026-04-16 14:30:17 +03:00

2209 lines
111 KiB
Python
Raw Blame History

# Copyright 2026-present ScyllaDB
#
# SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1
# Tests for the vector search feature. This is an Alternator extension
# that does not exist on DynamoDB, so most tests in this file are skipped
# when running against DynamoDB (a few tests that don't use the "vs" fixture
# can still run on DynamoDB).
import pytest
import time
import decimal
from decimal import Decimal
from contextlib import contextmanager
from functools import cache
from botocore.exceptions import ClientError
import boto3.dynamodb.types
from .util import random_string, new_test_table, unique_table_name, scylla_config_read, scylla_config_write, client_no_transform, is_aws
# Monkey-patch the boto3 library to stop doing its own error-checking on
# numbers. This works around a bug https://github.com/boto/boto3/issues/2500
# of incorrect checking of responses, and we also need to get boto3 to not do
# its own error checking of requests, to allow us to check the server's
# handling of such errors.
# This is needed at least for test_numeric_list_precision_range().
boto3.dynamodb.types.DYNAMODB_CONTEXT = decimal.Context(prec=100)
# We want to be able to run these tests using an unmodified boto3 library -
# which doesn't understand the new parameters that Alternator added to
# CreateTable, Query, and so on, and moreover will strip unexpected fields
# in Alternator's responses.
# So the following fixture "vs" is a DynamoDB API connection, similar to our
# usual "dynamodb" fixture, but modified to allow our new vector-search
# parameters in the requests and responses.
#
# Users can use exactly the same code to get vector search support in boto3,
# but the more "official" way would be to modify botocore's JSON configuration
# file, botocore/data/dynamodb/2012-08-10/service-2.json.
@pytest.fixture(scope="module")
def vs(new_dynamodb_session, dynamodb):
if is_aws(dynamodb):
pytest.skip('Scylla-only: vector search extensions not available on DynamoDB')
resource = new_dynamodb_session()
client = resource.meta.client
# Patch the client to support the new APIs:
# All the new parameter "shapes" that we will use below for the
# new parameters of the different operations:
new_shapes = {
# For CreateTable (and also DescribeTable's output)
'VectorIndexes': {
'type': 'list',
'member': {'shape': 'VectorIndex'},
},
'VectorIndex': {
'type': 'structure',
'members': {
'IndexName': {'shape': 'String'},
'VectorAttribute': {'shape': 'VectorAttribute'},
'Projection': {'shape': 'Projection'},
# The following two fields are only returned in DescribeTable's
# output, not accepted in CreateTable's input.
'IndexStatus': {'shape': 'String'},
'Backfilling': {'shape': 'BooleanObject'},
},
'required': ['IndexName', 'VectorAttribute'],
},
'VectorAttribute': {
'type': 'structure',
'members': {
'AttributeName': {'shape': 'String'},
'Dimensions': {'shape': 'Integer'},
},
'required': ['AttributeName', 'Dimensions'],
},
# For UpdateTable:
'VectorIndexUpdates': {
'type': 'list',
'member': {'shape': 'VectorIndexUpdate'},
},
'VectorIndexUpdate': {
'type': 'structure',
'members': {
'Create': {'shape': 'CreateVectorIndexAction'},
'Delete': {'shape': 'DeleteVectorIndexAction'},
}
},
'CreateVectorIndexAction': {
'type': 'structure',
'members': {
'IndexName': {'shape': 'String'},
'VectorAttribute': {'shape': 'VectorAttribute'},
'Projection': {'shape': 'Projection'},
},
'required': ['IndexName', 'VectorAttribute'],
},
'DeleteVectorIndexAction': {
'type': 'structure',
'members': {
'IndexName': {'shape': 'String'},
},
'required': ['IndexName'],
},
# For Query:
'VectorSearch': {
'type': 'structure',
'members': {
'QueryVector': {'shape': 'AttributeValue'},
},
'required': ['QueryVector'],
},
}
# Register the new shapes:
service_model = client.meta.service_model
shape_resolver = service_model._shape_resolver
for shape_name, shape_def in new_shapes.items():
shape_resolver._shape_map[shape_name] = shape_def
# Evict any cached shapes for these names
shape_resolver._shape_cache.pop(shape_name, None)
# Add a VectorIndexes parameter to CreateTable
create_table_op = service_model.operation_model('CreateTable')
input_shape = create_table_op.input_shape
input_shape._shape_model['members']['VectorIndexes'] = {
'shape': 'VectorIndexes'
}
input_shape._cache.pop('members', None)
# Add VectorIndexUpdates parameter to UpdateTable
update_table_op = service_model.operation_model('UpdateTable')
input_shape = update_table_op.input_shape
input_shape._shape_model['members']['VectorIndexUpdates'] = {
'shape': 'VectorIndexUpdates'
}
input_shape._cache.pop('members', None)
# Add a VectorSearch parameter to Query
query_op = service_model.operation_model('Query')
input_shape = query_op.input_shape
input_shape._shape_model['members']['VectorSearch'] = {
'shape': 'VectorSearch'
}
input_shape._cache.pop('members', None)
# Add a VectorIndexes field to "TableDescription", the shape returned
# by DescribeTable and also CreateTable
output_shape = shape_resolver.get_shape_by_name('TableDescription')
output_shape._shape_model['members']['VectorIndexes'] = {
'shape': 'VectorIndexes'
}
output_shape._cache.pop('members', None)
shape_resolver._shape_cache.pop('TableDescription', None)
yield resource
# A simple test for the vector type. In vector search, a vector is simply
# an array of known size that contains only numbers. In the DynamoDB API,
# there is no special "vector" type, it's just an regular "list" type,
# and the indexing code may later require that it contain only numbers or
# have a specific length. When this test was written, this vector is stored
# inefficiently as a JSON string with ASCII representation of numbers, but
# in the future, we may decide to recognize such numeric-only lists and
# store them on disk in an optimized way - and still this test will need
# to continue passing.
def test_vector_value(dynamodb, test_table_s):
p = random_string()
v = [Decimal("0"), Decimal("1.2"), Decimal("-2.3"), Decimal("1.2e10")]
test_table_s.put_item(Item={'p': p, 'v': v})
assert test_table_s.get_item(Key={'p': p}, ConsistentRead=True)['Item']['v'] == v
# Even if we will have an optimized storage format for a numeric-only list,
# we will still need to use some form of "decimal" type (variable precision
# based on decimal digits) to support DynamoDB's full numeric precision and
# range (38 decimal digits, exponent up to 125) even in a list. This test
# confirms that Alternator indeed allows that full precision and range
# (which is different from any hardware floating-point type) inside lists.
# See similar tests but for a single number in test_number.py.
def test_numeric_list_precision_range(test_table_s):
p = random_string()
v = [Decimal("3.1415926535897932384626433832795028841"),
Decimal("9.99999999e125")]
test_table_s.put_item(Item={'p': p, 'v': v})
assert test_table_s.get_item(Key={'p': p}, ConsistentRead=True)['Item']['v'] == v
# Test CreateTable creating a new table with a basic vector index. This test
# doesn't check that the vector index actually works - we'll do this in
# separate tests below. It just tests that the new Alternator-only
# CreateTable parameter "VectorIndexes" isn't rejected or otherwise fails.
# This test also doesn't cover all the different parameters inside
# VectorIndexes.
def test_createtable_vectorindexes(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 4}
}]) as table:
pass
# Test that in CreateTable's VectorIndexes, a IndexName and VectorAttribute
# is required. Inside the VectorAttribute, a AttributeName and Dimensions
# are required. With any of those fields missing we get a ValidationException.
def test_createtable_vectorindexes_missing_fields(vs):
# Note: in new_dynamodb_session in conftest.py, we used
# parameter_validation=False by default, so boto3 doesn't do the
# validation of missing parameters for us, which is good, because
# it allows us to send requests with missing fields and see the server
# catch that error.
for bad in bad_vector_indexes:
with pytest.raises(ClientError, match='ValidationException'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[bad]) as table:
pass
bad_vector_indexes = [
# everything missing:
{},
# VectorAttribute missing:
{'IndexName': 'hello'},
# IndexName missing:
{'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 4}},
# VectorAttribute missing parts:
{'IndexName': 'hello', 'VectorAttribute': {'Dimensions': 4}},
{'IndexName': 'hello', 'VectorAttribute': {'AttributeName': 'v'}},
]
# Check that we are not allowed to create two VectorIndexes with the same
# name.
def test_createtable_vectorindexes_same_name(vs):
with pytest.raises(ClientError, match='ValidationException.*Duplicate.*hello'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 4}
},
{ 'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 7}
}
]) as table:
pass
# Check that we are not allowed to a VectorIndexes with the same name as
# the name of another type of index - GSI or an LSI.
def test_createtable_vectorindexes_same_name_gsi(vs):
with pytest.raises(ClientError, match='ValidationException.*Duplicate.*hello'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
GlobalSecondaryIndexes=[
{ 'IndexName': 'hello',
'KeySchema': [{ 'AttributeName': 'p', 'KeyType': 'HASH' }],
'Projection': { 'ProjectionType': 'ALL' }
}],
VectorIndexes=[
{ 'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 7}
}]
) as table:
pass
def test_createtable_vectorindexes_same_name_lsi(vs):
with pytest.raises(ClientError, match='ValidationException.*Duplicate.*hello'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' },
{ 'AttributeName': 'c', 'KeyType': 'RANGE' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' },
{ 'AttributeName': 'c', 'AttributeType': 'S' },
{ 'AttributeName': 'x', 'AttributeType': 'S' }],
LocalSecondaryIndexes=[
{ 'IndexName': 'hello',
'KeySchema': [{ 'AttributeName': 'p', 'KeyType': 'HASH' },
{ 'AttributeName': 'x', 'KeyType': 'RANGE' }],
'Projection': { 'ProjectionType': 'ALL' }
}],
VectorIndexes=[
{ 'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 7}
}]
) as table:
pass
# Test that if a table is created to use vnodes instead of the modern default
# of tablets, then it can't use a vector index because vector index is
# officially supported only with tablets.
# When we finally remove vnode support from the code, this test should be
# deleted.
def test_createtable_vectorindexes_vnodes_forbidden(vs):
with pytest.raises(ClientError, match='ValidationException.*vnodes'):
with new_test_table(vs,
# set system:initial_tablets to a non-number to disable tablets:
Tags=[{'Key': 'system:initial_tablets', 'Value': 'none'}],
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 7}
}]
) as table:
pass
# Verify that a vector index's IndexName follows the same naming rules as
# table names - name length from 3 up to 192 (max_table_name_length) and
# match the regex [a-zA-Z0-9._-]+.
# Note that these rules are similar, but not identical, to the rules for
# IndexName for GSI/LSI (tested in test_gsi.py and test_lsi.py) - there,
# Alternator doesn't put the limit on the length of the GSI/LSI's IndexName,
# but puts a limit (222) on the sum of the table's name and GSI/LSI's name.
def test_createtable_vectorindexes_indexname_rules(vs):
# Forbidden names: shorter than 3 characters, longer than 192
# characters, or containing characters outside [a-zA-Z0-9._-].
# These names should be rejected
for bad_name in ['xy', 'x'*193, 'hello$world', 'hello world']:
with pytest.raises(ClientError, match='ValidationException.*IndexName'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': bad_name,
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 74 }
}]
) as table:
pass
# Allowed names: exactly 3 characters, 192 characters, and using
# all characters from [a-zA-Z0-9._-].
# This test is slightly slower than usual, because three tables and
# indexes will be successfully created and then immediately deleted.
for good_name in ['xyz', 'x'*192,
'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-']:
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': good_name,
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 74 }
}]
) as table:
pass
# Check that the "Dimensions" property in CreateTable's VectorIndexes's
# VectorAttribute must be an integer between 1 and 16000 (MAX_VECTOR_DIMENSION
# in the code).
MAX_VECTOR_DIMENSION = 16000
def test_createtable_vectorindexes_dimensions_rules(vs):
# Forbidden dimensions: non-integer, negative, zero, and above
# MAX_VECTOR_DIMENSION. These dimensions should be rejected.
for bad_dimensions in ['hello', 1.2, -17, 0, MAX_VECTOR_DIMENSION+1]:
with pytest.raises(ClientError, match='ValidationException.*Dimensions'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'vector_index',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': bad_dimensions }
}]
) as table:
pass
# Allowed dimensions: 1, MAX_VECTOR_DIMENSION:
for good_dimensions in [1, MAX_VECTOR_DIMENSION]:
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'vector_index',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': good_dimensions }
}]
) as table:
pass
# Check that the "AttributeName" property in CreateTable's VectorIndexes's
# VectorAttribute must not be a key column (of the base table or any of its
# GSIs or LSIs). This is because key columns have a declared type, which
# can't be a vector (a list), so making such a column the key of a vector
# index makes no sense.
def test_createtable_vectorindexes_attributename_key(vs):
# Forbidden AttributeName: base-table keys (hash and range), GSI keys,
# LSI keys:
for bad_attr in ['p', 'c', 'x', 'y', 'z']:
with pytest.raises(ClientError, match='ValidationException.*AttributeName'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' },
{ 'AttributeName': 'c', 'KeyType': 'RANGE' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' },
{ 'AttributeName': 'c', 'AttributeType': 'S' },
{ 'AttributeName': 'x', 'AttributeType': 'S' },
{ 'AttributeName': 'y', 'AttributeType': 'S' },
{ 'AttributeName': 'z', 'AttributeType': 'S' },
],
VectorIndexes=[
{ 'IndexName': 'vector_index',
'VectorAttribute': {'AttributeName': bad_attr, 'Dimensions': 42 }
}],
GlobalSecondaryIndexes=[
{ 'IndexName': 'gsi',
'KeySchema': [{ 'AttributeName': 'x', 'KeyType': 'HASH' },
{ 'AttributeName': 'y', 'KeyType': 'RANGE' }],
'Projection': { 'ProjectionType': 'ALL' }
}],
LocalSecondaryIndexes=[
{ 'IndexName': 'lsi',
'KeySchema': [{ 'AttributeName': 'p', 'KeyType': 'HASH' },
{ 'AttributeName': 'z', 'KeyType': 'RANGE' }],
'Projection': { 'ProjectionType': 'ALL' }
}],
) as table:
pass
# Check that the "AttributeName" property in CreateTable's VectorIndexes's
# VectorAttribute is an attribute name, limited exactly like ordinary (non-
# key) attributes to 65535 (DYNAMODB_NONKEY_ATTR_NAME_SIZE_MAX) bytes.
# Note that there is no limitation on which characters are allowed, so we
# don't check that.
def test_createtable_vectorindexes_attributename_len(vs):
# Forbidden AttributeName: empty string, string over 65535
for bad_attr in ['', 'x'*65536]:
with pytest.raises(ClientError, match='ValidationException.*AttributeName'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'vector_index',
'VectorAttribute': {'AttributeName': bad_attr, 'Dimensions': 42 }
}]
) as table:
pass
# Test that we can add two different vector indexes on the same table
# in CreateTable, but they must be on different attributes.
def test_createtable_vectorindexes_multiple(vs):
# Can create two vector indexes on two different attributes:
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'ind1',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 42 }
},
{ 'IndexName': 'ind2',
'VectorAttribute': {'AttributeName': 'y', 'Dimensions': 17 }
},
]) as table:
pass
# But can't create two vector indexes on the same attribute.
# Why don't we allow that? If the two indexes were to request different
# "dimensions", Alternator would not know which vector length to enforce
# when inserting values. But for simplicity (and avoiding wasted space
# and work) we decided not to allow two indexes on the same attribute,
# in any case.
with pytest.raises(ClientError, match='ValidationException.*Duplicate'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'ind1',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 42 }
},
{ 'IndexName': 'ind2',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 42 }
},
]) as table:
pass
# Test that vector indexes are correctly listed in DescribeTable:
def test_describetable_vectorindexes(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'ind1',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 42 }
},
{ 'IndexName': 'ind2',
'VectorAttribute': {'AttributeName': 'y', 'Dimensions': 17 }
},
]) as table:
desc = table.meta.client.describe_table(TableName=table.name)
assert 'Table' in desc
assert 'VectorIndexes' in desc['Table']
vector_indexes = desc['Table']['VectorIndexes']
assert len(vector_indexes) == 2
for vec in vector_indexes:
assert vec['IndexName'] == 'ind1' or vec['IndexName'] == 'ind2'
if vec['IndexName'] == 'ind1':
assert vec['VectorAttribute'] == {'AttributeName': 'x', 'Dimensions': 42}
else: # vec['IndexName'] == 'ind2':
assert vec['VectorAttribute'] == {'AttributeName': 'y', 'Dimensions': 17}
assert vec['Projection'] == {'ProjectionType': 'KEYS_ONLY'}
# Test that like DescribeTable, CreateTable also returns the VectorIndexes
# definition its response
def test_createtable_vectorindexes_returned(vs):
# To look at the response of CreateTable, we need to use the "client"
# interface, not the usual higher-level "resource" interface that we
# usually use in tests - because that doesn't return the actual response.
client = vs.meta.client
table_name = unique_table_name()
resp = client.create_table(
TableName=table_name,
BillingMode='PAY_PER_REQUEST',
KeySchema=[{ 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[{
'IndexName': 'ind',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 42 },
}])
try:
assert 'TableDescription' in resp
assert 'VectorIndexes' in resp['TableDescription']
vector_indexes = resp['TableDescription']['VectorIndexes']
assert len(vector_indexes) == 1
vec = vector_indexes[0]
assert vec['IndexName'] == 'ind'
assert vec['VectorAttribute'] == {'AttributeName': 'x', 'Dimensions': 42}
# Note that today, CreateTable just echoes back the parameters it got
# doesn't add default parameters, so we don't expect to see, for
# example, a "Projection" field in the response because we didn't send
# one. We may change this decision in the future.
#assert vec['Projection'] == {'ProjectionType': 'KEYS_ONLY'}
finally:
# In principle, we need to wait for the table to become ACTIVE before
# deleting it. But this test only runs on Alternator, where
# CreateTable is synchronous anyway, so we don't bother to add a
# waiting loop.
client.delete_table(TableName=table_name)
# Basic test for UpdateTable successfully adding a vector index
def test_updatetable_vectorindex_create(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
# There are no vector indexes yet:
desc = table.meta.client.describe_table(TableName=table.name)
assert 'Table' in desc
assert 'VectorIndexes' not in desc['Table']
# Add a vector index with UpdateTable
table.update(VectorIndexUpdates=[{'Create':
{ 'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 17 }
}}])
# Now describe_table should see the new vector index:
desc = table.meta.client.describe_table(TableName=table.name)
assert 'Table' in desc
assert 'VectorIndexes' in desc['Table']
vector_indexes = desc['Table']['VectorIndexes']
assert len(vector_indexes) == 1
vec = vector_indexes[0]
assert vec['IndexName'] == 'hello'
assert vec['VectorAttribute'] == {'AttributeName': 'x', 'Dimensions': 17}
# Basic test for UpdateTable successfully removing a vector index
def test_updatetable_vectorindex_delete(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[{
'IndexName': 'hello',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 42 }
}]) as table:
# There should be one vector index now:
desc = table.meta.client.describe_table(TableName=table.name)
assert 'Table' in desc
assert 'VectorIndexes' in desc['Table']
assert len(desc['Table']['VectorIndexes']) == 1
# Delete the vector index with UpdateTable
table.update(VectorIndexUpdates=[
{'Delete': { 'IndexName': 'hello' }}])
# Now describe_table should see no vector index:
desc = table.meta.client.describe_table(TableName=table.name)
assert 'Table' in desc
assert 'VectorIndexes' not in desc['Table']
# UpdateTable can't remove a vector index that doesn't exist. We get a
# ResourceNotFoundException.
def test_updatetable_vectorindex_delete_nonexistent(vs):
with pytest.raises(ClientError, match='ResourceNotFoundException'):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
table.update(VectorIndexUpdates=[
{'Delete': { 'IndexName': 'nonexistent' }}])
# Test that in UpdateTable's Create operation, a IndexName and VectorAttribute
# are required. Inside the VectorAttribute, a AttributeName and Dimensions
# are required. With any of those fields missing we get a ValidationException.
def test_updatetable_vectorindex_missing_fields(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
# Note: in new_dynamodb_session in conftest.py, we used
# parameter_validation=False by default, so boto3 doesn't do the
# validation of missing parameters for us, which is good, because
# it allows us to send requests with missing fields and see the server
# catch that error.
for bad in bad_vector_indexes:
with pytest.raises(ClientError, match='ValidationException.*VectorIndexUpdates'):
table.update(VectorIndexUpdates=[{'Create': bad}])
# Test that when adding a vector index with UpdateTable,
# 1. Its name cannot be the same as an existing vector index or GSI or LSI
# 2. Its attribute cannot be a key column (of base, GSI or LSI) or the
# attribute on an existing vector index
def test_updatetable_vectorindex_taken_name_or_attribute(vs):
# We create a table with vector index, GSI and LSI, so we can check
# all the desired cases on a single table.
with new_test_table(vs,
KeySchema=[
{ 'AttributeName': 'p', 'KeyType': 'HASH' },
{ 'AttributeName': 'c', 'KeyType': 'RANGE' }],
AttributeDefinitions=[
{ 'AttributeName': 'p', 'AttributeType': 'S' },
{ 'AttributeName': 'c', 'AttributeType': 'S' },
{ 'AttributeName': 'x', 'AttributeType': 'S' },
{ 'AttributeName': 'y', 'AttributeType': 'S' },
{ 'AttributeName': 'z', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'vec',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 13 }}],
GlobalSecondaryIndexes=[
{ 'IndexName': 'gsi',
'KeySchema': [
{ 'AttributeName': 'x', 'KeyType': 'HASH' },
{ 'AttributeName': 'y', 'KeyType': 'RANGE' }],
'Projection': { 'ProjectionType': 'ALL' }}],
LocalSecondaryIndexes=[
{ 'IndexName': 'lsi',
'KeySchema': [
{ 'AttributeName': 'p', 'KeyType': 'HASH' },
{ 'AttributeName': 'z', 'KeyType': 'RANGE' }],
'Projection': { 'ProjectionType': 'ALL' }
}],
) as table:
# IndexName already in use:
for bad_name in ['vec', 'gsi', 'lsi']:
with pytest.raises(ClientError, match='ValidationException.*already exists'):
table.update(VectorIndexUpdates=[{'Create':
{ 'IndexName': bad_name,
'VectorAttribute': {'AttributeName': 'xyz', 'Dimensions': 17 }
}}])
# AttributeName already in use:
for bad_attr in ['p', 'c', 'x', 'y', 'z', 'v']:
with pytest.raises(ClientError, match='ValidationException.*AttributeName'):
table.update(VectorIndexUpdates=[{'Create':
{ 'IndexName': 'newind',
'VectorAttribute': {'AttributeName': bad_attr, 'Dimensions': 17 }
}}])
# In test_updatetable_vectorindex_taken_name_or_attribute() above we tested
# that we can't add a vector index with the same name as an existing GSI or
# LSI. Here we check that the reverse also holds - we can't add a GSI with
# the same name as an existing vector index.
def test_updatetable_gsi_same_name_as_vector_index(vs):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[
{'IndexName': 'vec',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
with pytest.raises(ClientError, match='ValidationException.*already exists'):
table.meta.client.update_table(
TableName=table.name,
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
GlobalSecondaryIndexUpdates=[{'Create': {
'IndexName': 'vec',
'KeySchema': [{'AttributeName': 'p', 'KeyType': 'HASH'}],
'Projection': {'ProjectionType': 'ALL'}
}}])
# Similarly, we can't add a GSI on an attribute that's already used as a
# vector index attribute.
def test_updatetable_gsi_key_is_vector_attribute(vs):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[
{'IndexName': 'vec',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
# The attribute 'v' is already a vector index target - it cannot
# become the hash key of a new GSI.
with pytest.raises(ClientError, match='ValidationException.*AttributeDefinitions'):
table.meta.client.update_table(
TableName=table.name,
AttributeDefinitions=[{'AttributeName': 'v', 'AttributeType': 'S'}],
GlobalSecondaryIndexUpdates=[{'Create': {
'IndexName': 'gsi',
'KeySchema': [{'AttributeName': 'v', 'KeyType': 'HASH'}],
'Projection': {'ProjectionType': 'ALL'}
}}])
# Test that if a table is created to use vnodes instead of the modern default
# of tablets, then one can't add to it a vector index because vector index is
# officially supported only with tablets. This is the UpdateTable version
# of a similar test for CreateTable above.
# When we finally remove vnode support from the code, this test should be
# deleted.
def test_updatetable_vectorindex_vnodes_forbidden(vs):
with new_test_table(vs,
# set system:initial_tablets to a non-number to disable tablets:
Tags=[{'Key': 'system:initial_tablets', 'Value': 'none'}],
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
with pytest.raises(ClientError, match='ValidationException.*vnodes'):
table.update(VectorIndexUpdates=[{'Create':
{ 'IndexName': 'ind',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 17 }
}}])
# Similar to an above test for the CreateTable case, verify that also for
# UpdateTable create a new vector index, a vector index's IndexName must
# have length from 3 up to 192 (max_table_name_length) and match the regex
# [a-zA-Z0-9._-]+.
def test_updatetable_vectorindex_indexname_bad(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
# Forbidden names: shorter than 3 characters, longer than 192
# characters, or containing characters outside [a-zA-Z0-9._-].
# These names should be rejected
for bad_name in ['xy', 'x'*193, 'hello$world', 'hello world']:
with pytest.raises(ClientError, match='ValidationException.*IndexName'):
table.update(VectorIndexUpdates=[{'Create':
{ 'IndexName': bad_name,
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 17 }
}}])
# Similar to an above test for the CreateTable case, verify that also for
# UpdateTable create a new vector index, a vector index's Dimensions must
# be an integer between 1 and MAX_VECTOR_DIMENSION
def test_updatetable_vectorindex_dimensions_bad(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
for bad_dimensions in ['hello', 1.2, -17, 0, MAX_VECTOR_DIMENSION+1]:
with pytest.raises(ClientError, match='ValidationException.*Dimensions'):
table.update(VectorIndexUpdates=[{'Create':
{ 'IndexName': 'ind',
'VectorAttribute': {'AttributeName': 'x', 'Dimensions': bad_dimensions }
}}])
# Similar to an above test for the CreateTable case, verify that also for
# UpdateTable create a new vector index, a vector index's attribute must
# have between 1 and 65535 bytes.
# Note that we also checked above that it can't be one of the existing keys
# (of base table, GSI or LSI), or an already indexed vector column. Here we
# only test the allowed length limits.
def test_updatetable_vectorindex_attributename_bad_len(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
for bad_attr in ['', 'x'*65536]:
with pytest.raises(ClientError, match='ValidationException.*AttributeName'):
table.update(VectorIndexUpdates=[{'Create':
{ 'IndexName': 'ind',
'VectorAttribute': {'AttributeName': bad_attr, 'Dimensions': 17 }
}}])
# DynamoDB currently limits UpdateTable to only one GSI operation (Create
# or Delete), so we placed the same limit on VectorIndexUpdates - even though
# it's an array, it must have exactly one element. Let's validate this
# limitation is enforced - but if one day we decide to lift it, we can
# and delete this test.
def test_updatetable_vectorindex_just_one_update(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
# Zero operations aren't allowed - it's treated just like a missing
# VectorIndexUpdates, and therefore a do-nothing UpdateTable which
# is not allowed.
with pytest.raises(ClientError, match='ValidationException.*requires one'):
table.update(VectorIndexUpdates=[])
# Two "Create" aren't allowed.
# Again following DynamoDB's lead on GSI, interestingly in this case
# the error is LimitExceededException, not ValidationException.
with pytest.raises(ClientError, match='LimitExceededException.*allows one'):
table.update(VectorIndexUpdates=[
{'Create': {'IndexName': 'ind1', 'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 17 }}},
{'Create': {'IndexName': 'ind2', 'VectorAttribute': {'AttributeName': 'y', 'Dimensions': 17 }}}])
# Two "Delete" aren't allowed (they are rejected even before noticing
# that the indexes we ask to delete don't exist).
with pytest.raises(ClientError, match='LimitExceededException.*allows one'):
table.update(VectorIndexUpdates=[
{'Delete': {'IndexName': 'ind1'}},
{'Delete': {'IndexName': 'ind2'}}])
# Also one "Delete" and one "Create" isn't allowed
with pytest.raises(ClientError, match='LimitExceededException.*allows one'):
table.update(VectorIndexUpdates=[
{'Create': {'IndexName': 'ind1', 'VectorAttribute': {'AttributeName': 'x', 'Dimensions': 17 }}},
{'Delete': {'IndexName': 'ind2'}}])
# Also, it's not allowed to have in one UpdateTable request both a
# VectorIndexUpdates and a GlobalSecondaryIndexUpdates. There is no real
# reason why we can't support this, but since we already don't allow adding
# (or deleting) more than one GSI or more than one vector index in the same
# operation, it makes sense to disallow having both. If one day we decide to
# allow both in the same request, we can delete this test.
def test_updatetable_vector_and_gsi_same_request(vs):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
with pytest.raises(ClientError, match='LimitExceededException'):
table.meta.client.update_table(
TableName=table.name,
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexUpdates=[{'Create': {
'IndexName': 'vec',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}
}}],
GlobalSecondaryIndexUpdates=[{'Create': {
'IndexName': 'gsi',
'KeySchema': [{'AttributeName': 'p', 'KeyType': 'HASH'}],
'Projection': {'ProjectionType': 'ALL'}
}}])
# Test that PutItem still works as expected on a table with a vector index
# created by CreateTable or UpdateTable. It might not work if we set up CDC
# in a broken way that breaks writes.
def test_putitem_vectorindex_createtable(vs):
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }],
VectorIndexes=[
{ 'IndexName': 'vec',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3 }}]
) as table:
p = random_string()
item = {'p': p, 'v': [1,2,3]}
table.put_item(Item=item)
assert item == table.get_item(Key={'p': p}, ConsistentRead=True)['Item']
def test_putitem_vectorindex_updatetable(vs):
# Create the table without a vector index, and add it later:
with new_test_table(vs,
KeySchema=[ { 'AttributeName': 'p', 'KeyType': 'HASH' }],
AttributeDefinitions=[{ 'AttributeName': 'p', 'AttributeType': 'S' }]) as table:
table.update(VectorIndexUpdates=[
{'Create': {'IndexName': 'ind', 'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3 }}}])
# In general we may need to wait here until the vector index is
# ACTIVE, but currently in Alternator we don't need to wait.
p = random_string()
item = {'p': p, 'v': [1,2,3]}
table.put_item(Item=item)
assert item == table.get_item(Key={'p': p}, ConsistentRead=True)['Item']
# Simple test table with a vector index on a 3-dimensional vector column v
# Please note that because this is a shared table, tests that perform
# global queries on it, not filtering to a specific partition, may get
# results from other tests - so such tests will need to create their own
# table instead of using this shared one.
@pytest.fixture(scope="module")
def table_vs(vs):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
yield table
# Test that a Query with a VectorSearch parameter without an IndexName
# is rejected with a ValidationException.
def test_query_vectorsearch_missing_indexname(table_vs):
with pytest.raises(ClientError, match='ValidationException.*IndexName'):
table_vs.query(VectorSearch={'QueryVector': [1, 2, 3]})
# Test that a Query with a VectorSearch parameter with an IndexName
# which does not refer to a valid vector index is rejected with a
# ValidationException. Note that it doesn't really matter if IndexName
# refers to a garbage name or to a real GSI/LSI - the code just checks
# if it's a known vector index name.
def test_query_vectorsearch_wrong_indexname(table_vs):
with pytest.raises(ClientError, match='ValidationException.*is not a vector index'):
table_vs.query(IndexName='nonexistent',
VectorSearch={'QueryVector': [1, 2, 3]})
# Test that a Query on a vector index without a VectorSearch parameter is
# rejected. When VectorSearch isn't specified, the code expects a base table,
# LSI or GSI - which it won't find. But rather than reporting unhelpfully
# that an index by that name doesn't exist, we want to report that this index
# does exist - and is a vector index - so VectorSearch must be specified.
def test_query_vectorindex_no_vectorsearch(table_vs):
with pytest.raises(ClientError, match='ValidationException.*VectorSearch'):
table_vs.query(
IndexName='vind',
KeyConditionExpression='p = :p',
ExpressionAttributeValues={':p': 'x'},
)
# Test that a Query with a VectorSearch parameter that is missing the
# required QueryVector field is rejected with a ValidationException.
def test_query_vectorsearch_missing_queryvector(table_vs):
with pytest.raises(ClientError, match='ValidationException.*QueryVector'):
table_vs.query(
IndexName='vind',
VectorSearch={},
)
# Test that QueryVector must be a list of numbers, automatically (thanks to
# boto3) using the DynamoDB encoding {"L": [{"N": "1"}, ...]}, and must
# have the exact length defined as Dimensions of the vector index - which
# in table_vs is 3.
def test_query_vectorsearch_queryvector_bad(table_vs):
# A non-list QueryVector is rejected:
with pytest.raises(ClientError, match='ValidationException.*list of numbers'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': 'not a list'},
)
# A list of the right length but with non-numeric elements
# should be rejected:
with pytest.raises(ClientError, match='ValidationException.*only numbers'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 'b', 3]}
)
# A numeric list but with the wrong length is rejected:
with pytest.raises(ClientError, match='ValidationException.*length'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 2]}
)
with pytest.raises(ClientError, match='ValidationException.*length'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3, 4]}
)
# Test that a Query with VectorSearch requires a Limit parameter, which
# determines how many nearest neighbors to return. This is somewhat
# different from the usual meaning of "Limit" in a query which is
# optional, and used for pagination of results. Vector search does
# not currently support pagination.
def test_query_vectorsearch_limit_bad(table_vs):
# Limit cannot be missing:
with pytest.raises(ClientError, match='ValidationException.*Limit'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
)
# Limit must be a positive integer:
for bad_limit in ['hello', 1.5, 0, -3]:
with pytest.raises(ClientError, match='ValidationException.*Limit'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
Limit=bad_limit
)
# Test that a Query with VectorSearch does not support ConsistentRead=True,
# just like queries on a GSI.
def test_query_vectorsearch_consistent_read(table_vs):
with pytest.raises(ClientError, match='ValidationException.*Consistent'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
Limit=10,
ConsistentRead=True)
# Test that a Query with VectorSearch does not support pagination via
# ExclusiveStartKey.
def test_query_vectorsearch_exclusive_start_key(table_vs):
with pytest.raises(ClientError, match='ValidationException.*ExclusiveStartKey'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
Limit=10,
ExclusiveStartKey={'p': 'somekey'},
)
# Test that a Query with VectorSearch does not support ScanIndexForward.
# The ordering of vector search results is determined by vector distance,
# not by the sort key, so ScanIndexForward makes no sense and is rejected.
def test_query_vectorsearch_scan_index_forward(table_vs):
with pytest.raises(ClientError, match='ValidationException.*ScanIndexForward'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
Limit=10,
ScanIndexForward=True,
)
with pytest.raises(ClientError, match='ValidationException.*ScanIndexForward'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
Limit=10,
ScanIndexForward=False,
)
# Test that a Query with VectorSearch and an unused element in
# ExpressionAttributeValues is rejected.
def test_query_vectorsearch_unused_expression_attribute_values(table_vs):
with pytest.raises(ClientError, match='ValidationException.*val2'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
Limit=1,
ExpressionAttributeValues={':val2': 'a'},
)
# Test that a Query with VectorSearch and an unused element in
# ExpressionAttributeNames is rejected.
def test_query_vectorsearch_unused_expression_attribute_names(table_vs):
with pytest.raises(ClientError, match='ValidationException.*name2'):
table_vs.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]},
Limit=1,
ProjectionExpression='#name1',
ExpressionAttributeNames={'#name1': 'x', '#name2': 'y'},
)
# Helper function to check a vector store is configured in Scylla
# with the --vector-store-primary-uri option. This can be done, for
# example, by running test/alternator/run with the option "--vs".
# This function needs some table as a parameter; calling it again
# for the same table will use a cached result.
@cache
def vector_store_configured(table_vs):
# Issue a trial query to detect whether Scylla was started with a vector
# store URI. If we get an error message "Vector Store is disabled", it
# means the vector store is not configured. If we get any other error or
# success - it means the vector store is configured (but might not be
# ready yet - individual tests will use their own retry loops).
try:
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [0, 0, 0]},
Limit=1)
except ClientError as e:
if 'Vector Store is disabled' in e.response['Error']['Message']:
return False
return True
# Fixture to skip a test if the vector store is not configured.
# It is assumed that if Scylla is configured to use the vector store, then
# the reverse is also true - the vector store is configured to use Scylla,
# so we can check the end-to-end functionality.
@pytest.fixture(scope="module")
def needs_vector_store(table_vs):
if not vector_store_configured(table_vs):
pytest.skip('Vector Store is not configured (run with --vs)')
# The context manager unconfigured_vector_store() temporarily (for the
# duration of the "with" block) un-configures the vector store in Scylla -
# the vector_store_primary_uri configuration option. This allows testing the
# behavior when the vector store is not configured, even if we are testing
# on a setup where it is configured.
@contextmanager
def unconfigured_vector_store(vs):
# As mentioned in issue #28225, we can't write an empty string to the
# configuration due to a bug. But luckily, we can write any garbage which
# isn't a valid URI, and this will be considered unconfigured.
# We also can't restore an empty configuration due to #28225.
# When #28225 is fixed, this entire function can be simplified to just:
# with scylla_config_temporary_string(vs, 'vector_store_primary_uri', ''):
# yield
# Instead we need to use the following mess:
original_value = scylla_config_read(vs, 'vector_store_primary_uri')
if original_value == '""':
# nothing to do, or to restore
yield
return
assert original_value.startswith('"') and original_value.endswith('"')
original_value = original_value[1:-1]
scylla_config_write(vs, 'vector_store_primary_uri', 'garbage')
try:
yield
finally:
scylla_config_write(vs, 'vector_store_primary_uri', original_value)
# If the vector store is not configured, then Query with VectorSearch is
# rejected with a ValidationException saying "Vector Store is disabled".
def test_query_vector_store_disabled(vs, table_vs):
with unconfigured_vector_store(vs):
with pytest.raises(ClientError, match='ValidationException.*Vector Store is disabled'):
table_vs.query(IndexName='vind', VectorSearch={'QueryVector': [0, 0, 0]},
Limit=1)
# Test that even if the vector store is not configured, it is possible to
# create a vector index on the table - but DescribeTable will always show
# that it is CREATING, not ACTIVE.
# I'm not convinced it is a good idea to allow create vector indexes if
# the vector store isn't even configured in Scylla, but currently we do
# allow it.
def test_vectorindex_status_without_vector_store(vs):
with unconfigured_vector_store(vs):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
desc = table.meta.client.describe_table(TableName=table.name)
vector_indexes = desc['Table']['VectorIndexes']
assert len(vector_indexes) == 1
assert vector_indexes[0]['IndexName'] == 'vind'
assert vector_indexes[0]['IndexStatus'] == 'CREATING'
# Timeout (in seconds) used by the retry loops in tests that wait for the
# vector store to index data. Centralized here so it can be adjusted easily.
VECTOR_STORE_TIMEOUT = 20
# Test that a vector search Query returns the nearest-neighbour item.
# The vector store is eventually consistent: after put_item the ANN index
# takes time to reflect the new item, so we retry until it appears.
# A private table is used to avoid other tests' data interfering with the
# Limit=1 result. Data is inserted before the index is created so the
# vector store picks it up faster, by prefill scan rather than CDC.
def test_query_vector_prefill(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
p = random_string()
table.put_item(Item={'p': p, 'v': [1, 0, 0]})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1
)
if result.get('Items') and result['Items'][0]['p'] == p:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for vector store to return the expected item')
time.sleep(0.1)
# Same as test_query_vector_prefill but for a table with a clustering key, which
# exercises the separate code path in query_vector() for hash+range tables.
def test_query_vector_with_ck_prefill(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[
{'AttributeName': 'p', 'KeyType': 'HASH'},
{'AttributeName': 'c', 'KeyType': 'RANGE'}],
AttributeDefinitions=[
{'AttributeName': 'p', 'AttributeType': 'S'},
{'AttributeName': 'c', 'AttributeType': 'S'}]) as table:
p = random_string()
c = random_string()
table.put_item(Item={'p': p, 'c': c, 'v': [1, 0, 0]})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1
)
items = result.get('Items', [])
if items and items[0]['p'] == p and items[0]['c'] == c:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for vector store to return the expected item')
time.sleep(0.1)
# Utility function for waiting until the given vector index is ACTIVE, which
# means that when this function returns, we are guaranteed that:
# 1. Queries on this index will succeed.
# 2. The prefill scan of the existing table data has completed, so all items
# that existed in the table before the index was created have been indexed.
# This function uses DescribeTable and waits for the index's IndexStatus to
# become "ACTIVE". This is more elgant than waiting for an actual Query to
# succeed, and also doesn't require knowing the dimensions of this index to
# attempt a real Query.
def wait_for_vector_index_active(table, index_name):
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
desc = table.meta.client.describe_table(TableName=table.name)
for vi in desc.get('Table', {}).get('VectorIndexes', []):
if vi['IndexName'] == index_name and vi['IndexStatus'] == 'ACTIVE':
return
if time.monotonic() > deadline:
pytest.fail(f'Timed out waiting for vector index "{index_name}" to become ACTIVE')
time.sleep(0.1)
# Test that wait_for_vector_index_active(), waiting for IndexStatus==ACTIVE,
# indeed reliably waits for the index to be ready. A Query issued immediately
# after wait_for_vector_index_active() returns should succeed without any
# retry loop, and also returns the prefilled data.
def test_wait_for_vector_index_active(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
p = random_string()
table.put_item(Item={'p': p, 'v': [1, 0, 0]})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
wait_for_vector_index_active(table, 'vind')
# The index is now ACTIVE: the prefill scan has completed and the
# item we inserted is guaranteed to be indexed. Query without catching
# exceptions or retrying.
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1
)
assert result.get('Items') and result['Items'][0]['p'] == p
# The tests test_query_vector_prefill and test_query_vector_with_ck_prefill
# used string keys in the indexed table. In theory, there shouldn't be any
# difference in the vector store's behavior if the keys are of a different
# type (in addition to string, they can be numeric or binary). But in
# practice, the factor store does handle different key types differently,
# and this test used to fail before this was fixed.
# To save a bit of time, we don't test all combinations of hash and range
# key types but test each type at least once as a hash key and a range key.
@pytest.mark.skip(reason="Bug in vector store for non-string keys, fails very slowly so let's skip")
@pytest.mark.parametrize('hash_type,range_type', [
('N', None), ('B', None), ('S', 'N'), ('S', 'B'),
], ids=[
'N', 'B', 'SN', 'SB'])
def test_query_vector_prefill_key_types(vs, needs_vector_store, hash_type, range_type):
key_schema = [{'AttributeName': 'p', 'KeyType': 'HASH'}]
attr_defs = [{'AttributeName': 'p', 'AttributeType': hash_type}]
if range_type is not None:
key_schema.append({'AttributeName': 'c', 'KeyType': 'RANGE'})
attr_defs.append({'AttributeName': 'c', 'AttributeType': range_type})
key = {'S': 'hello', 'N': Decimal('42'), 'B': b'hello'}
with new_test_table(vs, KeySchema=key_schema,
AttributeDefinitions=attr_defs) as table:
p = key[hash_type]
item = {'p': p, 'v': [1, 0, 0]}
if range_type is not None:
c = key[range_type]
item['c'] = c
table.put_item(Item=item)
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
wait_for_vector_index_active(table, 'vind')
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1)
assert len(result['Items']) == 1 and result['Items'] == [item]
# Same as test_query_vector_prefill but whereas in test_query_vector_prefill
# the vector store reads the indexed data by scanning the table, here the
# vector index is created first and only later the data is written, so the
# vector store is expected to pick it up via CDC.
def test_query_vector_cdc(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
# Wait until the vector store is ready (prefill of the empty table
# has completed), to ensure the subsequent write is picked up via CDC.
wait_for_vector_index_active(table, 'vind')
# Now write the item. It should reach the vector store via CDC.
p = random_string()
table.put_item(Item={'p': p, 'v': [1, 0, 0]})
# Retry the query until the newly written item appears in the results.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1
)
if result.get('Items') and result['Items'][0]['p'] == p:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for vector store to return the expected item via CDC')
time.sleep(0.1)
# Similar test to test_query_vector_cdc, where an item is written after the
# vector index is created, but here the item is written using LWT (using a
# ConditionExpression that causes the request to be a read-modify-write
# operation so need to use LWT for most write isolation modes). This is
# important to test because LWT has different code path for recognizing we
# need to write to the CDC log).
def test_query_vector_cdc_lwt(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
wait_for_vector_index_active(table, 'vind')
# Write the item, with a ConditionExpression to guarantee LWT.
p = random_string()
table.put_item(Item={'p': p, 'v': [1, 0, 0]},
ConditionExpression='attribute_not_exists(p)')
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1)
if len(result['Items']) > 0:
break
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for vector store index an item via CDC')
time.sleep(0.1)
assert len(result['Items']) == 1 and result['Items'][0]['p'] == p
# Similar to test_query_vector_cdc, this test also that a vector search Query
# find data inserted after the index was created. But this test adds a twist:
# before creating the index, we insert a malformed value for the vector
# attribute (a string). We check that this malformed is ignored by the initial
# prefill scan, but should not prevent a later write with a well-formed vector
# from being indexed and returned by queries.
@pytest.mark.parametrize('malformed', ['garbage', [1,2]], ids=['string','wrong_length'])
def test_query_vector_cdc_malformed_prefill(vs, needs_vector_store, malformed):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
# Vector index is not yet enabled, so we can insert a string as the
# value of v, without validation.
p1 = random_string()
table.put_item(Item={'p': p1, 'v': malformed})
# Insert another item with a proper vector
p2 = random_string()
table.put_item(Item={'p': p2, 'v': [1, 0, 0]})
# Now create the vector index. The prefill scan will encounter the
# malformed item and must silently ignore it.
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
# Wait for the prefill scan to complete (index becomes ACTIVE).
wait_for_vector_index_active(table, 'vind')
# At this point only p2 should be indexed and returned by a query
result = table.query(IndexName='vind', VectorSearch={'QueryVector': [1, 0, 0]}, Limit=10)
assert len(result['Items']) == 1 and result['Items'][0]['p'] == p2
# Now replace the value of p1 by a properly formed vector. It should
# be eventually picked up by CDC and indexed by the vector index:
table.put_item(Item={'p': p1, 'v': [1, Decimal("0.1"), 0]})
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
result = table.query(IndexName='vind', VectorSearch={'QueryVector': [1, 0, 0]},
Limit=10)
if len(result['Items']) == 2 and {item['p'] for item in result['Items']} == {p2, p1}:
break
if time.monotonic() > deadline:
assert len(result['Items']) == 2 and {item['p'] for item in result['Items']} == {p2, p1}
break
time.sleep(0.1)
# Test like test_query_vector_prefill, but with a query returning multiple
# results. This helps us verify that:
# 1. "Limits" determines the number of results.
# 2. The query_vector() code correctly handles the need to read and
# return multiple items.
# 3. The multiple results are correctly sorted by distance (nearest first).
def test_query_vector_multiple_results(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
# Insert 4 items at known cosine distances from the query vector [1, 0, 0]:
# p1 at [1, 0, 0] - cosine distance 0 (closest, identical direction)
# p2 at [1, 0.1, 0] - cosine distance ~0.005 (2nd, slightly off-axis)
# p3 at [0, 1, 0] - cosine distance 1 (3rd, orthogonal)
# p4 at [-1, 0, 0] - cosine distance 2 (farthest, opposite direction)
# Data is inserted before the vector index is created so the vector
# store picks it up via scan rather than CDC, which finishes faster.
p1, p2, p3, p4 = random_string(), random_string(), random_string(), random_string()
table.put_item(Item={'p': p1, 'v': [Decimal("1"), Decimal("0"), Decimal("0")]})
table.put_item(Item={'p': p2, 'v': [Decimal("1"), Decimal("0.1"), Decimal("0")]})
table.put_item(Item={'p': p3, 'v': [Decimal("0"), Decimal("1"), Decimal("0")]})
table.put_item(Item={'p': p4, 'v': [Decimal("-1"), Decimal("0"), Decimal("0")]})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
expected_order = [p1, p2, p3]
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [Decimal("1"), Decimal("0"), Decimal("0")]},
Limit=3
)
items = result.get('Items', [])
got = [item['p'] for item in items if item['p'] in {p1, p2, p3, p4}]
if got == expected_order:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail(f'Timed out waiting for correct ordered results; last got: {got}, expected {expected_order}')
time.sleep(0.1)
# Same as test_query_vector_multiple_results but for a table with a
# clustering key, to exercise the hash+range code path in query_vector().
def test_query_vector_with_ck_multiple_results(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[
{'AttributeName': 'p', 'KeyType': 'HASH'},
{'AttributeName': 'c', 'KeyType': 'RANGE'}],
AttributeDefinitions=[
{'AttributeName': 'p', 'AttributeType': 'S'},
{'AttributeName': 'c', 'AttributeType': 'S'}]) as table:
p1, p2, p3, p4 = random_string(), random_string(), random_string(), random_string()
c1, c2, c3, c4 = random_string(), random_string(), random_string(), random_string()
table.put_item(Item={'p': p1, 'c': c1, 'v': [Decimal("1"), Decimal("0"), Decimal("0")]})
table.put_item(Item={'p': p2, 'c': c2, 'v': [Decimal("1"), Decimal("0.1"), Decimal("0")]})
table.put_item(Item={'p': p3, 'c': c3, 'v': [Decimal("0"), Decimal("1"), Decimal("0")]})
table.put_item(Item={'p': p4, 'c': c4, 'v': [Decimal("-1"), Decimal("0"), Decimal("0")]})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
expected_order = [(p1, c1), (p2, c2), (p3, c3)]
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [Decimal("1"), Decimal("0"), Decimal("0")]},
Limit=3
)
items = result.get('Items', [])
pcs = {(p1, c1), (p2, c2), (p3, c3), (p4, c4)}
got = [(item['p'], item['c']) for item in items if (item['p'], item['c']) in pcs]
if got == expected_order:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail(f'Timed out waiting for correct ordered results; last got: {got}, expected {expected_order}')
time.sleep(0.1)
# Test that a vector search Query returns, with Select='ALL_ATTRIBUTES', the
# full item content correctly (all attributes, correct key values) in the
# expected order - for multiple results. Two variants are tested via
# parametrize, to exercise two separate code paths in query_vector(), for
# tables with and without clustering keys:
# - no_ck: table with just a hash key
# - with_ck: table with a hash key and a range key
@pytest.mark.parametrize('have_ck', [False, True], ids=['no_ck', 'with_ck'])
def test_query_vector_full_items(vs, needs_vector_store, have_ck):
key_schema = [{'AttributeName': 'p', 'KeyType': 'HASH'}]
attr_defs = [{'AttributeName': 'p', 'AttributeType': 'S'}]
if have_ck:
key_schema.append({'AttributeName': 'c', 'KeyType': 'RANGE'})
attr_defs.append({'AttributeName': 'c', 'AttributeType': 'S'})
with new_test_table(vs,
KeySchema=key_schema,
AttributeDefinitions=attr_defs) as table:
# Build 3 items, each with distinct key(s), a vector, and extra attributes.
# A 4th item is inserted but should not appear with Limit=3.
if have_ck:
# deliberately use just two different p values, so some of the
# returned items have the same p but different c, to exercise yet
# another potentially different code path:
p1 = random_string()
p2 = random_string()
ps = [p1, p1, p2, p2]
else:
ps = [random_string() for _ in range(4)]
vectors = [
[Decimal("1"), Decimal("0"), Decimal("0")], # closest to query
[Decimal("1"), Decimal("0.1"), Decimal("0")], # 2nd
[Decimal("0"), Decimal("1"), Decimal("0")], # 3rd
[Decimal("-1"), Decimal("0"), Decimal("0")], # farthest, excluded
]
items = []
for i, (p, v) in enumerate(zip(ps, vectors)):
item = {'p': p, 'v': v, 'x': f'attr_{i}', 'y': Decimal(str(i * 10))}
if have_ck:
item['c'] = random_string()
items.append(item)
table.put_item(Item=item)
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
# The 3 nearest items in expected distance order (closest first).
expected_items = items[:3]
# Wait until the returned items match the expected list exactly,
# verifying both the full content of each item and their order.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [Decimal("1"), Decimal("0"), Decimal("0")]},
Limit=3,
Select='ALL_ATTRIBUTES'
)
if result.get('Items') == expected_items:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for vector store to return the expected items')
time.sleep(0.1)
# Test that PutItem rejects a vector attribute value that is invalid for
# the declared vector index on that attribute. The index on table_vs declares
# attribute 'v' as a 3-dimensional vector, so putting a non-list, a list of
# wrong length, a list with non-numeric elements, or a list containing a
# number that cannot be represented as a float must all be rejected.
#
# Note that this write rejection feature is nice to have (and mirrors what
# happens in GSI where writes with the wrong type for the indexed column
# are rejected), but was not really necessary: We could have allowed writes
# with the wrong type, and items with a wrong type would simply be ignored
# by the vector index and not returned in vector search results.
def test_putitem_vectorindex_bad_vector(table_vs):
p = random_string()
# Not a list - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.put_item(Item={'p': p, 'v': 'not a list'})
# A list of the wrong length - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.put_item(Item={'p': p, 'v': [1, 2]})
with pytest.raises(ClientError, match='ValidationException'):
table_vs.put_item(Item={'p': p, 'v': [1, 2, 3, 4]})
# A list of the right length but with a non-numeric element - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.put_item(Item={'p': p, 'v': [1, 'hello', 3]})
# A list whose numeric elements can't be represented as a 32-bit float
# (value out of float range) - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.put_item(Item={'p': p, 'v': [1, Decimal('1e100'), 3]})
# Same as test_putitem_vectorindex_bad_vector but using UpdateItem.
def test_updateitem_vectorindex_bad_vector(table_vs):
p = random_string()
# Not a list - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.update_item(Key={'p': p},
UpdateExpression='SET v = :val',
ExpressionAttributeValues={':val': 'not a list'})
# A list of the wrong length - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.update_item(Key={'p': p},
UpdateExpression='SET v = :val',
ExpressionAttributeValues={':val': [1, 2]})
with pytest.raises(ClientError, match='ValidationException'):
table_vs.update_item(Key={'p': p},
UpdateExpression='SET v = :val',
ExpressionAttributeValues={':val': [1, 2, 3, 4]})
# A list of the right length but with a non-numeric element - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.update_item(Key={'p': p},
UpdateExpression='SET v = :val',
ExpressionAttributeValues={':val': [1, 'hello', 3]})
# A list whose numeric elements can't be represented as a 32-bit float
# (value out of float range) - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
table_vs.update_item(Key={'p': p},
UpdateExpression='SET v = :val',
ExpressionAttributeValues={':val': [1, Decimal('1e100'), 3]})
# Same as test_putitem_vectorindex_bad_vector but using BatchWriteItem.
def test_batchwriteitem_vectorindex_bad_vector(table_vs):
p = random_string()
# Not a list - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
with table_vs.batch_writer() as batch:
batch.put_item(Item={'p': p, 'v': 'not a list'})
# A list of the wrong length - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
with table_vs.batch_writer() as batch:
batch.put_item(Item={'p': p, 'v': [1, 2]})
with pytest.raises(ClientError, match='ValidationException'):
with table_vs.batch_writer() as batch:
batch.put_item(Item={'p': p, 'v': [1, 2, 3, 4]})
# A list of the right length but with a non-numeric element - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
with table_vs.batch_writer() as batch:
batch.put_item(Item={'p': p, 'v': [1, 'hello', 3]})
# A list whose numeric elements can't be represented as a 32-bit float
# (value out of float range) - should be rejected:
with pytest.raises(ClientError, match='ValidationException'):
with table_vs.batch_writer() as batch:
batch.put_item(Item={'p': p, 'v': [1, Decimal('1e100'), 3]})
# Test that DeleteItem removes the item from the vector index.
# Two variants are tested via parametrize:
# - without clustering key (no_ck): deleting the only item in a partition
# generates a partition tombstone in CDC
# - with clustering key (with_ck): deleting a row generates a row tombstone
# in CDC, which is a different code path
@pytest.mark.parametrize('with_ck', [False, True], ids=['no_ck', 'with_ck'])
def test_deleteitem_vectorindex(vs, needs_vector_store, with_ck):
key_schema = [{'AttributeName': 'p', 'KeyType': 'HASH'}]
attr_defs = [{'AttributeName': 'p', 'AttributeType': 'S'}]
if with_ck:
key_schema.append({'AttributeName': 'c', 'KeyType': 'RANGE'})
attr_defs.append({'AttributeName': 'c', 'AttributeType': 'S'})
with new_test_table(vs,
KeySchema=key_schema,
AttributeDefinitions=attr_defs,
VectorIndexes=[
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
# Wait until the vector store is ready (empty table prefill done).
wait_for_vector_index_active(table, 'vind')
# Write the item and wait for it to appear in the vector index.
p = random_string()
item = {'p': p, 'v': [1, 0, 0]}
key = {'p': p}
if with_ck:
c = random_string()
item['c'] = c
key['c'] = c
table.put_item(Item=item)
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Select='ALL_ATTRIBUTES',
Limit=1)
if len(result['Items']) > 0:
assert result['Items'][0] == item
break
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for item to appear in vector index')
time.sleep(0.1)
# Delete the item and wait for it to disappear from the vector index.
table.delete_item(Key=key)
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1)
if len(result.get('Items', [])) == 0:
break
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for deleted item to disappear from vector index')
time.sleep(0.1)
# Test vector index with Alternator TTL together. A table is created without
# TTL enabled, data is inserted with expiration time set to the past (but
# expiration not yet enabled), and the item should still appear in vector
# search. Then TTL expiration is enabled and the item should disappear from
# the vector search once TTL deletes it and the deletion propagates via CDC.
# This test is skipped if alternator_ttl_period_in_seconds is not set to a
# low value because otherwise it would take too long to run.
# Two code paths are tested via parametrize:
# - without clustering key (no_ck): partition deletions in CDC.
# - with clustering key (with_ck): row deletions in CDC.
@pytest.mark.parametrize('have_ck', [False, True], ids=['no_ck', 'with_ck'])
def test_vector_with_ttl(vs, needs_vector_store, have_ck):
period = scylla_config_read(vs, 'alternator_ttl_period_in_seconds')
if period is None or float(period) > 1:
pytest.skip('need alternator_ttl_period_in_seconds <= 1 to run this test quickly')
key_schema = [{'AttributeName': 'p', 'KeyType': 'HASH'}]
attr_defs = [{'AttributeName': 'p', 'AttributeType': 'S'}]
if have_ck:
key_schema.append({'AttributeName': 'c', 'KeyType': 'RANGE'})
attr_defs.append({'AttributeName': 'c', 'AttributeType': 'S'})
with new_test_table(vs,
KeySchema=key_schema,
AttributeDefinitions=attr_defs,
VectorIndexes=[
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}
]) as table:
# Wait until the vector store is ready (prefill of the empty table
# has completed), to ensure the rest of the test doesn't need to
# the vector store not yet being up (we'll still need to wait for
# specific data to be indexed, but the index itself will be ready)
wait_for_vector_index_active(table, 'vind')
p = random_string()
item = {'p': p, 'expiration': int(time.time()) - 60, 'v': [1, 0, 0]}
if have_ck:
c = random_string()
item['c'] = c
# Insert an item with 'expiration' set to the past, before TTL is enabled.
# The item should still be visible (and indexed) because TTL is not yet
# configured on this table.
table.put_item(Item=item)
# Wait for the item to appear in vector search. Since TTL is not yet
# enabled, the item must be visible despite its past expiration time.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1)
if len(result.get('Items', [])) > 0:
assert result['Items'][0]['p'] == p
break
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for item to appear in vector search before TTL was enabled')
time.sleep(0.1)
# Now enable TTL on the 'expiration' attribute. The item has its
# expiration in the past, so TTL should delete it quickly.
table.meta.client.update_time_to_live(
TableName=table.name,
TimeToLiveSpecification={'AttributeName': 'expiration', 'Enabled': True})
# Wait for the item to disappear from vector search. TTL deletes the
# item from the database, and the deletion propagates to the vector
# store via CDC.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT + float(period)
while True:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Select='ALL_PROJECTED_ATTRIBUTES',
Limit=1
)
if len(result['Items']) == 0:
break
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for TTL-expired item to disappear from vector search')
time.sleep(0.1)
# Since we used Select='ALL_PROJECTED_ATTRIBUTES', the loop above
# already confirms the vector store removed the item (the results
# come directly from the vector store, not the base table).
assert result['Items'] == []
# Test support for "Select" parameter in vector search Query.
# We test all valid Select values and their effects on the returned items,
# as well as validation errors for invalid combinations.
# The first part tests validation errors (no vector store needed), and
# the second part tests correct results (needs vector store).
def test_query_vectorsearch_select_bad(table_vs):
# Unknown Select value
with pytest.raises(ClientError, match='ValidationException.*Select'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]}, Limit=1,
Select='GARBAGE')
# Select=SPECIFIC_ATTRIBUTES without ProjectionExpression or AttributesToGet
with pytest.raises(ClientError, match='ValidationException.*SPECIFIC_ATTRIBUTES'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]}, Limit=1,
Select='SPECIFIC_ATTRIBUTES')
# ProjectionExpression with Select=ALL_ATTRIBUTES is not allowed
with pytest.raises(ClientError, match='ValidationException.*SPECIFIC_ATTRIBUTES'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]}, Limit=1,
Select='ALL_ATTRIBUTES', ProjectionExpression='p')
# ProjectionExpression with Select=COUNT is not allowed
with pytest.raises(ClientError, match='ValidationException.*SPECIFIC_ATTRIBUTES'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]}, Limit=1,
Select='COUNT', ProjectionExpression='p')
# ProjectionExpression with Select=ALL_PROJECTED_ATTRIBUTES is not allowed
with pytest.raises(ClientError, match='ValidationException.*SPECIFIC_ATTRIBUTES'):
table_vs.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 2, 3]}, Limit=1,
Select='ALL_PROJECTED_ATTRIBUTES', ProjectionExpression='p')
def test_query_vectorsearch_select(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
p = random_string()
# Insert data before creating the vector index so the vector store
# picks it up via prefill scan rather than CDC (faster).
table.put_item(Item={'p': p, 'v': [1, 0, 0], 'x': 'hello', 'y': 'world'})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
# Wait for the item to appear in vector search.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1)
if result.get('Items') and result['Items'][0]['p'] == p:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for item to be indexed')
time.sleep(0.1)
# ALL_PROJECTED_ATTRIBUTES (default when no Select): returns only the
# primary key attributes.
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1)
assert result['Items'] == [{'p': p}]
# Explicit Select=ALL_PROJECTED_ATTRIBUTES:
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
Select='ALL_PROJECTED_ATTRIBUTES')
assert result['Items'] == [{'p': p}]
# Select=ALL_ATTRIBUTES: returns the full item.
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
Select='ALL_ATTRIBUTES')
assert result['Items'] == [{'p': p, 'v': [1, 0, 0], 'x': 'hello', 'y': 'world'}]
# Select=SPECIFIC_ATTRIBUTES with ProjectionExpression: returns only
# the specified attributes.
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
Select='SPECIFIC_ATTRIBUTES', ProjectionExpression='p, x')
assert result['Items'] == [{'p': p, 'x': 'hello'}]
# ProjectionExpression without Select: defaults to SPECIFIC_ATTRIBUTES.
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
ProjectionExpression='p, y')
assert result['Items'] == [{'p': p, 'y': 'world'}]
# Can also use ProjectionExpression with ExpressionAttributeNames:
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
ProjectionExpression='#name1, #name2',
ExpressionAttributeNames={'#name1': 'p', '#name2': 'x'})
assert result['Items'] == [{'p': p, 'x': 'hello'}]
# Select=SPECIFIC_ATTRIBUTES with AttributesToGet: returns only the
# specified attributes.
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
Select='SPECIFIC_ATTRIBUTES', AttributesToGet=['p', 'x'])
assert result['Items'] == [{'p': p, 'x': 'hello'}]
# AttributesToGet without Select: defaults to SPECIFIC_ATTRIBUTES.
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
AttributesToGet=['p', 'y'])
assert result['Items'] == [{'p': p, 'y': 'world'}]
# Select=COUNT: returns only the count, no items list.
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]}, Limit=1,
Select='COUNT')
assert 'Items' not in result
assert result['Count'] == 1
# Test that invalid Projection parameter values are rejected for both
# CreateTable and UpdateTable's vector index creation.
def test_vector_projection_bad(vs):
bad_projections = [
# 'not_an_object', # We can't check this with boto3
{'ProjectionType': 'GARBAGE'},
{}, # missing ProjectionType
]
for bad_projection in bad_projections:
with pytest.raises(ClientError, match='ValidationException.*Projection'):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[{
'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3},
'Projection': bad_projection,
}]) as table:
pass
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
for bad_projection in bad_projections:
with pytest.raises(ClientError, match='ValidationException.*Projection'):
table.update(VectorIndexUpdates=[{'Create': {
'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3},
'Projection': bad_projection,
}}])
# Test that a vector index created with Projection={'ProjectionType': 'KEYS_ONLY'}
# (via CreateTable or UpdateTable) works correctly:
# - The ProjectionType=KEYS_ONLY is accepted
# - Select=ALL_PROJECTED_ATTRIBUTES returns only the primary key attributes
# - Select=ALL_ATTRIBUTES returns all attributes
# ProjectionType=KEYS_ONLY matches the default vector index behavior, so it
# doesn't change results but must be accepted as a valid parameter.
@pytest.mark.parametrize('via_update', [False, True], ids=['createtable', 'updatetable'])
def test_vector_projection_keys_only(vs, needs_vector_store, via_update):
if via_update:
ctx = new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}])
else:
ctx = new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}],
VectorIndexes=[{
'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3},
'Projection': {'ProjectionType': 'KEYS_ONLY'},
}])
with ctx as table:
if via_update:
table.update(VectorIndexUpdates=[{'Create': {
'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3},
'Projection': {'ProjectionType': 'KEYS_ONLY'},
}}])
p = random_string()
table.put_item(Item={'p': p, 'v': [1, 0, 0], 'x': 'hello'})
wait_for_vector_index_active(table, 'vind')
# Select=ALL_PROJECTED_ATTRIBUTES returns only the primary key.
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1,
Select='ALL_PROJECTED_ATTRIBUTES')
assert result['Items'] == [{'p': p}]
# Select=ALL_ATTRIBUTES returns the full item.
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1,
Select='ALL_ATTRIBUTES')
assert result['Items'] == [{'p': p, 'v': [1, 0, 0], 'x': 'hello'}]
# As we saw in test_item.py::test_attribute_allowed_chars in the DynamoDB API
# attribute names can contain any characters whatsoever, including quotes,
# spaces, and even null bytes. Test that such crazy attribute names can be
# used as vector attributes in vector indexes, and that a vector index with
# such an attribute can be created and used successfully.
def test_vector_attribute_allowed_chars(vs, needs_vector_store):
# To check both scan-based prefill and CDC-based indexing, we create the
# table without a vector index and then add the vector index. Data that
# we added before creating the index needs scan, and data added later
# needs CDC. We want to ensure that both work correctly with such
# attribute names.
attribute_name = 'v with spaces and .-+-&*!#@$%^()\\ \' "quotes" and \0 null byte'
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
p1 = random_string()
table.put_item(Item={'p': p1, attribute_name: [1, 0, 0]})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': attribute_name, 'Dimensions': 3}}}])
wait_for_vector_index_active(table, 'vind')
# The previous item was indexed by a scan. Now let's add another item
# which will get indexed by CDC.
p2 = random_string()
table.put_item(Item={'p': p2, attribute_name: [0, 0, 1]})
# Wait until the CDC-indexed update (v=[0, 0, 1]) is reflected in the
# vector search results.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
result = table.query(IndexName='vind',
VectorSearch={'QueryVector': [0, 0, 1]}, Limit=2)
if 'Items' in result and len(result['Items']) == 2 and result['Items'][0]['p'] == p2 and result['Items'][1]['p'] == p1:
break
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for items to appear in vector search')
time.sleep(0.1)
# Test FilterExpression for post-filtering vector search results: After Limit
# results are found by the vector index and the full items are retrieved
# from the base table, items which do not match the given FilterExpression are
# removed. This means that fewer than Limit results may be returned. This
# matches DynamoDB's general Query behavior where the filtering is applied after
# Limit.
# Two Select values are tested (via parametrize):
# ALL_ATTRIBUTES: the matching items are returned in the Items list.
# COUNT: no items are returned, but the implementation still needs to retrieve
# full items (or at least the attributes needed by the filter) and
# count how many among the Limit candidates matched the filter.
# ScannedCount (number of pre-filtering results) and Count (number of post-
# filtering results) are returned in both cases and checked.
@pytest.mark.parametrize('select', ['ALL_ATTRIBUTES', 'COUNT'])
def test_query_vectorsearch_filter_expression(vs, needs_vector_store, select):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
# Insert all 5 items before the vector index is created so the vector
# store picks them up via prefill scan (faster than CDC).
# p_far is the furthest item and will not be among the 4 nearest
# neighbors returned with Limit=4.
p_keep1, p_keep2 = random_string(), random_string()
p_drop1, p_drop2 = random_string(), random_string()
p_far = random_string()
table.put_item(Item={'p': p_keep1, 'v': [1, 0, 0], 'x': 'keep'})
table.put_item(Item={'p': p_drop1, 'v': [1, Decimal("0.1"), 0], 'x': 'drop'})
table.put_item(Item={'p': p_keep2, 'v': [1, Decimal("0.2"), 0], 'x': 'keep'})
table.put_item(Item={'p': p_drop2, 'v': [1, Decimal("0.3"), 0], 'x': 'drop'})
table.put_item(Item={'p': p_far, 'v': [1, Decimal("0.4"), 0], 'x': 'keep'})
nearest_ps = {p_keep1, p_keep2, p_drop1, p_drop2} # 4 nearest neighbors
keep_ps = {p_keep1, p_keep2} # x='keep' items among 4 nearest neighbors
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
# Wait until nearest 4 items (nearest_ps) are visible in a query
# without a filter.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=4,
)
if {item['p'] for item in result.get('Items', [])} == nearest_ps:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for all items to be indexed')
time.sleep(0.1)
# Query with a FilterExpression that matches 2 of the 4 nearest
# candidates (Limit=4). We expect Count=2 and ScannedCount=4. Note
# that even though p_far also has x=keep, it was not among the 4
# nearest neighbors - so it will not be included.
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=4,
Select=select,
FilterExpression='x = :want',
ExpressionAttributeValues={':want': 'keep'},
)
assert result['Count'] == 2
assert result['ScannedCount'] == 4
if select == 'COUNT':
assert 'Items' not in result
else:
assert {item['p'] for item in result['Items']} == keep_ps
# Test FilterExpression for post-filtering vector search results with
# Select=SPECIFIC_ATTRIBUTES. Here the full items are not returned, but still
# need to be retrieved from the base table - including attributes which are
# needed by the filter but not returned in the final results.
def test_query_vectorsearch_filter_expression_specific_attributes(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
# Same 5-item setup as test_query_vectorsearch_filter_expression.
# p_far is the furthest and won't be among the 4 nearest with Limit=4.
p_keep1, p_keep2 = random_string(), random_string()
p_drop1, p_drop2 = random_string(), random_string()
p_far = random_string()
table.put_item(Item={'p': p_keep1, 'v': [1, 0, 0], 'x': 'keep'})
table.put_item(Item={'p': p_drop1, 'v': [1, Decimal("0.1"), 0], 'x': 'drop'})
table.put_item(Item={'p': p_keep2, 'v': [1, Decimal("0.2"), 0], 'x': 'keep'})
table.put_item(Item={'p': p_drop2, 'v': [1, Decimal("0.3"), 0], 'x': 'drop'})
table.put_item(Item={'p': p_far, 'v': [1, Decimal("0.4"), 0], 'x': 'keep'})
nearest_ps = {p_keep1, p_keep2, p_drop1, p_drop2}
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
# Wait until the 4 nearest items are visible without a filter.
deadline = time.monotonic() + VECTOR_STORE_TIMEOUT
while True:
try:
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=4,
)
if {item['p'] for item in result.get('Items', [])} == nearest_ps:
break
except ClientError:
pass
if time.monotonic() > deadline:
pytest.fail('Timed out waiting for all items to be indexed')
time.sleep(0.1)
# Query with Select=SPECIFIC_ATTRIBUTES projecting only 'p', but
# FilterExpression uses 'x' which is NOT in the projection. The
# implementation must still retrieve 'x' from the base table to
# evaluate the filter, even though 'x' is not returned to the caller.
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=4,
Select='SPECIFIC_ATTRIBUTES',
ProjectionExpression='p',
FilterExpression='x = :want',
ExpressionAttributeValues={':want': 'keep'},
)
assert result['Count'] == 2
assert result['ScannedCount'] == 4
# Items should contain only 'p' (the projected attribute), not 'x'
# (the filter attribute that was not projected).
assert result['Items'] == [{'p': p_keep1}, {'p': p_keep2}]
# Test FilterExpression with Select=SPECIFIC_ATTRIBUTES and a nested
# ProjectionExpression (e.g. 'x.a'). Only the requested sub-attribute
# should be returned, not the entire top-level attribute.
def test_query_vectorsearch_filter_expression_nested_projection(vs, needs_vector_store):
with new_test_table(vs,
KeySchema=[{'AttributeName': 'p', 'KeyType': 'HASH'}],
AttributeDefinitions=[{'AttributeName': 'p', 'AttributeType': 'S'}]) as table:
p = random_string()
# Item has a nested attribute 'x' with sub-attributes 'a' and 'b'.
# The FilterExpression uses 'y', which is not in the projection.
table.put_item(Item={'p': p, 'v': [1, 0, 0], 'x': {'a': 'keep', 'b': 'drop'}, 'y': 'pass'})
table.update(VectorIndexUpdates=[{'Create':
{'IndexName': 'vind',
'VectorAttribute': {'AttributeName': 'v', 'Dimensions': 3}}}])
wait_for_vector_index_active(table, 'vind')
# ProjectionExpression requests only the nested attribute 'x.a' (and 'p').
# FilterExpression uses 'y', which is not in the projection at all.
# The result should contain only 'p' and 'x': {'a': 'keep'} - the
# 'b' sub-attribute of 'x' must not appear, and 'y' must not appear.
result = table.query(
IndexName='vind',
VectorSearch={'QueryVector': [1, 0, 0]},
Limit=1,
Select='SPECIFIC_ATTRIBUTES',
ProjectionExpression='p, x.a',
FilterExpression='y = :want',
ExpressionAttributeValues={':want': 'pass'},
)
assert result['Count'] == 1
assert result['ScannedCount'] == 1
assert result['Items'] == [{'p': p, 'x': {'a': 'keep'}}]
# Test that garbage values (like "dog" or "Inf") for the "N"-typed numbers
# are not allowed as vector attribute values given as a list of numbers.
# They should be rejected with a validation error both before the index is
# created (this test) and after (the next test), because such values are not
# allowed as "N" variables - regardless of vector search.
# This test (the "before") doesn't need vector search and can also run on
# DynamoDB. It reproduces issue #8070 - where Alternator validates number
# values, but forget to validate numbers when they are inside a list.
@pytest.mark.xfail(reason='issue #8070 - Alternator did not validate "N" values inside lists')
def test_putitem_vector_bad_number_string_before(test_table_s):
p = random_string()
# boto3 normally validates number strings before sending them to the
# server, so we need client_no_transform to bypass that validation and
# let the server reject the bad values itself.
with client_no_transform(test_table_s.meta.client) as client:
for bad_num in ['dog', 'Inf', 'NaN', 'Infinity', '-Infinity']:
with pytest.raises(ClientError, match='ValidationException'):
client.put_item(
TableName=test_table_s.name,
Item={
'p': {'S': p},
'v': {'L': [{'N': '1'}, {'N': bad_num}, {'N': '0'}]},
})
def test_putitem_vector_bad_number_string_after(table_vs):
p = random_string()
# After the vector index is created, invalid "N" strings in a list
# must be rejected - they remain invalid DynamoDB numbers.
with client_no_transform(table_vs.meta.client) as client:
for bad_num in ['dog', 'Inf', 'NaN', 'Infinity', '-Infinity']:
with pytest.raises(ClientError, match='ValidationException'):
client.put_item(
TableName=table_vs.name,
Item={
'p': {'S': p},
'v': {'L': [{'N': '1'}, {'N': bad_num}, {'N': '0'}]},
})
# Test that a Query with a vector with a non-numeric "N" element, like "dog"
# or "Inf", is rejected with a validation error. Note that the Query path
# does not convert the numbers to Alternator's internal type ("decimal") so
# the validation path is different, so we need to check it.
def test_query_vectorsearch_queryvector_bad_number_string(table_vs, needs_vector_store):
# boto3 validates number strings before sending them, so we use
# client_no_transform to bypass that and let the server reject them.
with client_no_transform(table_vs.meta.client) as client:
for bad_num in ['dog', 'Inf', 'NaN', 'Infinity', '-Infinity']:
print(bad_num)
with pytest.raises(ClientError, match='ValidationException.*not a valid number'):
client.query(
TableName=table_vs.name,
IndexName='vind',
VectorSearch={'QueryVector': {'L': [{'N': '1'}, {'N': bad_num}, {'N': '0'}]}},
Limit=1,
)
##############################################################################
# CONTINUE HERE - MAKE A DECISION! PRE-FILTERING:
# Tests *pre-filtering* for filtering on projected attributes, which can be (if
# we continue the implementation) pushed to the vector store or right now -
# key columns.
# We need to decide: In DynamoDB pre-filtering is done with
# KeyConditionExpression NOT FilterExpression.
# KeyConditionExpression is the traditional approach in Query, but is a bit
# weird because these aren't really "keys" (even though in CQL CREATE TABLE
# syntax we pretend they are - and today, they really are keys).
# Do we want to force the user to put these pre-filtering in KeyConditionExpression
# instead of FilterExpression, and only allow FilterExpression for post-filtering
# that must be done in Scylla?
# Alternatively, we could put everything in FilterExpression. This is less
# with DynamoDB's usual semantics but simpler for users.
############################################################################
# WRITE TEST:
# Test that if the FilterExpression happens to only use projected attributes
# (by default this means key attributes) - or if we decide (see above) that
# it's KeyconditionExpression, then it can be, and is, sent to the vector
# store and performed there. We can check that this happens by noticing that
# we get a full LIMIT of results, and not less.
# WRITE TEST:
# Test FilterExpression for post-filtering vector search results with
# Select=ALL_PROJECTED_ATTRIBUTES where the filter only needs projected
# attributes (currently those are key attributes). Here it is important
# for efficency that the vector index applies the filter and we do not need
# to retrive the full items from the base table at all. We can verify that
# this code path was reached by checking that we got back LIMIT results, and
# not fewer.
# TODO: Like test_vector_projection_keys_only, write additional tests for
# ProjectionType=INCLUDE with NonKeyAttributes and ProjectionType=ALL.
# We don't yet support this feature, so I didn't bother to write such a
# test yet, but I can write an xfailing test because we know what we
# expect ALL_PROJECTED_ATTRIBUTES to return in that case.
# TODO: test enabling vector index and Alternator Streams together, and
# checking that Alternator Streams works as expected. Also we may need to
# do something to avoid vector search's favorite parameters like TTL and
# post-changes to take control - or vice versa we may get CDC which isn't
# good enough for vector search.
# Note that today, Alternator Streams only works with vnodes while vector
# search doesn't work with vnodes - so we can't actually check this
# combination! But we must check it when Alternator Streams finally supports
# tablets.
Reference in New Issue View Git Blame Copy Permalink