GeneralExample.java 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488
  1. /*
  2. * Licensed to the Apache Software Foundation (ASF) under one
  3. * or more contributor license agreements. See the NOTICE file
  4. * distributed with this work for additional information
  5. * regarding copyright ownership. The ASF licenses this file
  6. * to you under the Apache License, Version 2.0 (the
  7. * "License"); you may not use this file except in compliance
  8. * with the License. You may obtain a copy of the License at
  9. *
  10. * http://www.apache.org/licenses/LICENSE-2.0
  11. *
  12. * Unless required by applicable law or agreed to in writing,
  13. * software distributed under the License is distributed on an
  14. * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15. * KIND, either express or implied. See the License for the
  16. * specific language governing permissions and limitations
  17. * under the License.
  18. */
  19. package io.milvus;
  20. import io.milvus.client.MilvusServiceClient;
  21. import io.milvus.grpc.*;
  22. import io.milvus.param.*;
  23. import io.milvus.param.collection.*;
  24. import io.milvus.param.dml.*;
  25. import io.milvus.param.index.*;
  26. import io.milvus.param.partition.*;
  27. import io.milvus.Response.*;
  28. import java.nio.ByteBuffer;
  29. import java.util.*;
  30. public class GeneralExample {
  31. private static final MilvusServiceClient milvusClient;
  32. static {
  33. ConnectParam connectParam = ConnectParam.newBuilder()
  34. .withHost("localhost")
  35. .withPort(19530)
  36. .build();
  37. milvusClient = new MilvusServiceClient(connectParam);
  38. }
  39. private static final String COLLECTION_NAME = "TEST";
  40. private static final String ID_FIELD = "userID";
  41. private static final String VECTOR_FIELD = "userFace";
  42. private static final Integer VECTOR_DIM = 64;
  43. private static final String AGE_FIELD = "userAge";
  44. private static final String PROFILE_FIELD = "userProfile";
  45. private static final Integer BINARY_DIM = 128;
  46. private static final IndexType INDEX_TYPE = IndexType.IVF_FLAT;
  47. private static final String INDEX_PARAM = "{\"nlist\":128}";
  48. private static final Integer SEARCH_K = 5;
  49. private static final String SEARCH_PARAM = "{\"nprobe\":10}";
  50. private R<RpcStatus> createCollection() {
  51. System.out.println("========== createCollection() ==========");
  52. FieldType fieldType1 = FieldType.newBuilder()
  53. .withName(ID_FIELD)
  54. .withDescription("user identification")
  55. .withDataType(DataType.Int64)
  56. .withPrimaryKey(true)
  57. .withAutoID(true)
  58. .build();
  59. FieldType fieldType2 = FieldType.newBuilder()
  60. .withName(VECTOR_FIELD)
  61. .withDescription("face embedding")
  62. .withDataType(DataType.FloatVector)
  63. .withDimension(VECTOR_DIM)
  64. .build();
  65. FieldType fieldType3 = FieldType.newBuilder()
  66. .withName(AGE_FIELD)
  67. .withDescription("user age")
  68. .withDataType(DataType.Int8)
  69. .build();
  70. FieldType fieldType4 = FieldType.newBuilder()
  71. .withName(PROFILE_FIELD)
  72. .withDescription("user profile")
  73. .withDataType(DataType.BinaryVector)
  74. .withDimension(BINARY_DIM)
  75. .build();
  76. CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
  77. .withCollectionName(COLLECTION_NAME)
  78. .withDescription("customer info")
  79. .withShardsNum(2)
  80. .addFieldType(fieldType1)
  81. .addFieldType(fieldType2)
  82. .addFieldType(fieldType3)
  83. .addFieldType(fieldType4)
  84. .build();
  85. R<RpcStatus> response = milvusClient.createCollection(createCollectionReq);
  86. System.out.println(response);
  87. return response;
  88. }
  89. private R<RpcStatus> dropCollection() {
  90. System.out.println("========== dropCollection() ==========");
  91. R<RpcStatus> response = milvusClient.dropCollection(DropCollectionParam.newBuilder()
  92. .withCollectionName(COLLECTION_NAME)
  93. .build());
  94. System.out.println(response);
  95. return response;
  96. }
  97. private R<Boolean> hasCollection() {
  98. System.out.println("========== hasCollection() ==========");
  99. R<Boolean> response = milvusClient.hasCollection(HasCollectionParam.newBuilder()
  100. .withCollectionName(COLLECTION_NAME)
  101. .build());
  102. System.out.println(response);
  103. return response;
  104. }
  105. private R<RpcStatus> loadCollection() {
  106. System.out.println("========== loadCollection() ==========");
  107. R<RpcStatus> response = milvusClient.loadCollection(LoadCollectionParam.newBuilder()
  108. .withCollectionName(COLLECTION_NAME)
  109. .build());
  110. System.out.println(response);
  111. return response;
  112. }
  113. private R<RpcStatus> releaseCollection() {
  114. System.out.println("========== releaseCollection() ==========");
  115. R<RpcStatus> response = milvusClient.releaseCollection(ReleaseCollectionParam.newBuilder()
  116. .withCollectionName(COLLECTION_NAME)
  117. .build());
  118. System.out.println(response);
  119. return response;
  120. }
  121. private R<DescribeCollectionResponse> describeCollection() {
  122. System.out.println("========== describeCollection() ==========");
  123. R<DescribeCollectionResponse> response = milvusClient.describeCollection(DescribeCollectionParam.newBuilder()
  124. .withCollectionName(COLLECTION_NAME)
  125. .build());
  126. System.out.println(response);
  127. return response;
  128. }
  129. private R<GetCollectionStatisticsResponse> getCollectionStatistics() {
  130. System.out.println("========== getCollectionStatistics() ==========");
  131. R<GetCollectionStatisticsResponse> response = milvusClient.getCollectionStatistics(
  132. GetCollectionStatisticsParam.newBuilder()
  133. .withCollectionName(COLLECTION_NAME)
  134. .build());
  135. GetCollStatResponseWrapper wrapper = new GetCollStatResponseWrapper(response.getData());
  136. System.out.println("Collection row count: " + wrapper.GetRowCount());
  137. return response;
  138. }
  139. private R<ShowCollectionsResponse> showCollections() {
  140. System.out.println("========== showCollections() ==========");
  141. R<ShowCollectionsResponse> response = milvusClient.showCollections(ShowCollectionsParam.newBuilder()
  142. .build());
  143. System.out.println(response);
  144. return response;
  145. }
  146. private R<RpcStatus> createPartition(String partitionName) {
  147. System.out.println("========== createPartition() ==========");
  148. R<RpcStatus> response = milvusClient.createPartition(CreatePartitionParam.newBuilder()
  149. .withCollectionName(COLLECTION_NAME)
  150. .withPartitionName(partitionName)
  151. .build());
  152. System.out.println(response);
  153. return response;
  154. }
  155. private R<RpcStatus> dropPartition(String partitionName) {
  156. System.out.println("========== dropPartition() ==========");
  157. R<RpcStatus> response = milvusClient.dropPartition(DropPartitionParam.newBuilder()
  158. .withCollectionName(COLLECTION_NAME)
  159. .withPartitionName(partitionName)
  160. .build());
  161. System.out.println(response);
  162. return response;
  163. }
  164. private R<Boolean> hasPartition(String partitionName) {
  165. System.out.println("========== hasPartition() ==========");
  166. R<Boolean> response = milvusClient.hasPartition(HasPartitionParam.newBuilder()
  167. .withCollectionName(COLLECTION_NAME)
  168. .withPartitionName(partitionName)
  169. .build());
  170. System.out.println(response);
  171. return response;
  172. }
  173. private R<RpcStatus> releasePartition(String partitionName) {
  174. System.out.println("========== releasePartition() ==========");
  175. R<RpcStatus> response = milvusClient.releasePartitions(ReleasePartitionsParam.newBuilder()
  176. .withCollectionName(COLLECTION_NAME)
  177. .addPartitionName(partitionName)
  178. .build());
  179. System.out.println(response);
  180. return response;
  181. }
  182. private R<ShowPartitionsResponse> showPartitions() {
  183. System.out.println("========== showPartitions() ==========");
  184. R<ShowPartitionsResponse> response = milvusClient.showPartitions(ShowPartitionsParam.newBuilder()
  185. .withCollectionName(COLLECTION_NAME)
  186. .build());
  187. System.out.println(response);
  188. return response;
  189. }
  190. private R<RpcStatus> createIndex() {
  191. System.out.println("========== createIndex() ==========");
  192. R<RpcStatus> response = milvusClient.createIndex(CreateIndexParam.newBuilder()
  193. .withCollectionName(COLLECTION_NAME)
  194. .withFieldName(VECTOR_FIELD)
  195. .withIndexType(INDEX_TYPE)
  196. .withMetricType(MetricType.L2)
  197. .withExtraParam(INDEX_PARAM)
  198. .withSyncMode(Boolean.TRUE)
  199. .build());
  200. System.out.println(response);
  201. return response;
  202. }
  203. private R<RpcStatus> dropIndex() {
  204. System.out.println("========== dropIndex() ==========");
  205. R<RpcStatus> response = milvusClient.dropIndex(DropIndexParam.newBuilder()
  206. .withCollectionName(COLLECTION_NAME)
  207. .withFieldName(VECTOR_FIELD)
  208. .build());
  209. System.out.println(response);
  210. return response;
  211. }
  212. private R<DescribeIndexResponse> describeIndex() {
  213. System.out.println("========== describeIndex() ==========");
  214. R<DescribeIndexResponse> response = milvusClient.describeIndex(DescribeIndexParam.newBuilder()
  215. .withCollectionName(COLLECTION_NAME)
  216. .withFieldName(VECTOR_FIELD)
  217. .build());
  218. System.out.println(response);
  219. return response;
  220. }
  221. private R<GetIndexStateResponse> getIndexState() {
  222. System.out.println("========== getIndexState() ==========");
  223. R<GetIndexStateResponse> response = milvusClient.getIndexState(GetIndexStateParam.newBuilder()
  224. .withCollectionName(COLLECTION_NAME)
  225. .withFieldName(VECTOR_FIELD)
  226. .build());
  227. System.out.println(response);
  228. return response;
  229. }
  230. private R<GetIndexBuildProgressResponse> getIndexBuildProgress() {
  231. System.out.println("========== getIndexBuildProgress() ==========");
  232. R<GetIndexBuildProgressResponse> response = milvusClient.getIndexBuildProgress(
  233. GetIndexBuildProgressParam.newBuilder()
  234. .withCollectionName(COLLECTION_NAME)
  235. .build());
  236. System.out.println(response);
  237. return response;
  238. }
  239. private R<MutationResult> delete(String partitionName, String expr) {
  240. System.out.println("========== delete() ==========");
  241. DeleteParam build = DeleteParam.newBuilder()
  242. .withCollectionName(COLLECTION_NAME)
  243. .withPartitionName(partitionName)
  244. .withExpr(expr)
  245. .build();
  246. R<MutationResult> response = milvusClient.delete(build);
  247. System.out.println(response.getData());
  248. return response;
  249. }
  250. private R<SearchResults> searchFace(String expr) {
  251. System.out.println("========== searchFace() ==========");
  252. List<String> outFields = Collections.singletonList(AGE_FIELD);
  253. List<List<Float>> vectors = generateFloatVectors(5);
  254. SearchParam searchParam = SearchParam.newBuilder()
  255. .withCollectionName(COLLECTION_NAME)
  256. .withMetricType(MetricType.L2)
  257. .withOutFields(outFields)
  258. .withTopK(SEARCH_K)
  259. .withVectors(vectors)
  260. .withVectorFieldName(VECTOR_FIELD)
  261. .withExpr(expr)
  262. .withParams(SEARCH_PARAM)
  263. .build();
  264. R<SearchResults> response = milvusClient.search(searchParam);
  265. SearchResultsWrapper wrapper = new SearchResultsWrapper(response.getData().getResults());
  266. for (int i = 0; i < vectors.size(); ++i) {
  267. System.out.println("Search result of No." + i);
  268. List<SearchResultsWrapper.IDScore> scores = wrapper.GetIDScore(i);
  269. System.out.println(scores);
  270. }
  271. System.out.println(wrapper.GetFieldData(AGE_FIELD).getFieldData());
  272. return response;
  273. }
  274. private R<SearchResults> searchProfile(String expr) {
  275. System.out.println("========== searchProfile() ==========");
  276. List<String> outFields = Collections.singletonList(AGE_FIELD);
  277. List<ByteBuffer> vectors = generateBinaryVectors(5);
  278. SearchParam searchParam = SearchParam.newBuilder()
  279. .withCollectionName(COLLECTION_NAME)
  280. .withMetricType(MetricType.HAMMING)
  281. .withOutFields(outFields)
  282. .withTopK(SEARCH_K)
  283. .withVectors(vectors)
  284. .withVectorFieldName(PROFILE_FIELD)
  285. .withExpr(expr)
  286. .withParams(SEARCH_PARAM)
  287. .build();
  288. R<SearchResults> response = milvusClient.search(searchParam);
  289. SearchResultsWrapper wrapper = new SearchResultsWrapper(response.getData().getResults());
  290. for (int i = 0; i < vectors.size(); ++i) {
  291. System.out.println("Search result of No." + i);
  292. List<SearchResultsWrapper.IDScore> scores = wrapper.GetIDScore(i);
  293. System.out.println(scores);
  294. }
  295. System.out.println(wrapper.GetFieldData(AGE_FIELD).getFieldData());
  296. return response;
  297. }
  298. private R<CalcDistanceResults> calDistance() {
  299. System.out.println("========== calDistance() ==========");
  300. Random ran=new Random();
  301. List<Float> vector1 = new ArrayList<>();
  302. List<Float> vector2 = new ArrayList<>();
  303. for (int d = 0; d < VECTOR_DIM; ++d) {
  304. vector1.add(ran.nextFloat());
  305. vector2.add(ran.nextFloat());
  306. }
  307. CalcDistanceParam calcDistanceParam = CalcDistanceParam.newBuilder()
  308. .withVectorsLeft(Collections.singletonList(vector1))
  309. .withVectorsRight(Collections.singletonList(vector2))
  310. .withMetricType(MetricType.L2)
  311. .build();
  312. R<CalcDistanceResults> response = milvusClient.calcDistance(calcDistanceParam);
  313. System.out.println(response);
  314. return response;
  315. }
  316. private R<QueryResults> query(String expr) {
  317. System.out.println("========== query() ==========");
  318. List<String> fields = Arrays.asList(ID_FIELD, AGE_FIELD);
  319. QueryParam test = QueryParam.newBuilder()
  320. .withCollectionName(COLLECTION_NAME)
  321. .withExpr(expr)
  322. .withOutFields(fields)
  323. .build();
  324. R<QueryResults> response = milvusClient.query(test);
  325. QueryResultsWrapper wrapper = new QueryResultsWrapper(response.getData());
  326. System.out.println(ID_FIELD + ":" + wrapper.getFieldWrapper(ID_FIELD).getFieldData().toString());
  327. System.out.println(AGE_FIELD + ":" + wrapper.getFieldWrapper(AGE_FIELD).getFieldData().toString());
  328. System.out.println("Query row count: " + wrapper.getFieldWrapper(ID_FIELD).getRowCount());
  329. return response;
  330. }
  331. private R<MutationResult> insert(String partitionName, int count) {
  332. System.out.println("========== insert() ==========");
  333. List<List<Float>> vectors = generateFloatVectors(count);
  334. List<ByteBuffer> profiles = generateBinaryVectors(count);
  335. Random ran = new Random();
  336. List<Integer> ages = new ArrayList<>();
  337. for (long i = 0L; i < count; ++i) {
  338. ages.add(ran.nextInt(99));
  339. }
  340. List<InsertParam.Field> fields = new ArrayList<>();
  341. fields.add(new InsertParam.Field(VECTOR_FIELD, DataType.FloatVector, vectors));
  342. fields.add(new InsertParam.Field(PROFILE_FIELD, DataType.BinaryVector, profiles));
  343. fields.add(new InsertParam.Field(AGE_FIELD, DataType.Int8, ages));
  344. InsertParam insertParam = InsertParam.newBuilder()
  345. .withCollectionName(COLLECTION_NAME)
  346. .withPartitionName(partitionName)
  347. .withFields(fields)
  348. .build();
  349. return milvusClient.insert(insertParam);
  350. }
  351. private List<List<Float>> generateFloatVectors(int count) {
  352. Random ran = new Random();
  353. List<List<Float>> vectors = new ArrayList<>();
  354. for (int n = 0; n < count; ++n) {
  355. List<Float> vector = new ArrayList<>();
  356. for (int i = 0; i < VECTOR_DIM; ++i) {
  357. vector.add(ran.nextFloat());
  358. }
  359. vectors.add(vector);
  360. }
  361. return vectors;
  362. }
  363. private List<ByteBuffer> generateBinaryVectors(int count) {
  364. Random ran = new Random();
  365. List<ByteBuffer> vectors = new ArrayList<>();
  366. int byteCount = BINARY_DIM/8;
  367. for (int n = 0; n < count; ++n) {
  368. ByteBuffer vector = ByteBuffer.allocate(byteCount);
  369. for (int i = 0; i < byteCount; ++i) {
  370. vector.put((byte)ran.nextInt(Byte.MAX_VALUE));
  371. }
  372. vectors.add(vector);
  373. }
  374. return vectors;
  375. }
  376. public static void main(String[] args) {
  377. GeneralExample example = new GeneralExample();
  378. example.dropCollection();
  379. example.createCollection();
  380. example.hasCollection();
  381. example.describeCollection();
  382. example.showCollections();
  383. example.loadCollection();
  384. final String partitionName = "p1";
  385. example.createPartition(partitionName);
  386. example.hasPartition(partitionName);
  387. example.showPartitions();
  388. final int row_count = 10000;
  389. List<Long> deleteIds = new ArrayList<>();
  390. Random ran = new Random();
  391. for (int i = 0; i < 100; ++i) {
  392. R<MutationResult> result = example.insert(partitionName, row_count);
  393. InsertResultWrapper wrapper = new InsertResultWrapper(result.getData());
  394. List<Long> ids = wrapper.getLongIDs();
  395. deleteIds.add(ids.get(ran.nextInt(row_count)));
  396. }
  397. example.getCollectionStatistics();
  398. example.createIndex();
  399. example.describeIndex();
  400. example.getIndexBuildProgress();
  401. example.getIndexState();
  402. String deleteExpr = ID_FIELD + " in " + deleteIds.toString();
  403. example.delete(partitionName, deleteExpr);
  404. String queryExpr = AGE_FIELD + " == 60";
  405. example.query(queryExpr);
  406. String searchExpr = AGE_FIELD + " > 50";
  407. example.searchFace(searchExpr);
  408. searchExpr = AGE_FIELD + " <= 30";
  409. example.searchProfile(searchExpr);
  410. example.calDistance();
  411. example.releasePartition(partitionName);
  412. example.releaseCollection();
  413. example.dropPartition(partitionName);
  414. example.dropIndex();
  415. example.dropCollection();
  416. }
  417. }