Adding vector index shard key property, for DiskANN vector search (Azure#44007)

* Adding vector index shard key property

* Updating changelog

* Updating changelog

* Updating tests

* Resolving comments

* Resolving comments

* Updating tests

* fixing build

Author: aayush3011

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/VectorIndexTest.java

@@ -309,6 +309,7 @@ private void validateVectorIndexes(List<CosmosVectorIndexSpec> actual, List<Cosm
             if (Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.QUANTIZED_FLAT.toString()) ||
                 Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.DISK_ANN.toString())) {
                 assertThat(expected.get(i).getQuantizationSizeInBytes()).isEqualTo(actual.get(i).getQuantizationSizeInBytes());
+                assertThat(expected.get(i).getVectorIndexShardKeys()).isEqualTo(actual.get(i).getVectorIndexShardKeys());
             }
             if (Objects.equals(expected.get(i).getType(), CosmosVectorIndexType.DISK_ANN.toString())) {
                 assertThat(expected.get(i).getIndexingSearchListSize()).isEqualTo(actual.get(i).getIndexingSearchListSize());
@@ -326,12 +327,14 @@ private List<CosmosVectorIndexSpec> populateVectorIndexes() {
         cosmosVectorIndexSpec2.setPath("/vector2");
         cosmosVectorIndexSpec2.setType(CosmosVectorIndexType.QUANTIZED_FLAT.toString());
         cosmosVectorIndexSpec2.setQuantizationSizeInBytes(2);
+        cosmosVectorIndexSpec2.setVectorIndexShardKeys(Arrays.asList("/zipCode"));
 
         CosmosVectorIndexSpec cosmosVectorIndexSpec3 = new CosmosVectorIndexSpec();
         cosmosVectorIndexSpec3.setPath("/vector3");
         cosmosVectorIndexSpec3.setType(CosmosVectorIndexType.DISK_ANN.toString());
         cosmosVectorIndexSpec3.setQuantizationSizeInBytes(2);
         cosmosVectorIndexSpec3.setIndexingSearchListSize(30);
+        cosmosVectorIndexSpec3.setVectorIndexShardKeys(Arrays.asList("/country/city"));
 
         return Arrays.asList(cosmosVectorIndexSpec1, cosmosVectorIndexSpec2, cosmosVectorIndexSpec3);
     }

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Added the `vectorIndexShardKeys` to the vectorIndexSpec for QuantizedFlat and DiskANN vector search. - [PR 44007](https://github.com/Azure/azure-sdk-for-java/pull/44007)
 
 ### 4.68.0 (2025-03-20)
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java

@@ -153,6 +153,7 @@ public static final class Properties {
         public static final String DISTANCE_FUNCTION = "distanceFunction";
         public static final String VECTOR_QUANTIZATION_SIZE_IN_BYTES = "quantizationByteSize";
         public static final String VECTOR_INDEXING_SEARCH_LIST_SIZE =  "indexingSearchListSize";
+        public static final String VECTOR_INDEX_SHARD_KEYS = "vectorIndexShardKeys";
 
         // Unique index.
         public static final String UNIQUE_KEY_POLICY = "uniqueKeyPolicy";

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/IndexProperty.java

@@ -4,5 +4,6 @@
 
 public enum IndexProperty {
     INDEXING_SEARCH_LIST_SIZE,
-    QUANTIZATION_SIZE_IN_BYTES;
+    QUANTIZATION_SIZE_IN_BYTES,
+    VECTOR_INDEX_SHARD_KEYS;
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosVectorIndexSpec.java

@@ -9,6 +9,8 @@
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import java.util.List;
+
 import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
 
 /**
@@ -22,6 +24,8 @@ public final class CosmosVectorIndexSpec {
     private Integer quantizationSizeInBytes;
     @JsonInclude(JsonInclude.Include.NON_NULL)
     private Integer indexingSearchListSize;
+    @JsonInclude(JsonInclude.Include.NON_NULL)
+    private List<String> vectorIndexShardKeys;
     private final JsonSerializable jsonSerializable;
 
     /**
@@ -147,6 +151,38 @@ public CosmosVectorIndexSpec setIndexingSearchListSize(Integer indexingSearchLis
         return this;
     }
 
+    /**
+     * Gets the vector indexing shard keys
+     *
+     * @return vectorIndexShardKeys the list of string containing the shard keys used for partitioning the
+     * vector indexes. This applies to index types diskANN and quantizedFlat. The maximum allowed size for
+     * this array is currently limited to 1 - that is, there is only one allowed path.
+     */
+    public List<String> getVectorIndexShardKeys() {
+        if (this.vectorIndexShardKeys == null) {
+            this.vectorIndexShardKeys = this.jsonSerializable.getList(Constants.Properties.VECTOR_INDEX_SHARD_KEYS, String.class);
+        }
+        return this.vectorIndexShardKeys;
+    }
+
+    /**
+     * Sets the vector indexing shard keys
+     *
+     * @param vectorIndexShardKeys vectorIndexShardKeys the list of string containing the shard keys used for partitioning the
+     *                            vector indexes. This applies to index types diskANN and quantizedFlat. The maximum allowed size for
+     *                            this array is currently limited to 1 - that is, there is only one allowed path.
+     * @return CosmosVectorIndexSpec
+     */
+    public CosmosVectorIndexSpec setVectorIndexShardKeys(List<String> vectorIndexShardKeys) {
+        if (validateIndexType(IndexProperty.VECTOR_INDEX_SHARD_KEYS) && vectorIndexShardKeys != null) {
+            this.vectorIndexShardKeys = vectorIndexShardKeys;
+            this.jsonSerializable.set(Constants.Properties.VECTOR_INDEX_SHARD_KEYS, this.vectorIndexShardKeys);
+        } else {
+            this.vectorIndexShardKeys = null;
+        }
+        return this;
+    }
+
     void populatePropertyBag() {
         this.jsonSerializable.populatePropertyBag();
     }
@@ -157,7 +193,7 @@ JsonSerializable getJsonSerializable() {
 
     private Boolean validateIndexType(IndexProperty indexProperty) {
         String vectorIndexType = this.jsonSerializable.getString(Constants.Properties.VECTOR_INDEX_TYPE);
-        if (indexProperty.equals(IndexProperty.QUANTIZATION_SIZE_IN_BYTES)) {
+        if (indexProperty.equals(IndexProperty.QUANTIZATION_SIZE_IN_BYTES) || (indexProperty.equals(IndexProperty.VECTOR_INDEX_SHARD_KEYS))) {
             return vectorIndexType.equals(CosmosVectorIndexType.QUANTIZED_FLAT.toString()) ||
                 vectorIndexType.equals(CosmosVectorIndexType.DISK_ANN.toString());
         }