[Cosmos] DiskANN updates (Azure#38044)

* Update README.md

* add values to README, update tests

* Update README.md

* update samples

* Update sdk/cosmos/azure-cosmos/README.md

* remove vectorIndexShardKey

Author: simorenoh

sdk/cosmos/azure-cosmos/README.md

@@ -684,7 +684,7 @@ vector_embedding_policy = {
 ```
 
 Separately, vector indexes have been added to the already existing indexing_policy and only require two fields per index:
-the path to the relevant field to be used, and the type of index from the possible options (flat or quantizedFlat).
+the path to the relevant field to be used, and the type of index from the possible options - flat, quantizedFlat, or diskANN.
 A sample indexing policy with vector indexes would look like this:
 ```python
 indexing_policy = {
@@ -703,10 +703,28 @@ indexing_policy = {
         ],
         "vectorIndexes": [
             {"path": "/vector1", "type": "flat"},
-            {"path": "/vector2", "type": "quantizedFlat"}
+            {"path": "/vector2", "type": "quantizedFlat"},
+            {"path": "/vector3", "type": "diskANN"}
         ]
     }
 ```
+
+For vector index types of diskANN and quantizedFlat, there are additional options available as well. These are:
+
+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.
+
+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.
+```python
+indexing_policy = {
+        "automatic": True,
+        "indexingMode": "consistent",
+        "vectorIndexes": [
+            {"path": "/vector1", "type": "quantizedFlat", "quantizationByteSize": 8},
+            {"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 50}
+        ]
+    }
+```
+
 You would then pass in the relevant policies to your container creation method to ensure these configurations are used by it.
 The operation will fail if you pass new vector indexes to your indexing policy but forget to pass in an embedding policy.
 ```python

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

@@ -637,7 +637,7 @@ def use_geospatial_indexing_policy(db):
     try:
         delete_container_if_exists(db, CONTAINER_ID)
 
-        # Create a container with vector embedding policy and vector indexes
+        # Create a container with geospatial indexes
         indexing_policy = {
             'includedPaths': [
                 {'path': '/"Location"/?',
@@ -688,16 +688,23 @@ def use_vector_embedding_policy(db):
         # Create a container with vector embedding policy and vector indexes
         indexing_policy = {
             "vectorIndexes": [
-                {"path": "/embeddings", "type": "quantizedFlat"},
+                {"path": "/vector", "type": "quantizedFlat", "quantizationByteSize": 8},
+                {"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 50}
             ]
         }
         vector_embedding_policy = {
             "vectorEmbeddings": [
                 {
-                    "path": "/embeddings",
+                    "path": "/vector",
                     "dataType": "float32",
-                    "dimensions": 1000,
-                    "distanceFunction": "cosine"
+                    "dimensions": 256,
+                    "distanceFunction": "euclidean"
+                },
+                {
+                    "path": "/vector2",
+                    "dataType": "int8",
+                    "dimensions": 200,
+                    "distanceFunction": "dotproduct"
                 }
             ]
         }
@@ -751,7 +758,7 @@ def run_sample():
         client = obtain_client()
         fetch_all_databases(client)
 
-        # Create database if doesn't exist already.
+        # Create database if it doesn't exist already.
         created_db = create_database_if_not_exists(client, DATABASE_ID)
         print(created_db)
 

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

@@ -632,7 +632,7 @@ async def use_geospatial_indexing_policy(db):
     try:
         await delete_container_if_exists(db, CONTAINER_ID)
 
-        # Create a container with vector embedding policy and vector indexes
+        # Create a container with geospatial indexes
         indexing_policy = {
             'includedPaths': [
                 {'path': '/"Location"/?',
@@ -683,16 +683,23 @@ async def use_vector_embedding_policy(db):
         # Create a container with vector embedding policy and vector indexes
         indexing_policy = {
             "vectorIndexes": [
-                {"path": "/vector", "type": "quantizedFlat"},
+                {"path": "/vector", "type": "quantizedFlat", "quantizationByteSize": 8},
+                {"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 50}
             ]
         }
         vector_embedding_policy = {
             "vectorEmbeddings": [
                 {
                     "path": "/vector",
                     "dataType": "float32",
-                    "dimensions": 1000,
-                    "distanceFunction": "cosine"
+                    "dimensions": 256,
+                    "distanceFunction": "euclidean"
+                },
+                {
+                    "path": "/vector2",
+                    "dataType": "int8",
+                    "dimensions": 200,
+                    "distanceFunction": "dotproduct"
                 }
             ]
         }
@@ -746,7 +753,7 @@ async def run_sample():
         async with obtain_client() as client:
             await fetch_all_databases(client)
 
-            # Create database if doesn't exist already.
+            # Create database if it doesn't exist already.
             created_db = await client.create_database_if_not_exists(DATABASE_ID)
             print(created_db)
 

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

@@ -34,8 +34,8 @@ def test_create_vector_embedding_container(self):
         indexing_policy = {
             "vectorIndexes": [
                 {"path": "/vector1", "type": "flat"},
-                {"path": "/vector2", "type": "quantizedFlat"},
-                {"path": "/vector3", "type": "flat"}
+                {"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 8},
+                {"path": "/vector3", "type": "diskANN", "quantizationByteSize": 8, "indexingSearchListSize": 50}
             ]
         }
         vector_embedding_policy = {
@@ -69,6 +69,7 @@ def test_create_vector_embedding_container(self):
         )
         properties = created_container.read()
         assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
+        assert properties["indexingPolicy"]["vectorIndexes"] == indexing_policy["vectorIndexes"]
         self.test_db.delete_container(container_id)
 
     def test_fail_create_vector_indexing_policy(self):
@@ -131,6 +132,42 @@ def test_fail_create_vector_indexing_policy(self):
             assert e.status_code == 400
             assert "vector2 not matching in Embedding's path" in e.http_error_message
 
+        # Pass a vector indexing policy with wrong quantizationByteSize value
+        indexing_policy = {
+            "vectorIndexes": [
+                {"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 0}]
+        }
+        try:
+            self.test_db.create_container(
+                id='vector_container',
+                partition_key=PartitionKey(path="/id"),
+                indexing_policy=indexing_policy,
+                vector_embedding_policy=vector_embedding_policy
+            )
+            pytest.fail("Container creation should have failed for value mismatch.")
+        except exceptions.CosmosHttpResponseError as e:
+            assert e.status_code == 400
+            assert "QuantizationByteSize value :: 0 is out of range. The allowed range is between 1 and 256."\
+                   in e.http_error_message
+
+        # Pass a vector indexing policy with wrong indexingSearchListSize value
+        indexing_policy = {
+            "vectorIndexes": [
+                {"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 5}]
+        }
+        try:
+            self.test_db.create_container(
+                id='vector_container',
+                partition_key=PartitionKey(path="/id"),
+                indexing_policy=indexing_policy,
+                vector_embedding_policy=vector_embedding_policy
+            )
+            pytest.fail("Container creation should have failed for value mismatch.")
+        except exceptions.CosmosHttpResponseError as e:
+            assert e.status_code == 400
+            assert "IndexingSearchListSize value :: 5 is out of range. The allowed range is between 25 and 500."\
+                   in e.http_error_message
+
     def test_fail_replace_vector_indexing_policy(self):
         vector_embedding_policy = {
             "vectorEmbeddings": [

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

@@ -44,8 +44,8 @@ async def test_create_vector_embedding_container_async(self):
         indexing_policy = {
             "vectorIndexes": [
                 {"path": "/vector1", "type": "flat"},
-                {"path": "/vector2", "type": "quantizedFlat"},
-                {"path": "/vector3", "type": "flat"}
+                {"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 8},
+                {"path": "/vector3", "type": "diskANN", "quantizationByteSize": 8, "indexingSearchListSize": 50}
             ]
         }
         vector_embedding_policy = {
@@ -79,6 +79,7 @@ async def test_create_vector_embedding_container_async(self):
         )
         properties = await created_container.read()
         assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
+        assert properties["indexingPolicy"]["vectorIndexes"] == indexing_policy["vectorIndexes"]
         await self.test_db.delete_container(container_id)
 
     async def test_fail_create_vector_indexing_policy_async(self):
@@ -141,6 +142,42 @@ async def test_fail_create_vector_indexing_policy_async(self):
             assert e.status_code == 400
             assert "vector2 not matching in Embedding's path" in e.http_error_message
 
+        # Pass a vector indexing policy with wrong quantizationByteSize value
+        indexing_policy = {
+            "vectorIndexes": [
+                {"path": "/vector2", "type": "quantizedFlat", "quantizationByteSize": 0}]
+        }
+        try:
+            await self.test_db.create_container(
+                id='vector_container',
+                partition_key=PartitionKey(path="/id"),
+                indexing_policy=indexing_policy,
+                vector_embedding_policy=vector_embedding_policy
+            )
+            pytest.fail("Container creation should have failed for value mismatch.")
+        except exceptions.CosmosHttpResponseError as e:
+            assert e.status_code == 400
+            assert "QuantizationByteSize value :: 0 is out of range. The allowed range is between 1 and 256." \
+                   in e.http_error_message
+
+        # Pass a vector indexing policy with wrong indexingSearchListSize value
+        indexing_policy = {
+            "vectorIndexes": [
+                {"path": "/vector2", "type": "diskANN", "indexingSearchListSize": 5}]
+        }
+        try:
+            await self.test_db.create_container(
+                id='vector_container',
+                partition_key=PartitionKey(path="/id"),
+                indexing_policy=indexing_policy,
+                vector_embedding_policy=vector_embedding_policy
+            )
+            pytest.fail("Container creation should have failed for value mismatch.")
+        except exceptions.CosmosHttpResponseError as e:
+            assert e.status_code == 400
+            assert "IndexingSearchListSize value :: 5 is out of range. The allowed range is between 25 and 500." \
+                   in e.http_error_message
+
     async def test_fail_replace_vector_indexing_policy_async(self):
         vector_embedding_policy = {
             "vectorEmbeddings": [