VectorSearch.tsx 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426
  1. import { TextField, Typography } from '@material-ui/core';
  2. import { useTranslation } from 'react-i18next';
  3. import { useNavigationHook } from '../../hooks/Navigation';
  4. import { ALL_ROUTER_TYPES } from '../../router/Types';
  5. import CustomSelector from '../../components/customSelector/CustomSelector';
  6. import { useCallback, useEffect, useMemo, useState } from 'react';
  7. import SearchParams from './SearchParams';
  8. import { DEFAULT_METRIC_VALUE_MAP } from '../../consts/Milvus';
  9. import { FieldOption, SearchResultView, VectorSearchParam } from './Types';
  10. import MilvusGrid from '../../components/grid/Grid';
  11. import EmptyCard from '../../components/cards/EmptyCard';
  12. import icons from '../../components/icons/Icons';
  13. import { usePaginationHook } from '../../hooks/Pagination';
  14. import CustomButton from '../../components/customButton/CustomButton';
  15. import SimpleMenu from '../../components/menu/SimpleMenu';
  16. import { TOP_K_OPTIONS } from './Constants';
  17. import { Option } from '../../components/customSelector/Types';
  18. import { CollectionHttp } from '../../http/Collection';
  19. import { CollectionData, DataTypeEnum } from '../collections/Types';
  20. import { IndexHttp } from '../../http/Index';
  21. import { getVectorSearchStyles } from './Styles';
  22. import { parseValue } from '../../utils/Insert';
  23. import {
  24. classifyFields,
  25. getDefaultIndexType,
  26. getEmbeddingType,
  27. getNonVectorFieldsForFilter,
  28. getVectorFieldOptions,
  29. transferSearchResult,
  30. } from '../../utils/search';
  31. import { ColDefinitionsType } from '../../components/grid/Types';
  32. import Filter from '../../components/advancedSearch';
  33. import { Field } from '../../components/advancedSearch/Types';
  34. const VectorSearch = () => {
  35. useNavigationHook(ALL_ROUTER_TYPES.SEARCH);
  36. const { t: searchTrans } = useTranslation('search');
  37. const { t: btnTrans } = useTranslation('btn');
  38. const classes = getVectorSearchStyles();
  39. // data stored inside the component
  40. const [tableLoading, setTableLoading] = useState<boolean>(false);
  41. const [collections, setCollections] = useState<CollectionData[]>([]);
  42. const [selectedCollection, setSelectedCollection] = useState<string>('');
  43. const [fieldOptions, setFieldOptions] = useState<FieldOption[]>([]);
  44. // fields for advanced filter
  45. const [filterFields, setFilterFields] = useState<Field[]>([]);
  46. const [selectedField, setSelectedField] = useState<string>('');
  47. // search params form
  48. const [searchParam, setSearchParam] = useState<{ [key in string]: number }>(
  49. {}
  50. );
  51. // search params disable state
  52. const [paramDisabled, setParamDisabled] = useState<boolean>(true);
  53. // use null as init value before search, empty array means no results
  54. const [searchResult, setSearchResult] = useState<SearchResultView[] | null>(
  55. null
  56. );
  57. // default topK is 100
  58. const [topK, setTopK] = useState<number>(100);
  59. const [expression, setExpression] = useState<string>('');
  60. const [vectors, setVectors] = useState<string>('');
  61. const {
  62. pageSize,
  63. handlePageSize,
  64. currentPage,
  65. handleCurrentPage,
  66. total,
  67. data: result,
  68. } = usePaginationHook(searchResult || []);
  69. const searchDisabled = useMemo(() => {
  70. /**
  71. * before search, user must:
  72. * 1. enter vector value
  73. * 2. choose collection and field
  74. * 3. set extra search params
  75. */
  76. const isInvalid =
  77. vectors === '' ||
  78. selectedCollection === '' ||
  79. selectedField === '' ||
  80. paramDisabled;
  81. return isInvalid;
  82. }, [paramDisabled, selectedField, selectedCollection, vectors]);
  83. const collectionOptions: Option[] = useMemo(
  84. () =>
  85. collections.map(c => ({
  86. label: c._name,
  87. value: c._name,
  88. })),
  89. [collections]
  90. );
  91. const outputFields: string[] = useMemo(() => {
  92. const fields =
  93. collections.find(c => c._name === selectedCollection)?._fields || [];
  94. // vector field can't be output fields
  95. const invalidTypes = ['BinaryVector', 'FloatVector'];
  96. const nonVectorFields = fields.filter(
  97. field => !invalidTypes.includes(field._fieldType)
  98. );
  99. return nonVectorFields.map(f => f._fieldName);
  100. }, [selectedCollection, collections]);
  101. const colDefinitions: ColDefinitionsType[] = useMemo(() => {
  102. // filter id and score
  103. return searchResult && searchResult.length > 0
  104. ? Object.keys(searchResult[0])
  105. .filter(item => item !== 'id' && item !== 'score')
  106. .map(key => ({
  107. id: key,
  108. align: 'left',
  109. disablePadding: false,
  110. label: key,
  111. }))
  112. : [];
  113. }, [searchResult]);
  114. const { metricType, indexType, indexParams, fieldType, embeddingType } =
  115. useMemo(() => {
  116. if (selectedField !== '') {
  117. // field options must contain selected field, so selectedFieldInfo will never undefined
  118. const selectedFieldInfo = fieldOptions.find(
  119. f => f.value === selectedField
  120. );
  121. const index = selectedFieldInfo?.indexInfo;
  122. const embeddingType = getEmbeddingType(selectedFieldInfo!.fieldType);
  123. const metric =
  124. index?._metricType || DEFAULT_METRIC_VALUE_MAP[embeddingType];
  125. const indexParams = index?._indexParameterPairs || [];
  126. return {
  127. metricType: metric,
  128. indexType: index?._indexType || getDefaultIndexType(embeddingType),
  129. indexParams,
  130. fieldType: DataTypeEnum[selectedFieldInfo?.fieldType!],
  131. embeddingType,
  132. };
  133. }
  134. return {
  135. metricType: '',
  136. indexType: '',
  137. indexParams: [],
  138. fieldType: 0,
  139. embeddingType: DataTypeEnum.FloatVector,
  140. };
  141. }, [selectedField, fieldOptions]);
  142. // fetch data
  143. const fetchCollections = useCallback(async () => {
  144. const collections = await CollectionHttp.getCollections();
  145. setCollections(collections);
  146. }, []);
  147. const fetchFieldsWithIndex = useCallback(
  148. async (collectionName: string, collections: CollectionData[]) => {
  149. const fields =
  150. collections.find(c => c._name === collectionName)?._fields || [];
  151. const indexes = await IndexHttp.getIndexInfo(collectionName);
  152. const { vectorFields, nonVectorFields } = classifyFields(fields);
  153. // only vector type fields can be select
  154. const fieldOptions = getVectorFieldOptions(vectorFields, indexes);
  155. setFieldOptions(fieldOptions);
  156. // only non vector type fields can be advanced filter
  157. const filterFields = getNonVectorFieldsForFilter(nonVectorFields);
  158. setFilterFields(filterFields);
  159. },
  160. []
  161. );
  162. useEffect(() => {
  163. fetchCollections();
  164. }, [fetchCollections]);
  165. // get field options with index when selected collection changed
  166. useEffect(() => {
  167. if (selectedCollection !== '') {
  168. fetchFieldsWithIndex(selectedCollection, collections);
  169. }
  170. }, [selectedCollection, collections, fetchFieldsWithIndex]);
  171. // icons
  172. const VectorSearchIcon = icons.vectorSearch;
  173. const ResetIcon = icons.refresh;
  174. const ArrowIcon = icons.dropdown;
  175. // methods
  176. const handlePageChange = (e: any, page: number) => {
  177. handleCurrentPage(page);
  178. };
  179. const handleReset = () => {
  180. /**
  181. * reset search includes:
  182. * 1. reset vectors
  183. * 2. reset selected collection and field
  184. * 3. reset search params
  185. * 4. reset advanced filter expression
  186. * 5. clear search result
  187. */
  188. setVectors('');
  189. setSelectedField('');
  190. setSelectedCollection('');
  191. setSearchResult(null);
  192. setFilterFields([]);
  193. setExpression('');
  194. };
  195. const handleSearch = async (topK: number, expr = expression) => {
  196. const searhParamPairs = [
  197. // dynamic search params
  198. {
  199. key: 'params',
  200. value: JSON.stringify(searchParam),
  201. },
  202. {
  203. key: 'anns_field',
  204. value: selectedField,
  205. },
  206. {
  207. key: 'topk',
  208. value: topK,
  209. },
  210. {
  211. key: 'metric_type',
  212. value: metricType,
  213. },
  214. ];
  215. const params: VectorSearchParam = {
  216. output_fields: outputFields,
  217. expr,
  218. search_params: searhParamPairs,
  219. vectors: [parseValue(vectors)],
  220. vector_type: fieldType,
  221. };
  222. setTableLoading(true);
  223. try {
  224. const res = await CollectionHttp.vectorSearchData(
  225. selectedCollection,
  226. params
  227. );
  228. setTableLoading(false);
  229. const result = transferSearchResult(res.results);
  230. setSearchResult(result);
  231. } catch (err) {
  232. setTableLoading(false);
  233. }
  234. };
  235. const handleAdvancedFilterChange = (expression: string) => {
  236. setExpression(expression);
  237. if (!searchDisabled) {
  238. handleSearch(topK, expression);
  239. }
  240. };
  241. const handleVectorChange = (value: string) => {
  242. setVectors(value);
  243. };
  244. return (
  245. <section className="page-wrapper">
  246. {/* form section */}
  247. <form className={classes.form}>
  248. {/* vector value textarea */}
  249. <fieldset className="field">
  250. <Typography className="text">{searchTrans('firstTip')}</Typography>
  251. <TextField
  252. className="textarea"
  253. InputProps={{
  254. classes: {
  255. root: 'textfield',
  256. multiline: 'multiline',
  257. },
  258. }}
  259. multiline
  260. rows={5}
  261. placeholder={searchTrans('vectorPlaceholder')}
  262. value={vectors}
  263. onChange={(e: React.ChangeEvent<{ value: unknown }>) => {
  264. handleVectorChange(e.target.value as string);
  265. }}
  266. />
  267. </fieldset>
  268. {/* collection and field selectors */}
  269. <fieldset className="field field-second">
  270. <Typography className="text">{searchTrans('secondTip')}</Typography>
  271. <CustomSelector
  272. options={collectionOptions}
  273. wrapperClass={classes.selector}
  274. variant="filled"
  275. label={searchTrans(
  276. collectionOptions.length === 0 ? 'noCollection' : 'collection'
  277. )}
  278. disabled={collectionOptions.length === 0}
  279. value={selectedCollection}
  280. onChange={(e: { target: { value: unknown } }) => {
  281. const collection = e.target.value;
  282. setSelectedCollection(collection as string);
  283. // every time selected collection changed, reset field
  284. setSelectedField('');
  285. }}
  286. />
  287. <CustomSelector
  288. options={fieldOptions}
  289. // readOnly can't avoid all events, so we use disabled instead
  290. disabled={selectedCollection === ''}
  291. wrapperClass={classes.selector}
  292. variant="filled"
  293. label={searchTrans('field')}
  294. value={selectedField}
  295. onChange={(e: { target: { value: unknown } }) => {
  296. const field = e.target.value;
  297. setSelectedField(field as string);
  298. }}
  299. />
  300. </fieldset>
  301. {/* search params selectors */}
  302. <fieldset className="field field-params">
  303. <Typography className="text">{searchTrans('thirdTip')}</Typography>
  304. <SearchParams
  305. wrapperClass={classes.paramsWrapper}
  306. metricType={metricType!}
  307. embeddingType={
  308. embeddingType as
  309. | DataTypeEnum.BinaryVector
  310. | DataTypeEnum.FloatVector
  311. }
  312. indexType={indexType}
  313. indexParams={indexParams!}
  314. searchParamsForm={searchParam}
  315. handleFormChange={setSearchParam}
  316. topK={topK}
  317. setParamsDisabled={setParamDisabled}
  318. />
  319. </fieldset>
  320. </form>
  321. {/**
  322. * search toolbar section
  323. * including topK selector, advanced filter, search and reset btn
  324. */}
  325. <section className={classes.toolbar}>
  326. <div className="left">
  327. <Typography variant="h5" className="text">
  328. {`${searchTrans('result')}: `}
  329. </Typography>
  330. {/* topK selector */}
  331. <SimpleMenu
  332. label={searchTrans('topK', { number: topK })}
  333. menuItems={TOP_K_OPTIONS.map(item => ({
  334. label: item.toString(),
  335. callback: () => {
  336. setTopK(item);
  337. if (!searchDisabled) {
  338. handleSearch(item);
  339. }
  340. },
  341. wrapperClass: classes.menuItem,
  342. }))}
  343. buttonProps={{
  344. className: classes.menuLabel,
  345. endIcon: <ArrowIcon />,
  346. }}
  347. menuItemWidth="108px"
  348. />
  349. <Filter
  350. title="Advanced Filter"
  351. fields={filterFields}
  352. filterDisabled={selectedField === '' || selectedCollection === ''}
  353. onSubmit={handleAdvancedFilterChange}
  354. />
  355. </div>
  356. <div className="right">
  357. <CustomButton className="btn" onClick={handleReset}>
  358. <ResetIcon classes={{ root: 'icon' }} />
  359. {btnTrans('reset')}
  360. </CustomButton>
  361. <CustomButton
  362. variant="contained"
  363. disabled={searchDisabled}
  364. onClick={() => handleSearch(topK)}
  365. >
  366. {btnTrans('search')}
  367. </CustomButton>
  368. </div>
  369. </section>
  370. {/* search result table section */}
  371. {(searchResult && searchResult.length > 0) || tableLoading ? (
  372. <MilvusGrid
  373. toolbarConfigs={[]}
  374. colDefinitions={colDefinitions}
  375. rows={result}
  376. rowCount={total}
  377. primaryKey="rank"
  378. page={currentPage}
  379. onChangePage={handlePageChange}
  380. rowsPerPage={pageSize}
  381. setRowsPerPage={handlePageSize}
  382. openCheckBox={false}
  383. isLoading={tableLoading}
  384. />
  385. ) : (
  386. <EmptyCard
  387. wrapperClass={`page-empty-card`}
  388. icon={<VectorSearchIcon />}
  389. text={
  390. searchResult !== null
  391. ? searchTrans('empty')
  392. : searchTrans('startTip')
  393. }
  394. />
  395. )}
  396. </section>
  397. );
  398. };
  399. export default VectorSearch;