VectorIndexDefinition: Adds Support for Partitioned DiskANN

Microsoft.Azure.Cosmos/src/Fluent/Settings/VectorIndexDefinition.cs

@@ -55,6 +55,51 @@ public VectorIndexDefinition<T> Path(
             return this;
         }
 
+        /// <summary>
+        /// Configures the quantization byte size for the current <see cref="VectorIndexPath"/> definition.
+        /// </summary>
+        /// <param name="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. Note that, the allowed range for this parameter is between 1 and 3.
+        /// </param>
+        /// <returns>An instance of the current <see cref="VectorIndexDefinition{T}"/>.</returns>
+        public VectorIndexDefinition<T> WithQuantizationByteSize(
+            int quantizationByteSize)
+        {
+            this.vectorIndexPath.QuantizationByteSize = quantizationByteSize;
+            return this;
+        }
+
+        /// <summary>
+        /// Configures the indexing search list size for the current <see cref="VectorIndexPath"/> definition.
+        /// </summary>
+        /// <param name="indexingSearchListSize">
+        /// This represents 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 for this parameter is between 25 and 500.
+        /// </param>
+        /// <returns>An instance of the current <see cref="VectorIndexDefinition{T}"/>.</returns>
+        public VectorIndexDefinition<T> WithIndexingSearchListSize(
+            int indexingSearchListSize)
+        {
+            this.vectorIndexPath.IndexingSearchListSize = indexingSearchListSize;
+            return this;
+        }
+
+        /// <summary>
+        /// Configures the vector index shard key for the current <see cref="VectorIndexPath"/> definition.
+        /// </summary>
+        /// <param name="vectorIndexShardKey">
+        /// A string array containing the shard keys used for partitioning the vector indexes. This is an optional parameter and
+        /// applies to index types DiskANN and quantizedFlat.
+        /// </param>
+        /// <returns>An instance of the current <see cref="VectorIndexDefinition{T}"/>.</returns>
+        public VectorIndexDefinition<T> WithVectorIndexShardKey(
+            string[] vectorIndexShardKey)
+        {
+            this.vectorIndexPath.VectorIndexShardKey = vectorIndexShardKey ?? throw new ArgumentNullException(nameof(vectorIndexShardKey));
+            return this;
+        }
+
         /// <summary>
         /// Applies the current definition to the parent.
         /// </summary>

Microsoft.Azure.Cosmos/src/Resource/Settings/Embedding.cs

@@ -35,10 +35,10 @@ class Embedding : IEquatable<Embedding>
         public VectorDataType DataType { get; set; }
 
         /// <summary>
-        /// Gets or sets a long integer representing the dimensions of a vector. 
+        /// Gets or sets an integer representing the dimensions of a vector.
         /// </summary>
         [JsonProperty(PropertyName = "dimensions")]
-        public ulong Dimensions { get; set; }
+        public int Dimensions { get; set; }
 
         /// <summary>
         /// Gets or sets the <see cref="Cosmos.DistanceFunction"/> which is used to calculate the respective distance between the vectors. 

Microsoft.Azure.Cosmos/src/Resource/Settings/VectorDataType.cs

@@ -15,12 +15,6 @@ namespace Microsoft.Azure.Cosmos
 #endif
     enum VectorDataType
     {
-        /// <summary>
-        /// Represent a float16 data type.
-        /// </summary>
-        [EnumMember(Value = "float16")]
-        Float16,
-
         /// <summary>
         /// Represent a float32 data type.
         /// </summary>

Microsoft.Azure.Cosmos/src/Resource/Settings/VectorIndexPath.cs

@@ -32,11 +32,16 @@ namespace Microsoft.Azure.Cosmos
     ///         },
     ///         {
     ///             "path": "/vector2",
-    ///             "type": "flat"
+    ///             "type": "quantizedFlat",
+    ///             "quantizationByteSize": 3,
+    ///             "vectorIndexShardKey": ["/ZipCode"]
     ///         },
     ///         {
     ///             "path": "/embeddings/vector",
-    ///             "type": "flat"
+    ///             "type": "DiskANN",
+    ///             "quantizationByteSize": 2,
+    ///             "indexingSearchListSize": 100,
+    ///             "vectorIndexShardKey": ["/Country"]
     ///         }
     ///     ]
     /// }
@@ -49,6 +54,12 @@ namespace Microsoft.Azure.Cosmos
 #endif
     sealed class VectorIndexPath
     {
+        [JsonProperty(PropertyName = "indexingSearchListSize", NullValueHandling = NullValueHandling.Ignore)]
+        private int? indexingSearchListSizeInternal;
+
+        [JsonProperty(PropertyName = "quantizationByteSize", NullValueHandling = NullValueHandling.Ignore)]
+        private int? quantizationByteSizeInternal;
+
         /// <summary>
         /// Gets or sets the full path in a document used for vector indexing.
         /// </summary>
@@ -62,6 +73,34 @@ sealed class VectorIndexPath
         [JsonConverter(typeof(StringEnumConverter))]
         public VectorIndexType Type { get; set; }
 
+        /// <summary>
+        /// Gets or sets the quantization byte size for the vector index path. This is only applicable for the quantizedFlat and diskann vector index types.
+        /// The allowed range for this parameter is between 1 and 3.
+        /// </summary>
+        [JsonIgnore]
+        public int QuantizationByteSize
+        {
+            get => this.quantizationByteSizeInternal == null ? 0 : this.quantizationByteSizeInternal.Value;
+            set => this.quantizationByteSizeInternal = value;
+        }
+
+        /// <summary>
+        /// Gets or sets the indexing search list size for the vector index path. This is only applicable for the diskann vector index type.
+        /// The allowed range for this parameter is between 25 and 500.
+        /// </summary>
+        [JsonIgnore]
+        public int IndexingSearchListSize
+        {
+            get => this.indexingSearchListSizeInternal == null ? 0 : this.indexingSearchListSizeInternal.Value;
+            set => this.indexingSearchListSizeInternal = value;
+        }
+
+        /// <summary>
+        /// Gets or sets the vector index shard key for the vector index path. This is only applicable for the quantizedFlat and diskann vector index types.
+        /// </summary>
+        [JsonProperty(PropertyName = "vectorIndexShardKey", NullValueHandling = NullValueHandling.Ignore)]
+        public string[] VectorIndexShardKey { get; set; }
+
         /// <summary>
         /// This contains additional values for scenarios where the SDK is not aware of new fields. 
         /// This ensures that if resource is read and updated none of the fields will be lost in the process.

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Fluent/ContainerSettingsTests.cs

@@ -585,10 +585,15 @@ await databaseForVectorEmbedding.DefineContainer(containerName, partitionKeyPath
                                 .Path(vector1Path, VectorIndexType.Flat)
                              .Attach()
                             .WithVectorIndex()
-                                .Path(vector2Path, VectorIndexType.Flat)
+                                .Path(vector2Path, VectorIndexType.QuantizedFlat)
+                                .WithQuantizationByteSize(3)
+                                .WithVectorIndexShardKey(new string[] { "/Country" })
                              .Attach()
                             .WithVectorIndex()
-                                .Path(vector3Path, VectorIndexType.Flat)
+                                .Path(vector3Path, VectorIndexType.DiskANN)
+                                .WithQuantizationByteSize(2)
+                                .WithIndexingSearchListSize(5)
+                                .WithVectorIndexShardKey(new string[] { "/ZipCode" })
                              .Attach()
                         .Attach()
                         .CreateAsync();
@@ -610,9 +615,14 @@ await databaseForVectorEmbedding.DefineContainer(containerName, partitionKeyPath
                 Assert.AreEqual(vector1Path, containerSettings.IndexingPolicy.VectorIndexes[0].Path);
                 Assert.AreEqual(VectorIndexType.Flat, containerSettings.IndexingPolicy.VectorIndexes[0].Type);
                 Assert.AreEqual(vector2Path, containerSettings.IndexingPolicy.VectorIndexes[1].Path);
-                Assert.AreEqual(VectorIndexType.Flat, containerSettings.IndexingPolicy.VectorIndexes[1].Type);
+                Assert.AreEqual(VectorIndexType.QuantizedFlat, containerSettings.IndexingPolicy.VectorIndexes[1].Type);
+                Assert.AreEqual(3, containerSettings.IndexingPolicy.VectorIndexes[1].QuantizationByteSize);
+                CollectionAssert.AreEqual(new string[] { "/Country" }, containerSettings.IndexingPolicy.VectorIndexes[1].VectorIndexShardKey);
                 Assert.AreEqual(vector3Path, containerSettings.IndexingPolicy.VectorIndexes[2].Path);
-                Assert.AreEqual(VectorIndexType.Flat, containerSettings.IndexingPolicy.VectorIndexes[2].Type);
+                Assert.AreEqual(VectorIndexType.DiskANN, containerSettings.IndexingPolicy.VectorIndexes[2].Type);
+                Assert.AreEqual(2, containerSettings.IndexingPolicy.VectorIndexes[2].QuantizationByteSize);
+                Assert.AreEqual(5, containerSettings.IndexingPolicy.VectorIndexes[2].IndexingSearchListSize);
+                CollectionAssert.AreEqual(new string[] { "/ZipCode" }, containerSettings.IndexingPolicy.VectorIndexes[2].VectorIndexShardKey);
             }
             finally
             {