Support filter partitions in search page (#680)

* add parition selectors

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

* finish

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

---------

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

client/src/i18n/cn/search.ts

@@ -28,7 +28,9 @@ const searchTrans = {
   outputFields: '输出字段',
   consistency: '一致性',
   graphNodeHoverTip: '双击以查看更多',
-  inputVectorTip: '向量或实体ID',
+  inputVectorPlaceHolder: '向量或实体ID',
+  partitionFilter: '分区过滤',
+  loading: '加载中...',
 };
 
 export default searchTrans;

client/src/i18n/en/search.ts

@@ -28,7 +28,9 @@ const searchTrans = {
   outputFields: 'Outputs',
   consistency: 'Consistency',
   graphNodeHoverTip: 'Double click to explore more',
-  inputVectorTip: 'Vector or entity id',
+  inputVectorPlaceHolder: 'Vector or entity id',
+  partitionFilter: 'Partition Filter',
+  loading: 'Loading...',
 };
 
 export default searchTrans;

client/src/pages/databases/Databases.tsx

@@ -167,6 +167,7 @@ const Databases = () => {
             ...prevParams,
             {
               collection: c,
+              partitions: [],
               searchParams: c.schema.vectorFields.map(v => {
                 return {
                   anns_field: v.name,

client/src/pages/databases/collections/search/PartitionsSelector.tsx

@@ -0,0 +1,92 @@
+import { useState } from 'react';
+import Autocomplete from '@mui/material/Autocomplete';
+import CircularProgress from '@mui/material/CircularProgress';
+import { PartitionService } from '@/http';
+import { PartitionData } from '@server/types';
+import CustomInput from '@/components/customInput/CustomInput';
+import { useTranslation } from 'react-i18next';
+
+interface PartitionsSelectorProps {
+  collectionName: string;
+  selected: PartitionData[];
+  setSelected: (value: PartitionData[]) => void;
+}
+
+export default function PartitionsSelector(props: PartitionsSelectorProps) {
+  // i18n
+  const { t: searchTrans } = useTranslation('search');
+  // default loading
+  const DEFAULT_LOADING_OPTIONS: readonly PartitionData[] = [
+    { name: searchTrans('loading'), id: -1, rowCount: -1, createdTime: '' },
+  ];
+
+  // props
+  const { collectionName, selected, setSelected } = props;
+  // state
+  const [open, setOpen] = useState(false);
+  const [options, setOptions] = useState<readonly PartitionData[]>(
+    DEFAULT_LOADING_OPTIONS
+  );
+  const [loading, setLoading] = useState(false);
+
+  const handleOpen = () => {
+    setOpen(true);
+    (async () => {
+      try {
+        const res = await PartitionService.getPartitions(collectionName);
+        setLoading(false);
+        setOptions([...res]);
+      } catch (err) {
+        setLoading(false);
+      }
+    })();
+  };
+
+  const handleClose = () => {
+    setOpen(false);
+    setOptions(DEFAULT_LOADING_OPTIONS);
+  };
+
+  return (
+    <Autocomplete
+      open={open}
+      multiple
+      limitTags={2}
+      color="primary"
+      disableCloseOnSelect
+      onOpen={handleOpen}
+      onClose={handleClose}
+      onChange={(_, value) => {
+        setSelected(value);
+      }}
+      value={selected}
+      isOptionEqualToValue={(option, value) =>
+        option && value && option.name === value.name
+      }
+      getOptionLabel={option => (option && option.name) || ''}
+      options={options}
+      loading={loading}
+      renderInput={params => {
+        return loading ? (
+          <CircularProgress color="inherit" size={20} />
+        ) : (
+          <CustomInput
+            textConfig={{
+              ...params,
+              label: searchTrans('partitionFilter'),
+              key: 'partitionFilter',
+              className: 'input',
+              value: params.inputProps.value,
+              disabled: false,
+              variant: 'filled',
+              required: false,
+              InputLabelProps: { shrink: true },
+            }}
+            checkValid={() => true}
+            type="text"
+          />
+        );
+      }}
+    />
+  );
+}

client/src/pages/databases/collections/search/Search.tsx

@@ -22,6 +22,7 @@ import VectorInputBox from './VectorInputBox';
 import StatusIcon, { LoadingType } from '@/components/status/StatusIcon';
 import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
 import CustomInput from '@/components/customInput/CustomInput';
+import PartitionsSelector from './PartitionsSelector';
 import {
   formatFieldType,
   cloneObj,
@@ -163,6 +164,16 @@ const Search = (props: CollectionDataProps) => {
     [JSON.stringify(searchParams)]
   );
 
+  // on partitions change
+  const onPartitionsChange = useCallback(
+    (value: any) => {
+      const s = cloneObj(searchParams) as SearchParamsType;
+      s.partitions = value;
+      setSearchParams({ ...s });
+    },
+    [JSON.stringify(searchParams)]
+  );
+
   // set search result
   const setSearchResult = useCallback(
     (props: { results: SearchResultView[]; latency: number }) => {
@@ -380,6 +391,9 @@ const Search = (props: CollectionDataProps) => {
     disableSearchTooltip = searchTrans('noVectorToSearch');
   }
 
+  // enable partition filter
+  const enablePartitionsFilter = !collection.schema.enablePartitionKey;
+
   return (
     <div className={classes.root}>
       {collection && (
@@ -456,6 +470,14 @@ const Search = (props: CollectionDataProps) => {
                 </Accordion>
               );
             })}
+
+            {enablePartitionsFilter && (
+              <PartitionsSelector
+                collectionName={collectionName}
+                selected={searchParams.partitions}
+                setSelected={onPartitionsChange}
+              />
+            )}
           </div>
 
           <div className={classes.searchControls}>

client/src/pages/databases/collections/search/VectorInputBox.tsx

@@ -1,4 +1,4 @@
-import { useRef, useEffect, useState } from 'react';
+import { useRef, useEffect } from 'react';
 import { useTranslation } from 'react-i18next';
 import { EditorState, Compartment } from '@codemirror/state';
 import { EditorView, keymap, placeholder } from '@codemirror/view';
@@ -121,9 +121,6 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
   const { searchParams, onChange, collection } = props;
   const { field, data } = searchParams;
 
-  // UI states
-  const [isFocused, setIsFocused] = useState(false);
-
   // classes
   const classes = getQueryStyles();
 
@@ -208,7 +205,7 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
         extensions: [
           minimalSetup,
           javascript(),
-          placeholder('vector or entity id'),
+          placeholder(searchTrans('inputVectorPlaceHolder')),
           linter(view => {
             const text = view.state.doc.toString();
 
@@ -275,7 +272,10 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
               }
             }
             if (update.focusChanged) {
-              setIsFocused(update.view.hasFocus);
+              editorEl.current?.classList.toggle(
+                'focused',
+                update.view.hasFocus
+              );
             }
           }),
         ],
@@ -315,10 +315,5 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
     }
   }, [theme.palette.mode]);
 
-  return (
-    <div
-      className={`${classes.vectorInputBox} ${isFocused ? 'focused' : ''}`}
-      ref={editorEl}
-    ></div>
-  );
+  return <div className={classes.vectorInputBox} ref={editorEl}></div>;
 }

client/src/pages/databases/types.ts

@@ -1,4 +1,4 @@
-import { FieldObject, CollectionObject } from '@server/types';
+import { FieldObject, CollectionObject, PartitionData } from '@server/types';
 
 export type SearchSingleParams = {
   anns_field: string;
@@ -51,6 +51,7 @@ export type GraphData = {
 
 export type SearchParams = {
   collection: CollectionObject;
+  partitions: PartitionData[];
   searchParams: SearchSingleParams[];
   globalParams: GlobalParams;
   searchResult: SearchResultView[] | null;

client/src/utils/search.ts

@@ -43,6 +43,10 @@ export const buildSearchParams = (searchParams: SearchParams) => {
     consistency_level: searchParams.globalParams.consistency_level,
   };
 
+  if (searchParams.partitions.length > 0) {
+    params.partition_names = searchParams.partitions.map(p => p.name);
+  }
+
   // group_by_field if exists
   if (searchParams.globalParams.group_by_field) {
     params.group_by_field = searchParams.globalParams.group_by_field;

server/src/collections/collections.service.ts

@@ -131,6 +131,9 @@ export class CollectionsService {
 
     // add extra data to schema
     res.schema.hasVectorIndex = vectorFields.every(v => v.index);
+    res.schema.enablePartitionKey = res.schema.fields.some(
+      v => v.is_partition_key
+    );
     res.schema.scalarFields = scalarFields;
     res.schema.vectorFields = vectorFields;
     res.schema.dynamicFields = res.schema.enable_dynamic_field

server/src/types/collections.type.ts

@@ -31,6 +31,7 @@ export interface SchemaObject extends CollectionSchema {
   scalarFields: FieldObject[];
   dynamicFields: FieldObject[];
   hasVectorIndex: boolean;
+  enablePartitionKey: boolean;
 }
 
 export interface DescribeCollectionRes extends DescribeCollectionResponse {