[Cosmos] JS SDK support for partitioned DiskANN (#31839)

### Packages impacted by this PR
@azure/cosmos

### Describe the problem that is addressed by this PR
In this PR we are adding three optional parameters in the Indexing
Policy in the vector Index to support partitioned DiskANN. The three
optional parameters are:

quantizationByteSize : The number of bytes used in product quantization
of the vectors.This is an optional parameter and applies to index types
DiskANN and quantizedFlat. The allowed range for this parameter is
between 1 and 200.
vectorIndexShardKey : The list of string containing the shard keys used
for partitioning the vector indexes. This is an optional parameter and
applies to index types DiskANN and quantizedFlat.
indexingSearchListSize : The size of the candidate list of approximate
neighbors stored while building the diskANN index as part of the
optimization processes.This is an optional parameter and applies to
index type DiskANN only. The allowed range is between 25 and 500.
The changes have been tested using a Test account, since as of now the
backend changes aren't live yet.

### What are the possible designs available to address the problem? If
there are more than one possible design, why was the one in this PR
chosen?


### Are there test cases added in this PR? _(If not, why?)_
Yes

### Provide a list of related PRs _(if any)_


### Command used to generate this PR:**_(Applicable only to SDK release
request PRs)_

### Checklists
- [ ] Added impacted package name to the issue description
- [ ] Does this PR needs any fixes in the SDK Generator?** _(If so,
create an Issue in the
[Autorest/typescript](https://github.com/Azure/autorest.typescript)
repository and link it here)_
- [ ] Added a changelog (if necessary)

---------

Co-authored-by: Ujjwal Soni <[email protected]>
Co-authored-by: Manik Khandelwal <[email protected]>

sdk/cosmosdb/cosmos/CHANGELOG.md

@@ -6,6 +6,7 @@
 
 - Full Text Support: This feature adds support for full text search policy and indexing policy. It also enables performing full text search queries. [docs](https://learn.microsoft.com/azure/cosmos-db/gen-ai/full-text-search)
 - Hybrid Search Support: This feature adds support for performing hybrid search queries. [docs](https://learn.microsoft.com/azure/cosmos-db/gen-ai/hybrid-search)
+- Added support for three optional properties to support `quantizedFlat` and `diskANN` vector indexing policies. The properties are: `quantizationByteSize`, `vectorIndexShardKey` and `indexingSearchListSize`.
 
 ## 4.1.1 (2024-08-30)
 

sdk/cosmosdb/cosmos/review/cosmos.api.md

@@ -2598,8 +2598,11 @@ export interface VectorEmbeddingPolicy {
 
 // @public
 export interface VectorIndex {
+    indexingSearchListSize?: number;
     path: string;
+    quantizationByteSize?: number;
     type: VectorIndexType;
+    vectorIndexShardKey?: string[];
 }
 
 // @public

sdk/cosmosdb/cosmos/samples-dev/ContainerManagement.ts

@@ -161,8 +161,19 @@ async function run(): Promise<void> {
     spatialIndexes: [{ path: "/location/*", types: ["Point", "Polygon"] }] as SpatialIndex[],
     vectorIndexes: [
       { path: "/vector1", type: VectorIndexType.Flat },
-      { path: "/vector2", type: VectorIndexType.QuantizedFlat },
-      { path: "/vector3", type: VectorIndexType.DiskANN },
+      {
+        path: "/vector2",
+        type: VectorIndexType.QuantizedFlat,
+        quantizationByteSize: 2,
+        vectorIndexShardKey: ["/Country"],
+      },
+      {
+        path: "/vector3",
+        type: VectorIndexType.DiskANN,
+        quantizationByteSize: 2,
+        indexingSearchListSize: 50,
+        vectorIndexShardKey: ["/ZipCode"],
+      },
     ],
   };
 

sdk/cosmosdb/cosmos/samples/v4/javascript/ContainerManagement.js

@@ -157,8 +157,19 @@ async function run() {
     spatialIndexes: [{ path: "/location/*", types: ["Point", "Polygon"] }],
     vectorIndexes: [
       { path: "/vector1", type: VectorIndexType.Flat },
-      { path: "/vector2", type: VectorIndexType.QuantizedFlat },
-      { path: "/vector3", type: VectorIndexType.DiskANN },
+      {
+        path: "/vector2",
+        type: VectorIndexType.QuantizedFlat,
+        quantizationByteSize: 2,
+        vectorIndexShardKey: ["/Country"],
+      },
+      {
+        path: "/vector3",
+        type: VectorIndexType.DiskANN,
+        quantizationByteSize: 2,
+        indexingSearchListSize: 50,
+        vectorIndexShardKey: ["/ZipCode"],
+      },
     ],
   };
 

sdk/cosmosdb/cosmos/samples/v4/typescript/src/ContainerManagement.ts

@@ -161,8 +161,19 @@ async function run(): Promise<void> {
     spatialIndexes: [{ path: "/location/*", types: ["Point", "Polygon"] }] as SpatialIndex[],
     vectorIndexes: [
       { path: "/vector1", type: VectorIndexType.Flat },
-      { path: "/vector2", type: VectorIndexType.QuantizedFlat },
-      { path: "/vector3", type: VectorIndexType.DiskANN },
+      {
+        path: "/vector2",
+        type: VectorIndexType.QuantizedFlat,
+        quantizationByteSize: 2,
+        vectorIndexShardKey: ["/Country"],
+      },
+      {
+        path: "/vector3",
+        type: VectorIndexType.DiskANN,
+        quantizationByteSize: 2,
+        indexingSearchListSize: 50,
+        vectorIndexShardKey: ["/ZipCode"],
+      },
     ],
   };
 

sdk/cosmosdb/cosmos/src/documents/IndexingPolicy.ts

@@ -70,7 +70,26 @@ export interface VectorIndex {
    * Currently, flat, diskANN, and quantizedFlat are supported.
    */
   type: VectorIndexType;
+  /**
+   * The number of bytes used in product quantization of the vectors.
+   * This is an optional parameter and applies to index types DiskANN and quantizedFlat.
+   * The allowed range for this parameter is between 1 and min(Dimensions, 512).
+   */
+  quantizationByteSize?: number;
+  /**
+   * The list of string containing the shard keys used for partitioning the vector indexes.
+   * This is an optional parameter and applies to index types DiskANN and quantizedFlat.
+   */
+  vectorIndexShardKey?: string[];
+  /**
+   * The size of the candidate list of approximate neighbors stored while building
+   * the diskANN index as part of the optimization processes.
+   * This is an optional parameter and applies to index type DiskANN only.
+   * The allowed range is between 25 and 500.
+   */
+  indexingSearchListSize?: number;
 }
+
 /**
  * Represents the index type of the vector.
  */

sdk/cosmosdb/cosmos/test/public/functional/NonStreamingQueryPolicy.spec.ts

@@ -65,6 +65,79 @@ describe("Vector search feature", async () => {
       assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[2].path === "/vector3");
     });
 
+    // skipping the test case for now. Will enable it once the changes are live on backend
+    it.skip("validate VectorEmbeddingPolicy", async function () {
+      const indexingPolicy: IndexingPolicy = {
+        vectorIndexes: [
+          { path: "/vector1", type: VectorIndexType.Flat },
+          {
+            path: "/vector2",
+            type: VectorIndexType.QuantizedFlat,
+            quantizationByteSize: 1,
+            vectorIndexShardKey: ["/Country"],
+          },
+          {
+            path: "/vector3",
+            type: VectorIndexType.DiskANN,
+            quantizationByteSize: 2,
+            indexingSearchListSize: 100,
+            vectorIndexShardKey: ["/ZipCode"],
+          },
+        ],
+      };
+      const vectorEmbeddingPolicy: VectorEmbeddingPolicy = {
+        vectorEmbeddings: [
+          {
+            path: "/vector1",
+            dataType: VectorEmbeddingDataType.Float32,
+            dimensions: 500,
+            distanceFunction: VectorEmbeddingDistanceFunction.Euclidean,
+          },
+          {
+            path: "/vector2",
+            dataType: VectorEmbeddingDataType.Int8,
+            dimensions: 200,
+            distanceFunction: VectorEmbeddingDistanceFunction.DotProduct,
+          },
+          {
+            path: "/vector3",
+            dataType: VectorEmbeddingDataType.UInt8,
+            dimensions: 400,
+            distanceFunction: VectorEmbeddingDistanceFunction.Cosine,
+          },
+        ],
+      };
+      const containerName = "JSApp-vector embedding container";
+      // create container
+      const { resource: containerdef } = await database.containers.createIfNotExists({
+        id: containerName,
+        vectorEmbeddingPolicy: vectorEmbeddingPolicy,
+        indexingPolicy: indexingPolicy,
+      });
+
+      assert(containerdef.indexingPolicy !== undefined);
+      assert(containerdef.vectorEmbeddingPolicy !== undefined);
+      assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings.length === 3);
+      assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[0].path === "/vector1");
+      assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[1].path === "/vector2");
+      assert(containerdef.vectorEmbeddingPolicy.vectorEmbeddings[2].path === "/vector3");
+
+      assert(containerdef.indexingPolicy.vectorIndexes.length === 3);
+      assert(containerdef.indexingPolicy.vectorIndexes[0].path === "/vector1");
+      assert(containerdef.indexingPolicy.vectorIndexes[1].path === "/vector2");
+      assert(containerdef.indexingPolicy.vectorIndexes[2].path === "/vector3");
+
+      assert(containerdef.indexingPolicy.vectorIndexes[0].type === VectorIndexType.Flat);
+      assert(containerdef.indexingPolicy.vectorIndexes[1].type === VectorIndexType.QuantizedFlat);
+      assert(containerdef.indexingPolicy.vectorIndexes[2].type === VectorIndexType.DiskANN);
+
+      assert(containerdef.indexingPolicy.vectorIndexes[1].quantizationByteSize === 1);
+      assert(containerdef.indexingPolicy.vectorIndexes[2].quantizationByteSize === 2);
+      assert(containerdef.indexingPolicy.vectorIndexes[2].indexingSearchListSize === 100);
+      assert(containerdef.indexingPolicy.vectorIndexes[1].vectorIndexShardKey[0] === "/Country");
+      assert(containerdef.indexingPolicy.vectorIndexes[2].vectorIndexShardKey[0] === "/ZipCode");
+    });
+
     it("should fail to create vector indexing policy", async function () {
       const vectorEmbeddingPolicy: VectorEmbeddingPolicy = {
         vectorEmbeddings: [