From 10cda214013fcea51ad030f6b8d0d5a563fdb0b7 Mon Sep 17 00:00:00 2001
From: asier-isayas <asier.isayas@gmail.com>
Date: Fri, 2 May 2025 11:05:40 -0400
Subject: [PATCH] Add Vector Index Shard Key option on container creation
 (#2097)

* Add vector index shard key

* npm run format

* rename shard key to vector index shard key

* add tooltip for quantization byte size

* change text for GSI and container in VectorEmbedding Policy

---------

Co-authored-by: Asier Isayas <aisayas@microsoft.com>
---
 src/Contracts/DataModels.ts                   |  2 +-
 .../Controls/Settings/SettingsComponent.tsx   |  1 +
 .../ContainerPolicyComponent.tsx              |  1 +
 .../VectorEmbeddingPoliciesComponent.tsx      | 66 ++++++++-----------
 .../AddGlobalSecondaryIndexPanel.tsx          |  1 +
 .../Components/VectorSearchComponent.tsx      |  3 +
 6 files changed, 36 insertions(+), 38 deletions(-)

diff --git a/src/Contracts/DataModels.ts b/src/Contracts/DataModels.ts
index 58e412b76..7bf489fd6 100644
--- a/src/Contracts/DataModels.ts
+++ b/src/Contracts/DataModels.ts
@@ -210,7 +210,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/Settings/SettingsComponent.tsx b/src/Explorer/Controls/Settings/SettingsComponent.tsx
index 684858c56..65000e7ee 100644
--- a/src/Explorer/Controls/Settings/SettingsComponent.tsx
+++ b/src/Explorer/Controls/Settings/SettingsComponent.tsx
@@ -1215,6 +1215,7 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       isFullTextSearchEnabled: this.isFullTextSearchEnabled,
       shouldDiscardContainerPolicies: this.state.shouldDiscardContainerPolicies,
       resetShouldDiscardContainerPolicyChange: this.resetShouldDiscardContainerPolicies,
+      isGlobalSecondaryIndex: this.isGlobalSecondaryIndex,
     };
 
     const indexingPolicyComponentProps: IndexingPolicyComponentProps = {
diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/ContainerPolicyComponent.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/ContainerPolicyComponent.tsx
index f4db62ec6..1257d4a65 100644
--- a/src/Explorer/Controls/Settings/SettingsSubComponents/ContainerPolicyComponent.tsx
+++ b/src/Explorer/Controls/Settings/SettingsSubComponents/ContainerPolicyComponent.tsx
@@ -22,6 +22,7 @@ export interface ContainerPolicyComponentProps {
   isFullTextSearchEnabled: boolean;
   shouldDiscardContainerPolicies: boolean;
   resetShouldDiscardContainerPolicyChange: () => void;
+  isGlobalSecondaryIndex?: boolean;
 }
 
 export const ContainerPolicyComponent: React.FC<ContainerPolicyComponentProps> = ({
diff --git a/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx b/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx
index 03d1ab1e6..ba0c5b62b 100644
--- a/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx
+++ b/src/Explorer/Controls/VectorSearch/VectorEmbeddingPoliciesComponent.tsx
@@ -9,6 +9,7 @@ import {
   Stack,
   TextField,
 } from "@fluentui/react";
+import { InfoTooltip } from "Common/Tooltip/InfoTooltip";
 import { VectorEmbedding, VectorIndex } from "Contracts/DataModels";
 import { CollapsibleSectionComponent } from "Explorer/Controls/CollapsiblePanel/CollapsibleSectionComponent";
 import {
@@ -29,6 +30,7 @@ export interface IVectorEmbeddingPoliciesComponentProps {
   discardChanges?: boolean;
   onChangesDiscarded?: () => void;
   disabled?: boolean;
+  isGlobalSecondaryIndex?: boolean;
 }
 
 export interface VectorEmbeddingPolicyData {
@@ -39,8 +41,7 @@ export interface VectorEmbeddingPolicyData {
   indexType: VectorIndex["type"] | "none";
   pathError: string;
   dimensionsError: string;
-  diskANNShardKey?: string;
-  diskANNShardKeyError?: string;
+  vectorIndexShardKey?: string[];
   indexingSearchListSize?: number;
   indexingSearchListSizeError?: string;
   quantizationByteSize?: number;
@@ -87,6 +88,7 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
   discardChanges,
   onChangesDiscarded,
   disabled,
+  isGlobalSecondaryIndex,
 }): JSX.Element => {
   const onVectorEmbeddingPathError = (path: string, index?: number): string => {
     let error = "";
@@ -132,12 +134,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 +143,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 +183,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 +245,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,
@@ -292,6 +286,11 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
     setVectorEmbeddingPolicyData(vectorEmbeddings);
   };
 
+  const getQuantizationByteSizeTooltipContent = (): string => {
+    const containerName: string = isGlobalSecondaryIndex ? "global secondary index" : "container";
+    return `This is dynamically set by the ${containerName} if left blank, or it can be set to a fixed number`;
+  };
+
   return (
     <Stack tokens={{ childrenGap: 4 }}>
       {vectorEmbeddingPolicyData &&
@@ -402,6 +401,7 @@ export const VectorEmbeddingPoliciesComponent: FunctionComponent<IVectorEmbeddin
                         styles={labelStyles}
                       >
                         Quantization byte size
+                        <InfoTooltip>{getQuantizationByteSizeTooltipContent()}</InfoTooltip>
                       </Label>
                       <TextField
                         disabled={
@@ -431,26 +431,18 @@ 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>
+                    <Stack style={{ marginLeft: "10px" }}>
+                      <Label disabled={disabled || vectorEmbeddingPolicy.indexType !== "diskANN"} styles={labelStyles}>
+                        Vector index 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 || "")}
-                        onChange={(event: React.ChangeEvent<HTMLInputElement>) =>
-                          onDiskANNShardKeyChange(index, event)
-                        }
+                        value={String(vectorEmbeddingPolicy.vectorIndexShardKey?.[0] ?? "")}
+                        onChange={(event: React.ChangeEvent<HTMLInputElement>) => onShardKeyChange(index, event)}
                       />
                     </Stack>
-                    */}
                   </Stack>
                 )}
               </Stack>
diff --git a/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/AddGlobalSecondaryIndexPanel.tsx b/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/AddGlobalSecondaryIndexPanel.tsx
index ba49299fb..f3d3a3341 100644
--- a/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/AddGlobalSecondaryIndexPanel.tsx
+++ b/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/AddGlobalSecondaryIndexPanel.tsx
@@ -388,6 +388,7 @@ export const AddGlobalSecondaryIndexPanel = (props: AddGlobalSecondaryIndexPanel
                 setVectorIndexingPolicy,
                 vectorPolicyValidated,
                 setVectorPolicyValidated,
+                isGlobalSecondaryIndex: true,
               }}
             />
           )}
diff --git a/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/Components/VectorSearchComponent.tsx b/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/Components/VectorSearchComponent.tsx
index 440d51ea4..8b5ad9a28 100644
--- a/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/Components/VectorSearchComponent.tsx
+++ b/src/Explorer/Panes/AddGlobalSecondaryIndexPanel/Components/VectorSearchComponent.tsx
@@ -14,6 +14,7 @@ export interface VectorSearchComponentProps {
   vectorIndexingPolicy: VectorIndex[];
   setVectorIndexingPolicy: React.Dispatch<React.SetStateAction<VectorIndex[]>>;
   setVectorPolicyValidated: React.Dispatch<React.SetStateAction<boolean>>;
+  isGlobalSecondaryIndex?: boolean;
 }
 
 export const VectorSearchComponent = (props: VectorSearchComponentProps): JSX.Element => {
@@ -23,6 +24,7 @@ export const VectorSearchComponent = (props: VectorSearchComponentProps): JSX.El
     vectorIndexingPolicy,
     setVectorIndexingPolicy,
     setVectorPolicyValidated,
+    isGlobalSecondaryIndex,
   } = props;
 
   return (
@@ -49,6 +51,7 @@ export const VectorSearchComponent = (props: VectorSearchComponentProps): JSX.El
                 setVectorIndexingPolicy(vectorIndexingPolicy);
                 setVectorPolicyValidated(vectorPolicyValidated);
               }}
+              isGlobalSecondaryIndex={isGlobalSecondaryIndex}
             />
           </Stack>
         </Stack>