Browse Source

Merge pull request #125 from nameczz/main

add search api
Tumao 4 years ago
parent
commit
fe6f6e2e4c

+ 14 - 1
server/src/collections/collections.controller.ts

@@ -12,7 +12,12 @@ import {
 } from '@nestjs/common';
 import { ApiTags } from '@nestjs/swagger';
 import { CollectionsService } from './collections.service';
-import { CreateCollection, InsertData, ShowCollections } from './dto';
+import {
+  CreateCollection,
+  InsertData,
+  ShowCollections,
+  VectorSearch,
+} from './dto';
 
 @ApiTags('collections')
 @Controller('collections')
@@ -85,4 +90,12 @@ export class CollectionsController {
       ...data,
     });
   }
+
+  @Post(':name/search')
+  async vectorSearch(@Param('name') name: string, @Body() data: VectorSearch) {
+    return await this.collectionsService.vectorSearch({
+      collection_name: name,
+      ...data,
+    });
+  }
 }

+ 7 - 0
server/src/collections/collections.service.ts

@@ -9,6 +9,7 @@ import {
   InsertReq,
   LoadCollectionReq,
   ReleaseLoadCollectionReq,
+  SearchReq,
 } from '@zilliz/milvus2-sdk-node/dist/milvus/types';
 import { throwErrorFromSDK } from '../utils/Error';
 import { findKeyValue } from '../utils/Helper';
@@ -73,6 +74,12 @@ export class CollectionsService {
     return res;
   }
 
+  async vectorSearch(data: SearchReq) {
+    const res = await this.milvusClient.search(data);
+    throwErrorFromSDK(res.status);
+    return res;
+  }
+
   /**
    * We do not throw error for this.
    * Because if collection dont have index, it will throw error.

+ 64 - 2
server/src/collections/dto.ts

@@ -6,14 +6,22 @@ import {
   IsArray,
   ArrayNotEmpty,
   IsEnum,
+  ArrayMinSize,
 } from 'class-validator';
 import {
   FieldType,
   ShowCollectionsType,
 } from '@zilliz/milvus2-sdk-node/dist/milvus/types/Collection'; // todo: need improve like export types in root file.
-import { DataType } from '@zilliz/milvus2-sdk-node/dist/milvus/types/Common';
+import {
+  DataType,
+  KeyValuePair,
+} from '@zilliz/milvus2-sdk-node/dist/milvus/types/Common';
 import { ApiProperty } from '@nestjs/swagger';
 
+enum VectorTypes {
+  Binary = DataType.BinaryVector,
+  Float = DataType.FloatVector,
+}
 export class CreateCollection {
   @ApiProperty({
     description: 'Milvus collection name',
@@ -34,7 +42,6 @@ export class CreateCollection {
 
   @ApiProperty({
     description: 'Field data type',
-    enum: DataType,
   })
   @IsArray()
   @ArrayNotEmpty()
@@ -67,3 +74,58 @@ export class InsertData {
   })
   readonly fields_data: any[];
 }
+
+export class VectorSearch {
+  @ApiProperty({
+    description: 'Milvus collection name',
+  })
+  @IsString()
+  @IsNotEmpty({
+    message: 'collection_name is empty',
+  })
+  collection_name: string;
+
+  @ApiProperty({
+    description: 'Partition in this collection',
+  })
+  @IsOptional()
+  partition_names?: string[];
+
+  @ApiProperty({
+    description: 'Non vector fields filter.',
+  })
+  @IsString()
+  expr?: string;
+
+  @ApiProperty({
+    description: 'Vector search params',
+    default: [{ key: 'metric_type', value: 'L2' }],
+  })
+  @IsArray()
+  @ArrayMinSize(1)
+  search_params: KeyValuePair[];
+
+  @ApiProperty({
+    description: 'Searched vector value',
+    default: [[1, 2, 3, 4]],
+  })
+  @IsArray()
+  @ArrayMinSize(1)
+  vectors: number[][];
+
+  @ApiProperty({
+    description:
+      'Define what non vector fields you want return in search results',
+    default: ['a', 'b'],
+  })
+  @IsArray()
+  @IsOptional()
+  output_fields?: string[];
+
+  @ApiProperty({
+    description: 'Only support 101(binary) or 100(float)',
+    enum: VectorTypes,
+  })
+  @IsEnum(VectorTypes, { message: 'Type allow all->0 inmemory->1' })
+  vector_type: DataType.BinaryVector | DataType.FloatVector;
+}