Support create GPU index (#551)

Signed-off-by: shanghaikid <jiangruiyi@gmail.com>

client/src/consts/Milvus.ts

@@ -121,7 +121,12 @@ export type searchKeywordsType =
   | 'radius'
   | 'range_filter'
   | 'drop_ratio_search'
-  | 'filter';
+  | 'filter'
+  | 'itopk_size'
+  | 'search_width'
+  | 'min_iterations'
+  | 'max_iterations'
+  | 'team_size';
 
 export type indexConfigType = {
   [x: string]: {
@@ -132,9 +137,13 @@ export type indexConfigType = {
 
 // index
 export const FLOAT_INDEX_CONFIG: indexConfigType = {
-  SCANN: {
-    create: ['nlist', 'with_raw_data'],
-    search: ['nprobe'],
+  AUTOINDEX: {
+    create: [],
+    search: ['level'],
+  },
+  HNSW: {
+    create: ['M', 'efConstruction'],
+    search: ['ef'],
   },
   IVF_FLAT: {
     create: ['nlist'],
@@ -152,22 +161,50 @@ export const FLOAT_INDEX_CONFIG: indexConfigType = {
     create: [],
     search: ['nprobe'],
   },
-  HNSW: {
-    create: ['M', 'efConstruction'],
-    search: ['ef'],
-  },
-  AUTOINDEX: {
-    create: [],
-    search: ['level'],
+  SCANN: {
+    create: ['nlist', 'with_raw_data'],
+    search: ['nprobe'],
   },
+};
+
+export const DISK_INDEX_CONFIG: indexConfigType = {
   DISKANN: {
     create: [],
     search: ['search_list'],
   },
 };
 
+export const GPU_INDEX_CONFIG: indexConfigType = {
+  GPU_CAGRA: {
+    create: [
+      'intermediate_graph_degree',
+      'graph_degree',
+      'build_algo',
+      'cache_dataset_on_device',
+    ],
+    search: [
+      'itopk_size',
+      'search_width',
+      'min_iterations',
+      'max_iterations',
+      'team_size',
+    ],
+  },
+  GPU_IVF_FLAT: {
+    create: ['nlist'],
+    search: ['nprobe'],
+  },
+  GPU_IVF_PQ: {
+    create: ['nlist', 'm', 'nbits'],
+    search: ['nprobe'],
+  },
+  GPU_BRUTE_FORCE: {
+    create: [],
+    search: [],
+  },
+};
+
 export const BINARY_INDEX_CONFIG: indexConfigType = {
-  // },
   BIN_FLAT: {
     create: [],
     search: [],
@@ -193,6 +230,8 @@ export const INDEX_CONFIG: indexConfigType = {
   ...FLOAT_INDEX_CONFIG,
   ...BINARY_INDEX_CONFIG,
   ...SPARSE_INDEX_CONFIG,
+  ...DISK_INDEX_CONFIG,
+  ...GPU_INDEX_CONFIG,
 };
 
 export const COLLECTION_NAME_REGX = /^[0-9,a-z,A-Z$_]+$/;
@@ -224,6 +263,14 @@ export const INDEX_OPTIONS_MAP = {
       value: INDEX_TYPES_ENUM.MARISA_TRIE,
     },
   ],
+  ['DISK']: Object.keys(DISK_INDEX_CONFIG).map(v => ({
+    label: v,
+    value: v,
+  })),
+  ['GPU']: Object.keys(GPU_INDEX_CONFIG).map(v => ({
+    label: v,
+    value: v,
+  })),
 };
 
 // search params default value map

client/src/i18n/cn/index.ts

@@ -1,6 +1,10 @@
 const indexTrans = {
   type: '索引类型',
   param: '索引参数',
+  inMemory: '内存索引',
+  disk: '磁盘索引',
+  gpu: 'GPU索引',
+  scalar: '标量索引',
 
   create: '创建索引',
   index: '索引',

client/src/i18n/en/index.ts

@@ -1,6 +1,10 @@
 const indexTrans = {
   type: 'Index Type',
   param: 'Index Parameters',
+  inMemory: 'In-Memory Index',
+  disk: 'Disk Index',
+  gpu: 'GPU Index',
+  scalar: 'Scalar Index',
 
   create: 'Create Index',
   index: 'Index',

client/src/pages/databases/collections/overview/CreateForm.tsx

@@ -4,7 +4,8 @@ import { useTranslation } from 'react-i18next';
 import { ITextfieldConfig } from '@/components/customInput/Types';
 import CustomInput from '@/components/customInput/CustomInput';
 import CustomSelector from '@/components/customSelector/CustomSelector';
-import { Option } from '@/components/customSelector/Types';
+import CustomGroupedSelect from '@/components/customSelector/CustomGroupedSelect';
+import { Option, GroupOption } from '@/components/customSelector/Types';
 import { FormHelperType } from '../../../../types/Common';
 
 const useStyles = makeStyles((theme: Theme) => ({
@@ -26,7 +27,7 @@ const useStyles = makeStyles((theme: Theme) => ({
 const CreateForm = (
   props: FormHelperType & {
     metricOptions: Option[];
-    indexOptions: Option[];
+    indexOptions: GroupOption[];
     indexParams: string[];
     indexTypeChange?: (type: string) => void;
   }
@@ -76,6 +77,16 @@ const CreateForm = (
       };
 
       if (type === 'number') {
+        config.validations!.push({
+          rule: 'number',
+          errorText: warningTrans('number', { name: label }),
+        });
+      }
+      if (
+        type === 'number' &&
+        typeof min === 'number' &&
+        typeof max === 'number'
+      ) {
         config.validations!.push({
           rule: 'range',
           errorText: warningTrans('range', { min, max }),
@@ -157,6 +168,24 @@ const CreateForm = (
         key: 'with_raw_data',
         type: 'bool',
       }),
+      intermediate_graph_degree: generateConfig({
+        label: 'intermediate_graph_degree',
+        key: 'intermediate_graph_degree',
+      }),
+      graph_degree: generateConfig({
+        label: 'graph_degree',
+        key: 'graph_degree',
+      }),
+      build_algo: generateConfig({
+        label: 'build_algo',
+        key: 'build_algo',
+        type: 'text',
+      }),
+      cache_dataset_on_device: generateConfig({
+        label: 'cache_dataset_on_device',
+        key: 'cache_dataset_on_device',
+        type: 'bool',
+      }),
     };
 
     const result: ITextfieldConfig[] = [];
@@ -182,10 +211,10 @@ const CreateForm = (
 
   return (
     <div className={`${classes.wrapper} ${wrapperClass}`}>
-      <CustomSelector
+      <CustomGroupedSelect
         label={indexTrans('type')}
-        value={formValue.index_type}
         options={indexOptions}
+        value={formValue.index_type}
         onChange={(e: { target: { value: unknown } }) => {
           const type = e.target.value;
           updateForm('index_type', type as string);
@@ -193,8 +222,7 @@ const CreateForm = (
           updateForm('metric_type', metricOptions[0].value as string);
           indexTypeChange && indexTypeChange(type as string);
         }}
-        variant="filled"
-        wrapperClass={classes.select}
+        className={classes.select}
       />
       <CustomInput
         type="text"

client/src/pages/databases/collections/overview/Create.tsx → client/src/pages/databases/collections/overview/CreateIndexDialog.tsx

@@ -1,8 +1,9 @@
-import { useEffect, useMemo, useState } from 'react';
+import { useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { CodeLanguageEnum, CodeViewData } from '@/components/code/Types';
 import DialogTemplate from '@/components/customDialog/DialogTemplate';
 import CustomSwitch from '@/components/customSwitch/CustomSwitch';
+import { Option } from '@/components/customSelector/Types';
 import {
   INDEX_CONFIG,
   INDEX_OPTIONS_MAP,
@@ -103,6 +104,10 @@ const CreateIndex = (props: {
     knng: '',
     drop_ratio_build: '0.5',
     with_raw_data: 'true',
+    intermediate_graph_degree: '128',
+    graph_degree: '64',
+    build_algo: 'IVF_PQ',
+    cache_dataset_on_device: 'false',
   });
 
   // control whether show code mode
@@ -138,26 +143,53 @@ const CreateIndex = (props: {
     return extraParams;
   }, [indexCreateParams, indexSetting]);
 
+  const getOptions = (label: string, children: Option[]) => [
+    { label, children },
+  ];
+
   const indexOptions = useMemo(() => {
-    switch (fieldType) {
-      case DataTypeStringEnum.BinaryVector:
-        return INDEX_OPTIONS_MAP[DataTypeEnum.BinaryVector];
-      case DataTypeStringEnum.FloatVector:
-      case DataTypeStringEnum.Float16Vector:
-      case DataTypeStringEnum.BFloat16Vector:
-        return INDEX_OPTIONS_MAP[DataTypeEnum.FloatVector];
-      case DataTypeStringEnum.SparseFloatVector:
-        return INDEX_OPTIONS_MAP[DataTypeEnum.SparseFloatVector];
-      case DataTypeStringEnum.VarChar:
-        return INDEX_OPTIONS_MAP[DataTypeEnum.VarChar];
+    if (VectorTypes.includes(dataType)) {
+      switch (fieldType) {
+        case DataTypeStringEnum.BinaryVector:
+          return [
+            ...getOptions(
+              indexTrans('inMemory'),
+              INDEX_OPTIONS_MAP[DataTypeEnum.BinaryVector]
+            ),
+          ];
+        case DataTypeStringEnum.SparseFloatVector:
+          return [
+            ...getOptions(
+              indexTrans('inMemory'),
+              INDEX_OPTIONS_MAP[DataTypeEnum.SparseFloatVector]
+            ),
+          ];
 
-      default:
-        return [
-          { label: 'INVERTED', value: INDEX_TYPES_ENUM.INVERTED },
-          { label: 'STL sort', value: INDEX_TYPES_ENUM.SORT },
-        ];
+        default:
+          return [
+            ...getOptions(
+              indexTrans('inMemory'),
+              INDEX_OPTIONS_MAP[DataTypeEnum.FloatVector]
+            ),
+            ...getOptions(indexTrans('disk'), INDEX_OPTIONS_MAP['DISK']),
+            ...getOptions(indexTrans('gpu'), INDEX_OPTIONS_MAP['GPU']),
+          ];
+      }
+    } else {
+      switch (fieldType) {
+        case DataTypeStringEnum.VarChar:
+          return getOptions(
+            indexTrans('scalar'),
+            INDEX_OPTIONS_MAP[DataTypeEnum.VarChar]
+          );
+        default:
+          return getOptions(indexTrans('scalar'), [
+            { label: 'INVERTED', value: INDEX_TYPES_ENUM.INVERTED },
+            { label: 'STL sort', value: INDEX_TYPES_ENUM.SORT },
+          ]);
+      }
     }
-  }, [fieldType]);
+  }, [fieldType, dataType, fieldName]);
 
   const checkedForm = useMemo(() => {
     if (!VectorTypes.includes(dataType)) {
@@ -241,8 +273,8 @@ const CreateIndex = (props: {
     }, 0);
   };
 
-  const handleCreateIndex = () => {
-    handleCreate(extraParams, indexSetting.index_name);
+  const handleCreateIndex = async () => {
+    await handleCreate(extraParams, indexSetting.index_name);
   };
 
   const handleShowCode = (event: React.ChangeEvent<{ checked: boolean }>) => {

client/src/pages/databases/collections/overview/IndexTypeElement.tsx

@@ -13,7 +13,7 @@ import {
   DataTypeStringEnum,
   DataTypeEnum,
 } from '@/consts';
-import CreateIndex from './Create';
+import CreateIndexDialog from './CreateIndexDialog';
 import { FieldObject } from '@server/types';
 import CustomButton from '@/components/customButton/CustomButton';
 
@@ -106,7 +106,7 @@ const IndexTypeElement: FC<{
       type: 'custom',
       params: {
         component: (
-          <CreateIndex
+          <CreateIndexDialog
             collectionName={collectionName}
             fieldName={field.name}
             dataType={field.dataType as unknown as DataTypeEnum}

client/src/pages/search/SearchParams.tsx

@@ -305,6 +305,63 @@ const SearchParams: FC<SearchParamsProps> = ({
             handleInputChange('drop_ratio_search', value);
           },
         },
+        itopk_size: {
+          label: 'itopk_size',
+          key: 'itopk_size',
+          value: searchParamsForm['itopk_size'] ?? '',
+          isInt: true,
+          type: 'number',
+          required: false,
+          handleChange: value => {
+            handleInputChange('itopk_size', value);
+          },
+        },
+        search_width: {
+          label: 'search_width',
+          key: 'search_width',
+          value: searchParamsForm['search_width'] ?? '',
+          isInt: true,
+          type: 'number',
+          required: false,
+          handleChange: value => {
+            handleInputChange('search_width', value);
+          },
+        },
+        min_iterations: {
+          label: 'min_iterations',
+          key: 'min_iterations',
+          value: searchParamsForm['min_iterations'] ?? '0',
+          isInt: true,
+          type: 'number',
+          required: false,
+          handleChange: value => {
+            handleInputChange('min_iterations', value);
+          },
+        },
+        max_iterations: {
+          label: 'max_iterations',
+          key: 'max_iterations',
+          value: searchParamsForm['max_iterations'] ?? '0',
+          isInt: true,
+          type: 'number',
+          required: false,
+          handleChange: value => {
+            handleInputChange('max_iterations', value);
+          },
+        },
+        team_size: {
+          label: 'team_size',
+          key: 'team_size',
+          value: searchParamsForm['team_size'] ?? '0',
+          min: 2,
+          max: 32,
+          isInt: true,
+          type: 'number',
+          required: false,
+          handleChange: value => {
+            handleInputChange('team_size', value);
+          },
+        },
       };
 
       const param = configParamMap[paramKey];