From f57691a229ff10adada6f6a59c47c6574f284825 Mon Sep 17 00:00:00 2001
From: Asier Isayas <aisayas@microsoft.com>
Date: Mon, 7 Apr 2025 11:12:34 -0400
Subject: [PATCH] Add vector index shard key

---
 src/Contracts/DataModels.ts                   |  2 +-
 .../VectorEmbeddingPoliciesComponent.tsx      | 48 +++++++------------
 2 files changed, 19 insertions(+), 31 deletions(-)

diff --git a/src/Contracts/DataModels.ts b/src/Contracts/DataModels.ts
index 3b3ab5027..08d0ac3f8 100644
--- a/src/Contracts/DataModels.ts
+++ b/src/Contracts/DataModels.ts
@@ -207,7 +207,7 @@ export interface IndexingPolicy {
 export interface VectorIndex {
   path: string;
   type: "flat" | "diskANN" | "quantizedFlat";
-  diskANNShardKey?: string;
+  vectorIndexShardKey?: string[];
   indexingSearchListSize?: number;
   quantizationByteSize?: number;
 }
diff --git a/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx b/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx
index 03d1ab1e6..b391ce498 100644
--- a/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx
+++ b/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx
@@ -39,8 +39,7 @@ export interface VectorEmbeddingPolicyData {
   indexType: VectorIndex["type"] | "none";
   pathError: string;
   dimensionsError: string;
-  diskANNShardKey?: string;
-  diskANNShardKeyError?: string;
+  vectorIndexShardKey?: string[];
   indexingSearchListSize?: number;
   indexingSearchListSizeError?: string;
   quantizationByteSize?: number;
@@ -132,12 +131,6 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
     return error;
   };
 
-  //TODO: no restrictions yet due to this field being removed for now.
-  // Uncomment and replace with validation code when field is reinstated
-  // const onDiskANNShardKeyError = (shardKey: string): string => {
-  //   return "";
-  // };
-
   const initializeData = (vectorEmbeddings: VectorEmbedding[], vectorIndexes: VectorIndex[]) => {
     const mergedData: VectorEmbeddingPolicyData[] = [];
     vectorEmbeddings.forEach((embedding) => {
@@ -147,6 +140,7 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
         indexType: matchingIndex?.type || "none",
         indexingSearchListSize: matchingIndex?.indexingSearchListSize || undefined,
         quantizationByteSize: matchingIndex?.quantizationByteSize || undefined,
+        vectorIndexShardKey: matchingIndex?.vectorIndexShardKey || undefined,
         pathError: onVectorEmbeddingPathError(embedding.path),
         dimensionsError: onVectorEmbeddingDimensionError(embedding.dimensions, matchingIndex?.type || "none"),
       });
@@ -186,6 +180,7 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
             type: policy.indexType,
             indexingSearchListSize: policy.indexingSearchListSize,
             quantizationByteSize: policy.quantizationByteSize,
+            vectorIndexShardKey: policy.vectorIndexShardKey
           }) as VectorIndex,
       );
     const validationPassed = vectorEmbeddingPolicyData.every(
@@ -247,20 +242,16 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
     setVectorEmbeddingPolicyData(vectorEmbeddings);
   };
 
-  // TODO: uncomment after Ignite
-  // DiskANNShardKey was removed for Ignite due to backend problems. Leaving this here as it will be reinstated immediately after Ignite
-  // const onDiskANNShardKeyChange = (index: number, event: React.ChangeEvent<HTMLInputElement>) => {
-  //   const value = event.target.value.trim();
-  //   const vectorEmbeddings = [...vectorEmbeddingPolicyData];
-  //   if (!vectorEmbeddings[index]?.diskANNShardKey && !value.startsWith("/")) {
-  //     vectorEmbeddings[index].diskANNShardKey = "/" + value;
-  //   } else {
-  //     vectorEmbeddings[index].diskANNShardKey = value;
-  //   }
-  //   const error = onDiskANNShardKeyError(value);
-  //   vectorEmbeddings[index].diskANNShardKeyError = error;
-  //   setVectorEmbeddingPolicyData(vectorEmbeddings);
-  // }
+  const onShardKeyChange = (index: number, event: React.ChangeEvent<HTMLInputElement>) => {
+    const value = event.target.value.trim();
+    const vectorEmbeddings = [...vectorEmbeddingPolicyData];
+    if (!vectorEmbeddings[index]?.vectorIndexShardKey?.[0] && !value.startsWith("/")) {
+      vectorEmbeddings[index].vectorIndexShardKey = ["/" + value];
+    } else {
+      vectorEmbeddings[index].vectorIndexShardKey = [value];
+    }
+    setVectorEmbeddingPolicyData(vectorEmbeddings);
+  }
 
   const onVectorEmbeddingPolicyChange = (
     index: number,
@@ -431,26 +422,23 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
                         }
                       />
                     </Stack>
-                    {/*TODO: uncomment after Ignite */}
-                    {/* DiskANNShardKey was removed for Ignite due to backend problems. Leaving this here as it will be reinstated immediately after Ignite
                     <Stack
                       style={{ marginLeft: "10px" }}
                     >
                       <Label
                         disabled={disabled || vectorEmbeddingPolicy.indexType !== "diskANN"}
                         styles={labelStyles}
-                      >DiskANN shard key</Label>
+                      >Shard key</Label>
                       <TextField
                         disabled={disabled || vectorEmbeddingPolicy.indexType !== "diskANN"}
-                        id={`vector-policy-diskANNShardKey-${index + 1}`}
+                        id={`vector-policy-vectorIndexShardKey-${index + 1}`}
                         styles={textFieldStyles}
-                        value={String(vectorEmbeddingPolicy.diskANNShardKey || "")}
+                        value={String(vectorEmbeddingPolicy.vectorIndexShardKey?.[0] ?? "")}
                         onChange={(event: React.ChangeEvent<HTMLInputElement>) =>
-                          onDiskANNShardKeyChange(index, event)
+                          onShardKeyChange(index, event)
                         }
                       />
-                    </Stack>
-                    */}
+                    </Stack>                  
                   </Stack>
                 )}
               </Stack>