Skip to content

Commit 41cec39

Browse files
simorenohsimorenoh
authored
[Cosmos] DiskANN updates (#38044)
* Update README.md * add values to README, update tests * Update README.md * update samples * Update sdk/cosmos/azure-cosmos/README.md * remove vectorIndexShardKey
1 parent 11fe0c6 commit 41cec39

File tree

5 files changed

+123
-17
lines changed

5 files changed

+123
-17
lines changed

‎sdk/cosmos/azure-cosmos/README.md

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -684,7 +684,7 @@ vector_embedding_policy = {
684684
```
685685

686686
Separately, vector indexes have been added to the already existing indexing_policy and only require two fields per index:
687-
the path to the relevant field to be used, and the type of index from the possible options (flat or quantizedFlat).
687+
the path to the relevant field to be used, and the type of index from the possible options - flat, quantizedFlat, or diskANN.
688688
A sample indexing policy with vector indexes would look like this:
689689
```python
690690
indexing_policy = {
@@ -703,10 +703,28 @@ indexing_policy = {
703703
],
704704
"vectorIndexes": [
705705
{"path": "/vector1", "type": "flat"},
706-
{"path": "/vector2", "type": "quantizedFlat"}
706+
{"path": "/vector2", "type": "quantizedFlat"},
707+
{"path": "/vector3", "type": "diskANN"}
707708
]
708709
}
709710
```
711+
712+
For vector index types of diskANN and quantizedFlat, there are additional options available as well. These are:
713+
714+
quantizationByteSize - the number of bytes used in product quantization of the vectors. A larger value may result in better recall for vector searches at the expense of latency. This applies to index types diskANN and quantizedFlat. The allowed range is between 1 and the minimum between 512 and the vector dimensions. The default value is 64.
715+
716+
indexingSearchListSize - which represents the size of the candidate list of approximate neighbors stored while building the diskANN index as part of the optimization processes. This applies only to index type diskANN. The allowed range is between 25 and 500.
717+
```python
718+
indexing_policy = {
719+
"automatic": True,
720+
"indexingMode": "consistent",
721+
"vectorIndexes": [
722+
{"path": "/vector1", "type": "quantizedFlat", "quantizationByteSize": 8},
723+
{"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 50}
724+
]
725+
}
726+
```
727+
710728
You would then pass in the relevant policies to your container creation method to ensure these configurations are used by it.
711729
The operation will fail if you pass new vector indexes to your indexing policy but forget to pass in an embedding policy.
712730
```python

‎sdk/cosmos/azure-cosmos/samples/index_management.py

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -637,7 +637,7 @@ def use_geospatial_indexing_policy(db):
637637
try:
638638
delete_container_if_exists(db, CONTAINER_ID)
639639

640-
# Create a container with vector embedding policy and vector indexes
640+
# Create a container with geospatial indexes
641641
indexing_policy = {
642642
'includedPaths': [
643643
{'path': '/"Location"/?',
@@ -688,16 +688,23 @@ def use_vector_embedding_policy(db):
688688
# Create a container with vector embedding policy and vector indexes
689689
indexing_policy = {
690690
"vectorIndexes": [
691-
{"path": "/embeddings", "type": "quantizedFlat"},
691+
{"path": "/vector", "type": "quantizedFlat", "quantizationByteSize": 8},
692+
{"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 50}
692693
]
693694
}
694695
vector_embedding_policy = {
695696
"vectorEmbeddings": [
696697
{
697-
"path": "/embeddings",
698+
"path": "/vector",
698699
"dataType": "float32",
699-
"dimensions": 1000,
700-
"distanceFunction": "cosine"
700+
"dimensions": 256,
701+
"distanceFunction": "euclidean"
702+
},
703+
{
704+
"path": "/vector2",
705+
"dataType": "int8",
706+
"dimensions": 200,
707+
"distanceFunction": "dotproduct"
701708
}
702709
]
703710
}
@@ -751,7 +758,7 @@ def run_sample():
751758
client = obtain_client()
752759
fetch_all_databases(client)
753760

754-
# Create database if doesn't exist already.
761+
# Create database if it doesn't exist already.
755762
created_db = create_database_if_not_exists(client, DATABASE_ID)
756763
print(created_db)
757764

‎sdk/cosmos/azure-cosmos/samples/index_management_async.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -632,7 +632,7 @@ async def use_geospatial_indexing_policy(db):
632632
try:
633633
await delete_container_if_exists(db, CONTAINER_ID)
634634

635-
# Create a container with vector embedding policy and vector indexes
635+
# Create a container with geospatial indexes
636636
indexing_policy = {
637637
'includedPaths': [
638638
{'path': '/"Location"/?',
@@ -683,16 +683,23 @@ async def use_vector_embedding_policy(db):
683683
# Create a container with vector embedding policy and vector indexes
684684
indexing_policy = {
685685
"vectorIndexes": [
686-
{"path": "/vector", "type": "quantizedFlat"},
686+
{"path": "/vector", "type": "quantizedFlat", "quantizationByteSize": 8},
687+
{"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 50}
687688
]
688689
}
689690
vector_embedding_policy = {
690691
"vectorEmbeddings": [
691692
{
692693
"path": "/vector",
693694
"dataType": "float32",
694-
"dimensions": 1000,
695-
"distanceFunction": "cosine"
695+
"dimensions": 256,
696+
"distanceFunction": "euclidean"
697+
},
698+
{
699+
"path": "/vector2",
700+
"dataType": "int8",
701+
"dimensions": 200,
702+
"distanceFunction": "dotproduct"
696703
}
697704
]
698705
}
@@ -746,7 +753,7 @@ async def run_sample():
746753
async with obtain_client() as client:
747754
await fetch_all_databases(client)
748755

749-
# Create database if doesn't exist already.
756+
# Create database if it doesn't exist already.
750757
created_db = await client.create_database_if_not_exists(DATABASE_ID)
751758
print(created_db)
752759

‎sdk/cosmos/azure-cosmos/test/test_vector_policy.py

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ def test_create_vector_embedding_container(self):
3434
indexing_policy = {
3535
"vectorIndexes": [
3636
{"path": "/vector1", "type": "flat"},
37-
{"path": "/vector2", "type": "quantizedFlat"},
38-
{"path": "/vector3", "type": "flat"}
37+
{"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 8},
38+
{"path": "/vector3", "type": "diskANN", "quantizationByteSize": 8, "indexingSearchListSize": 50}
3939
]
4040
}
4141
vector_embedding_policy = {
@@ -69,6 +69,7 @@ def test_create_vector_embedding_container(self):
6969
)
7070
properties = created_container.read()
7171
assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
72+
assert properties["indexingPolicy"]["vectorIndexes"] == indexing_policy["vectorIndexes"]
7273
self.test_db.delete_container(container_id)
7374

7475
def test_fail_create_vector_indexing_policy(self):
@@ -131,6 +132,42 @@ def test_fail_create_vector_indexing_policy(self):
131132
assert e.status_code == 400
132133
assert "vector2 not matching in Embedding's path" in e.http_error_message
133134

135+
# Pass a vector indexing policy with wrong quantizationByteSize value
136+
indexing_policy = {
137+
"vectorIndexes": [
138+
{"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 0}]
139+
}
140+
try:
141+
self.test_db.create_container(
142+
id='vector_container',
143+
partition_key=PartitionKey(path="/id"),
144+
indexing_policy=indexing_policy,
145+
vector_embedding_policy=vector_embedding_policy
146+
)
147+
pytest.fail("Container creation should have failed for value mismatch.")
148+
except exceptions.CosmosHttpResponseError as e:
149+
assert e.status_code == 400
150+
assert "QuantizationByteSize value :: 0 is out of range. The allowed range is between 1 and 256."\
151+
in e.http_error_message
152+
153+
# Pass a vector indexing policy with wrong indexingSearchListSize value
154+
indexing_policy = {
155+
"vectorIndexes": [
156+
{"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 5}]
157+
}
158+
try:
159+
self.test_db.create_container(
160+
id='vector_container',
161+
partition_key=PartitionKey(path="/id"),
162+
indexing_policy=indexing_policy,
163+
vector_embedding_policy=vector_embedding_policy
164+
)
165+
pytest.fail("Container creation should have failed for value mismatch.")
166+
except exceptions.CosmosHttpResponseError as e:
167+
assert e.status_code == 400
168+
assert "IndexingSearchListSize value :: 5 is out of range. The allowed range is between 25 and 500."\
169+
in e.http_error_message
170+
134171
def test_fail_replace_vector_indexing_policy(self):
135172
vector_embedding_policy = {
136173
"vectorEmbeddings": [

‎sdk/cosmos/azure-cosmos/test/test_vector_policy_async.py

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,8 @@ async def test_create_vector_embedding_container_async(self):
4444
indexing_policy = {
4545
"vectorIndexes": [
4646
{"path": "/vector1", "type": "flat"},
47-
{"path": "/vector2", "type": "quantizedFlat"},
48-
{"path": "/vector3", "type": "flat"}
47+
{"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 8},
48+
{"path": "/vector3", "type": "diskANN", "quantizationByteSize": 8, "indexingSearchListSize": 50}
4949
]
5050
}
5151
vector_embedding_policy = {
@@ -79,6 +79,7 @@ async def test_create_vector_embedding_container_async(self):
7979
)
8080
properties = await created_container.read()
8181
assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
82+
assert properties["indexingPolicy"]["vectorIndexes"] == indexing_policy["vectorIndexes"]
8283
await self.test_db.delete_container(container_id)
8384

8485
async def test_fail_create_vector_indexing_policy_async(self):
@@ -141,6 +142,42 @@ async def test_fail_create_vector_indexing_policy_async(self):
141142
assert e.status_code == 400
142143
assert "vector2 not matching in Embedding's path" in e.http_error_message
143144

145+
# Pass a vector indexing policy with wrong quantizationByteSize value
146+
indexing_policy = {
147+
"vectorIndexes": [
148+
{"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 0}]
149+
}
150+
try:
151+
await self.test_db.create_container(
152+
id='vector_container',
153+
partition_key=PartitionKey(path="/id"),
154+
indexing_policy=indexing_policy,
155+
vector_embedding_policy=vector_embedding_policy
156+
)
157+
pytest.fail("Container creation should have failed for value mismatch.")
158+
except exceptions.CosmosHttpResponseError as e:
159+
assert e.status_code == 400
160+
assert "QuantizationByteSize value :: 0 is out of range. The allowed range is between 1 and 256." \
161+
in e.http_error_message
162+
163+
# Pass a vector indexing policy with wrong indexingSearchListSize value
164+
indexing_policy = {
165+
"vectorIndexes": [
166+
{"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 5}]
167+
}
168+
try:
169+
await self.test_db.create_container(
170+
id='vector_container',
171+
partition_key=PartitionKey(path="/id"),
172+
indexing_policy=indexing_policy,
173+
vector_embedding_policy=vector_embedding_policy
174+
)
175+
pytest.fail("Container creation should have failed for value mismatch.")
176+
except exceptions.CosmosHttpResponseError as e:
177+
assert e.status_code == 400
178+
assert "IndexingSearchListSize value :: 5 is out of range. The allowed range is between 25 and 500." \
179+
in e.http_error_message
180+
144181
async def test_fail_replace_vector_indexing_policy_async(self):
145182
vector_embedding_policy = {
146183
"vectorEmbeddings": [

0 commit comments

Comments
 (0)