From 9ba7513504a06a833734827aadb6951444167458 Mon Sep 17 00:00:00 2001 From: ryjiang Date: Wed, 30 Oct 2024 17:24:10 +0800 Subject: [PATCH] Support filter partitions in search page (#680) * add parition selectors Signed-off-by: shanghaikid * finish Signed-off-by: shanghaikid --------- Signed-off-by: shanghaikid --- client/src/i18n/cn/search.ts | 4 +- client/src/i18n/en/search.ts | 4 +- client/src/pages/databases/Databases.tsx | 1 + .../collections/search/PartitionsSelector.tsx | 92 +++++++++++++++++++ .../databases/collections/search/Search.tsx | 22 +++++ .../collections/search/VectorInputBox.tsx | 19 ++-- client/src/pages/databases/types.ts | 3 +- client/src/utils/search.ts | 4 + server/src/collections/collections.service.ts | 3 + server/src/types/collections.type.ts | 1 + 10 files changed, 138 insertions(+), 15 deletions(-) create mode 100644 client/src/pages/databases/collections/search/PartitionsSelector.tsx diff --git a/client/src/i18n/cn/search.ts b/client/src/i18n/cn/search.ts index 36f49c98..cf9740fe 100644 --- a/client/src/i18n/cn/search.ts +++ b/client/src/i18n/cn/search.ts @@ -28,7 +28,9 @@ const searchTrans = { outputFields: '输出字段', consistency: '一致性', graphNodeHoverTip: '双击以查看更多', - inputVectorTip: '向量或实体ID', + inputVectorPlaceHolder: '向量或实体ID', + partitionFilter: '分区过滤', + loading: '加载中...', }; export default searchTrans; diff --git a/client/src/i18n/en/search.ts b/client/src/i18n/en/search.ts index da15369a..7cfbf058 100644 --- a/client/src/i18n/en/search.ts +++ b/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; diff --git a/client/src/pages/databases/Databases.tsx b/client/src/pages/databases/Databases.tsx index 4667d9ed..a02f906d 100644 --- a/client/src/pages/databases/Databases.tsx +++ b/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, diff --git a/client/src/pages/databases/collections/search/PartitionsSelector.tsx b/client/src/pages/databases/collections/search/PartitionsSelector.tsx new file mode 100644 index 00000000..794ecc7f --- /dev/null +++ b/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( + 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 ( + { + 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 ? ( + + ) : ( + true} + type="text" + /> + ); + }} + /> + ); +} diff --git a/client/src/pages/databases/collections/search/Search.tsx b/client/src/pages/databases/collections/search/Search.tsx index 0a6ad1e1..0279d243 100644 --- a/client/src/pages/databases/collections/search/Search.tsx +++ b/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 (
{collection && ( @@ -456,6 +470,14 @@ const Search = (props: CollectionDataProps) => { ); })} + + {enablePartitionsFilter && ( + + )}
diff --git a/client/src/pages/databases/collections/search/VectorInputBox.tsx b/client/src/pages/databases/collections/search/VectorInputBox.tsx index 7f7e9387..89ac7974 100644 --- a/client/src/pages/databases/collections/search/VectorInputBox.tsx +++ b/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 ( -
- ); + return
; } diff --git a/client/src/pages/databases/types.ts b/client/src/pages/databases/types.ts index d076034d..48027cd8 100644 --- a/client/src/pages/databases/types.ts +++ b/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; diff --git a/client/src/utils/search.ts b/client/src/utils/search.ts index 5b699e91..439b95e3 100644 --- a/client/src/utils/search.ts +++ b/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; diff --git a/server/src/collections/collections.service.ts b/server/src/collections/collections.service.ts index dcc50f5d..dcff4ce6 100644 --- a/server/src/collections/collections.service.ts +++ b/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 diff --git a/server/src/types/collections.type.ts b/server/src/types/collections.type.ts index 087f875a..90fef236 100644 --- a/server/src/types/collections.type.ts +++ b/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 {