MilvusClientDockerTest.java 70 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684
  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.client;
  20. import com.alibaba.fastjson.JSON;
  21. import com.alibaba.fastjson.JSONObject;
  22. import com.google.common.util.concurrent.ListenableFuture;
  23. import io.milvus.grpc.*;
  24. import io.milvus.param.*;
  25. import io.milvus.param.collection.*;
  26. //import io.milvus.param.credential.*;
  27. import io.milvus.param.highlevel.dml.DeleteIdsParam;
  28. import io.milvus.param.highlevel.dml.GetIdsParam;
  29. import io.milvus.param.highlevel.dml.response.DeleteResponse;
  30. import io.milvus.param.highlevel.dml.response.GetResponse;
  31. import io.milvus.param.dml.InsertParam;
  32. import io.milvus.param.dml.QueryParam;
  33. import io.milvus.param.dml.SearchParam;
  34. import io.milvus.param.index.CreateIndexParam;
  35. import io.milvus.param.index.DescribeIndexParam;
  36. import io.milvus.param.index.DropIndexParam;
  37. import io.milvus.param.index.GetIndexStateParam;
  38. import io.milvus.param.partition.GetPartitionStatisticsParam;
  39. import io.milvus.param.partition.ShowPartitionsParam;
  40. import io.milvus.response.*;
  41. import org.apache.commons.text.RandomStringGenerator;
  42. import org.apache.logging.log4j.LogManager;
  43. import org.apache.logging.log4j.Logger;
  44. import org.codehaus.plexus.util.FileUtils;
  45. import org.junit.jupiter.api.Assertions;
  46. import org.junit.jupiter.api.BeforeAll;
  47. import org.junit.jupiter.api.AfterAll;
  48. import org.junit.jupiter.api.Test;
  49. import java.nio.ByteBuffer;
  50. import java.util.*;
  51. import java.util.concurrent.ExecutionException;
  52. import java.util.concurrent.TimeUnit;
  53. import static org.junit.Assert.assertEquals;
  54. import static org.junit.Assert.assertTrue;
  55. class MilvusClientDockerTest {
  56. private static final Logger logger = LogManager.getLogger("MilvusClientTest");
  57. private static MilvusClient client;
  58. private static RandomStringGenerator generator;
  59. private static final int dimension = 128;
  60. private static final Boolean useDockerCompose = Boolean.TRUE;
  61. private static void startDockerContainer() {
  62. if (!useDockerCompose) {
  63. return;
  64. }
  65. // start the test container
  66. Runtime runtime = Runtime.getRuntime();
  67. String bashCommand = "docker-compose up -d";
  68. try {
  69. logger.info(bashCommand);
  70. Process pro = runtime.exec(bashCommand);
  71. int status = pro.waitFor();
  72. if (status != 0) {
  73. logger.error("Failed to start docker compose, status " + status);
  74. }
  75. logger.info("Docker compose is up");
  76. } catch (Throwable t) {
  77. logger.error("Failed to execute docker compose up", t);
  78. }
  79. ConnectParam connectParam = connectParamBuilder().withAuthorization("root", "Milvus").build();
  80. MilvusServiceClient tempClient = new MilvusServiceClient(connectParam);
  81. long waitTime = 0;
  82. while (true) {
  83. // although milvus container is alive, it is still in initializing,
  84. // connection will failed and get error "proxy not health".
  85. // check health state for every few seconds, until the server is ready.
  86. long checkInterval = 3;
  87. try {
  88. TimeUnit.SECONDS.sleep(checkInterval);
  89. } catch (InterruptedException t) {
  90. logger.error("Interrupted", t);
  91. break;
  92. }
  93. try{
  94. R<CheckHealthResponse> resp = tempClient.checkHealth();
  95. if (resp.getData().getIsHealthy()) {
  96. logger.info(String.format("Milvus service is ready after %d seconds", waitTime));
  97. break;
  98. }
  99. logger.info("Milvus service is not ready, waiting...");
  100. } catch (Throwable t) {
  101. logger.error("Milvus service is in initialize, not able to connect", t);
  102. }
  103. waitTime += checkInterval;
  104. if (waitTime > 120) {
  105. logger.error(String.format("Milvus service failed to start within %d seconds", waitTime));
  106. break;
  107. }
  108. }
  109. }
  110. private static void stopDockerContainer() {
  111. if (!useDockerCompose) {
  112. return;
  113. }
  114. // stop all test dockers
  115. Runtime runtime = Runtime.getRuntime();
  116. String bashCommand = "docker-compose down";
  117. try {
  118. logger.info("Milvus service stopping...");
  119. TimeUnit.SECONDS.sleep(5);
  120. logger.info(bashCommand);
  121. Process pro = runtime.exec(bashCommand);
  122. int status = pro.waitFor();
  123. if (status != 0) {
  124. logger.error("Failed to stop test docker containers" + pro.getOutputStream().toString());
  125. }
  126. } catch (Throwable t) {
  127. logger.error("Failed to execute docker compose down", t);
  128. }
  129. // clean up log dir
  130. runtime = Runtime.getRuntime();
  131. bashCommand = "docker-compose rm";
  132. try {
  133. logger.info(bashCommand);
  134. Process pro = runtime.exec(bashCommand);
  135. int status = pro.waitFor();
  136. if (status != 0) {
  137. logger.error("Failed to clean up test docker containers" + pro.getOutputStream().toString());
  138. }
  139. logger.error("Clean up volume directory of Docker");
  140. FileUtils.cleanDirectory("volumes");
  141. } catch (Throwable t) {
  142. logger.error("Failed to remove docker compose volume", t);
  143. }
  144. }
  145. @BeforeAll
  146. public static void setUp() {
  147. startDockerContainer();
  148. ConnectParam connectParam = connectParamBuilder().withAuthorization("root", "Milvus").build();
  149. client = new MilvusServiceClient(connectParam);
  150. generator = new RandomStringGenerator.Builder().withinRange('a', 'z').build();
  151. }
  152. @AfterAll
  153. public static void tearDown() {
  154. if (client != null) {
  155. client.close();
  156. }
  157. stopDockerContainer();
  158. }
  159. protected static ConnectParam.Builder connectParamBuilder() {
  160. return connectParamBuilder("localhost", 19530);
  161. }
  162. private static ConnectParam.Builder connectParamBuilder(String host, int port) {
  163. return ConnectParam.newBuilder().withHost(host).withPort(port);
  164. }
  165. private List<List<Float>> generateFloatVectors(int count) {
  166. Random ran = new Random();
  167. List<List<Float>> vectors = new ArrayList<>();
  168. for (int n = 0; n < count; ++n) {
  169. List<Float> vector = new ArrayList<>();
  170. for (int i = 0; i < dimension; ++i) {
  171. vector.add(ran.nextFloat());
  172. }
  173. vectors.add(vector);
  174. }
  175. return vectors;
  176. }
  177. private List<List<Float>> normalizeFloatVectors(List<List<Float>> src) {
  178. for (List<Float> vector : src) {
  179. double total = 0.0;
  180. for (Float val : vector) {
  181. total = total + val * val;
  182. }
  183. float squre = (float) Math.sqrt(total);
  184. for (int i = 0; i < vector.size(); ++i) {
  185. vector.set(i, vector.get(i) / squre);
  186. }
  187. }
  188. return src;
  189. }
  190. private List<ByteBuffer> generateBinaryVectors(int count) {
  191. Random ran = new Random();
  192. List<ByteBuffer> vectors = new ArrayList<>();
  193. int byteCount = dimension / 8;
  194. for (int n = 0; n < count; ++n) {
  195. ByteBuffer vector = ByteBuffer.allocate(byteCount);
  196. for (int i = 0; i < byteCount; ++i) {
  197. vector.put((byte) ran.nextInt(Byte.MAX_VALUE));
  198. }
  199. vectors.add(vector);
  200. }
  201. return vectors;
  202. }
  203. @Test
  204. void testFloatVectors() {
  205. String randomCollectionName = generator.generate(10);
  206. // collection schema
  207. String field1Name = "long_field";
  208. String field2Name = "vec_field";
  209. String field3Name = "bool_field";
  210. String field4Name = "double_field";
  211. String field5Name = "int_field";
  212. List<FieldType> fieldsSchema = new ArrayList<>();
  213. fieldsSchema.add(FieldType.newBuilder()
  214. .withPrimaryKey(true)
  215. .withAutoID(false)
  216. .withDataType(DataType.Int64)
  217. .withName(field1Name)
  218. .withDescription("identity")
  219. .build());
  220. fieldsSchema.add(FieldType.newBuilder()
  221. .withDataType(DataType.FloatVector)
  222. .withName(field2Name)
  223. .withDescription("face")
  224. .withDimension(dimension)
  225. .build());
  226. fieldsSchema.add(FieldType.newBuilder()
  227. .withDataType(DataType.Bool)
  228. .withName(field3Name)
  229. .withDescription("gender")
  230. .build());
  231. fieldsSchema.add(FieldType.newBuilder()
  232. .withDataType(DataType.Double)
  233. .withName(field4Name)
  234. .withDescription("weight")
  235. .build());
  236. fieldsSchema.add(FieldType.newBuilder()
  237. .withDataType(DataType.Int8)
  238. .withName(field5Name)
  239. .withDescription("age")
  240. .build());
  241. // create collection
  242. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  243. .withCollectionName(randomCollectionName)
  244. .withDescription("test")
  245. .withFieldTypes(fieldsSchema)
  246. .build();
  247. R<RpcStatus> createR = client.createCollection(createParam);
  248. assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  249. R<DescribeCollectionResponse> response = client.describeCollection(DescribeCollectionParam.newBuilder()
  250. .withCollectionName(randomCollectionName)
  251. .build());
  252. DescCollResponseWrapper desc = new DescCollResponseWrapper(response.getData());
  253. System.out.println(desc.toString());
  254. R<ShowPartitionsResponse> spResp = client.showPartitions(ShowPartitionsParam.newBuilder()
  255. .withCollectionName(randomCollectionName)
  256. .build());
  257. System.out.println(spResp);
  258. ShowPartResponseWrapper wra = new ShowPartResponseWrapper(spResp.getData());
  259. List<ShowPartResponseWrapper.PartitionInfo> parts = wra.getPartitionsInfo();
  260. System.out.println("Partition num: "+parts.size());
  261. // insert data
  262. int rowCount = 10000;
  263. List<Long> ids = new ArrayList<>();
  264. List<Boolean> genders = new ArrayList<>();
  265. List<Double> weights = new ArrayList<>();
  266. List<Short> ages = new ArrayList<>();
  267. for (long i = 0L; i < rowCount; ++i) {
  268. ids.add(i);
  269. genders.add(i % 3 == 0 ? Boolean.TRUE : Boolean.FALSE);
  270. weights.add(((double) (i + 1) / 100));
  271. ages.add((short) ((i + 1) % 99));
  272. }
  273. List<List<Float>> vectors = generateFloatVectors(rowCount);
  274. List<InsertParam.Field> fieldsInsert = new ArrayList<>();
  275. fieldsInsert.add(new InsertParam.Field(field1Name, ids));
  276. fieldsInsert.add(new InsertParam.Field(field5Name, ages));
  277. fieldsInsert.add(new InsertParam.Field(field4Name, weights));
  278. fieldsInsert.add(new InsertParam.Field(field3Name, genders));
  279. fieldsInsert.add(new InsertParam.Field(field2Name, vectors));
  280. InsertParam insertParam = InsertParam.newBuilder()
  281. .withCollectionName(randomCollectionName)
  282. .withFields(fieldsInsert)
  283. .build();
  284. R<MutationResult> insertR = client.withTimeout(10, TimeUnit.SECONDS).insert(insertParam);
  285. assertEquals(R.Status.Success.getCode(), insertR.getStatus().intValue());
  286. MutationResultWrapper insertResultWrapper = new MutationResultWrapper(insertR.getData());
  287. System.out.println(insertResultWrapper.getInsertCount() + " rows inserted");
  288. // get collection statistics
  289. R<GetCollectionStatisticsResponse> statR = client.getCollectionStatistics(GetCollectionStatisticsParam
  290. .newBuilder()
  291. .withCollectionName(randomCollectionName)
  292. .withFlush(true)
  293. .build());
  294. assertEquals(R.Status.Success.getCode(), statR.getStatus().intValue());
  295. GetCollStatResponseWrapper stat = new GetCollStatResponseWrapper(statR.getData());
  296. System.out.println("Collection row count: " + stat.getRowCount());
  297. // get partition statistics
  298. R<GetPartitionStatisticsResponse> statPartR = client.getPartitionStatistics(GetPartitionStatisticsParam
  299. .newBuilder()
  300. .withCollectionName(randomCollectionName)
  301. .withPartitionName("_default") // each collection has '_default' partition
  302. .withFlush(true)
  303. .build());
  304. assertEquals(R.Status.Success.getCode(), statPartR.getStatus().intValue());
  305. GetPartStatResponseWrapper statPart = new GetPartStatResponseWrapper(statPartR.getData());
  306. System.out.println("Partition row count: " + statPart.getRowCount());
  307. // create index on scalar field
  308. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  309. .withCollectionName(randomCollectionName)
  310. .withFieldName(field5Name)
  311. .withIndexType(IndexType.STL_SORT)
  312. .withSyncMode(Boolean.TRUE)
  313. .build();
  314. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  315. assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  316. // create index on vector field
  317. indexParam = CreateIndexParam.newBuilder()
  318. .withCollectionName(randomCollectionName)
  319. .withFieldName(field2Name)
  320. .withIndexName("abv")
  321. .withIndexType(IndexType.HNSW)
  322. .withMetricType(MetricType.L2)
  323. .withExtraParam("{\"M\":16,\"efConstruction\":64}")
  324. .withSyncMode(Boolean.TRUE)
  325. .withSyncWaitingInterval(500L)
  326. .withSyncWaitingTimeout(30L)
  327. .build();
  328. createIndexR = client.createIndex(indexParam);
  329. assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  330. // get index description
  331. DescribeIndexParam descIndexParam = DescribeIndexParam.newBuilder()
  332. .withCollectionName(randomCollectionName)
  333. .withIndexName(indexParam.getIndexName())
  334. .build();
  335. R<DescribeIndexResponse> descIndexR = client.describeIndex(descIndexParam);
  336. assertEquals(R.Status.Success.getCode(), descIndexR.getStatus().intValue());
  337. DescIndexResponseWrapper indexDesc = new DescIndexResponseWrapper(descIndexR.getData());
  338. System.out.println("Index description: " + indexDesc.toString());
  339. // load collection
  340. R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
  341. .withCollectionName(randomCollectionName)
  342. .build());
  343. assertEquals(R.Status.Success.getCode(), loadR.getStatus().intValue());
  344. // show collections
  345. R<ShowCollectionsResponse> showR = client.showCollections(ShowCollectionsParam.newBuilder()
  346. .addCollectionName(randomCollectionName)
  347. .build());
  348. assertEquals(R.Status.Success.getCode(), showR.getStatus().intValue());
  349. ShowCollResponseWrapper info = new ShowCollResponseWrapper(showR.getData());
  350. System.out.println("Collection info: " + info.toString());
  351. // show partitions
  352. R<ShowPartitionsResponse> showPartR = client.showPartitions(ShowPartitionsParam.newBuilder()
  353. .withCollectionName(randomCollectionName)
  354. .addPartitionName("_default") // each collection has a '_default' partition
  355. .build());
  356. assertEquals(R.Status.Success.getCode(), showPartR.getStatus().intValue());
  357. ShowPartResponseWrapper infoPart = new ShowPartResponseWrapper(showPartR.getData());
  358. System.out.println("Partition info: " + infoPart.toString());
  359. // query vectors to verify
  360. List<Long> queryIDs = new ArrayList<>();
  361. List<Double> compareWeights = new ArrayList<>();
  362. int nq = 5;
  363. Random ran = new Random();
  364. int randomIndex = ran.nextInt(rowCount - nq);
  365. for (int i = randomIndex; i < randomIndex + nq; ++i) {
  366. queryIDs.add(ids.get(i));
  367. compareWeights.add(weights.get(i));
  368. }
  369. String expr = field1Name + " in " + queryIDs.toString();
  370. List<String> outputFields = Arrays.asList(field1Name, field2Name, field3Name, field4Name, field5Name);
  371. QueryParam queryParam = QueryParam.newBuilder()
  372. .withCollectionName(randomCollectionName)
  373. .withExpr(expr)
  374. .withOutFields(outputFields)
  375. .build();
  376. R<QueryResults> queryR = client.query(queryParam);
  377. assertEquals(R.Status.Success.getCode(), queryR.getStatus().intValue());
  378. // verify query result
  379. QueryResultsWrapper queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
  380. for (String fieldName : outputFields) {
  381. FieldDataWrapper wrapper = queryResultsWrapper.getFieldWrapper(fieldName);
  382. System.out.println("Query data of " + fieldName + ", row count: " + wrapper.getRowCount());
  383. System.out.println(wrapper.getFieldData());
  384. assertEquals(nq, wrapper.getFieldData().size());
  385. if (fieldName.compareTo(field1Name) == 0) {
  386. List<?> out = queryResultsWrapper.getFieldWrapper(field1Name).getFieldData();
  387. assertEquals(nq, out.size());
  388. for (Object o : out) {
  389. long id = (Long) o;
  390. assertTrue(queryIDs.contains(id));
  391. }
  392. }
  393. }
  394. // Note: the query() return vectors are not in same sequence to the input
  395. // here we cannot compare vector one by one
  396. // the boolean also cannot be compared
  397. if (outputFields.contains(field2Name)) {
  398. assertTrue(queryResultsWrapper.getFieldWrapper(field2Name).isVectorField());
  399. List<?> out = queryResultsWrapper.getFieldWrapper(field2Name).getFieldData();
  400. assertEquals(nq, out.size());
  401. }
  402. if (outputFields.contains(field3Name)) {
  403. List<?> out = queryResultsWrapper.getFieldWrapper(field3Name).getFieldData();
  404. assertEquals(nq, out.size());
  405. }
  406. if (outputFields.contains(field4Name)) {
  407. List<?> out = queryResultsWrapper.getFieldWrapper(field4Name).getFieldData();
  408. assertEquals(nq, out.size());
  409. for (Object o : out) {
  410. double d = (Double) o;
  411. assertTrue(compareWeights.contains(d));
  412. }
  413. }
  414. // query with offset and limit
  415. int queryLimit = 5;
  416. expr = field5Name + " > 1"; // "age > 1"
  417. queryParam = QueryParam.newBuilder()
  418. .withCollectionName(randomCollectionName)
  419. .withExpr(expr)
  420. .withOffset(100L)
  421. .withLimit((long) queryLimit)
  422. .build();
  423. queryR = client.query(queryParam);
  424. assertEquals(R.Status.Success.getCode(), queryR.getStatus().intValue());
  425. queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
  426. // we didn't set the output fields, only primary key field is returned
  427. List<?> out = queryResultsWrapper.getFieldWrapper(field1Name).getFieldData();
  428. assertEquals(queryLimit, out.size());
  429. // pick some vectors to search
  430. List<Long> targetVectorIDs = new ArrayList<>();
  431. List<List<Float>> targetVectors = new ArrayList<>();
  432. for (int i = randomIndex; i < randomIndex + nq; ++i) {
  433. targetVectorIDs.add(ids.get(i));
  434. targetVectors.add(vectors.get(i));
  435. }
  436. int topK = 5;
  437. SearchParam searchParam = SearchParam.newBuilder()
  438. .withCollectionName(randomCollectionName)
  439. .withMetricType(MetricType.L2)
  440. .withTopK(topK)
  441. .withVectors(targetVectors)
  442. .withVectorFieldName(field2Name)
  443. .withParams("{\"ef\":64}")
  444. .addOutField(field4Name)
  445. .build();
  446. R<SearchResults> searchR = client.search(searchParam);
  447. // System.out.println(searchR);
  448. assertEquals(R.Status.Success.getCode(), searchR.getStatus().intValue());
  449. // verify the search result
  450. SearchResultsWrapper results = new SearchResultsWrapper(searchR.getData().getResults());
  451. for (int i = 0; i < targetVectors.size(); ++i) {
  452. List<SearchResultsWrapper.IDScore> scores = results.getIDScore(i);
  453. System.out.println("The result of No." + i + " target vector(ID = " + targetVectorIDs.get(i) + "):");
  454. System.out.println(scores);
  455. assertEquals(targetVectorIDs.get(i).longValue(), scores.get(0).getLongID());
  456. }
  457. List<?> fieldData = results.getFieldData(field4Name, 0);
  458. assertEquals(topK, fieldData.size());
  459. fieldData = results.getFieldData(field4Name, nq - 1);
  460. assertEquals(topK, fieldData.size());
  461. // release collection
  462. ReleaseCollectionParam releaseCollectionParam = ReleaseCollectionParam.newBuilder()
  463. .withCollectionName(randomCollectionName).build();
  464. R<RpcStatus> releaseCollectionR = client.releaseCollection(releaseCollectionParam);
  465. assertEquals(R.Status.Success.getCode(), releaseCollectionR.getStatus().intValue());
  466. // drop index
  467. DropIndexParam dropIndexParam = DropIndexParam.newBuilder()
  468. .withCollectionName(randomCollectionName)
  469. .withIndexName(indexParam.getIndexName())
  470. .build();
  471. R<RpcStatus> dropIndexR = client.dropIndex(dropIndexParam);
  472. assertEquals(R.Status.Success.getCode(), dropIndexR.getStatus().intValue());
  473. // drop collection
  474. DropCollectionParam dropParam = DropCollectionParam.newBuilder()
  475. .withCollectionName(randomCollectionName)
  476. .build();
  477. R<RpcStatus> dropR = client.dropCollection(dropParam);
  478. assertEquals(R.Status.Success.getCode(), dropR.getStatus().intValue());
  479. }
  480. @Test
  481. void testBinaryVectors() {
  482. String randomCollectionName = generator.generate(10);
  483. // collection schema
  484. String field1Name = "field1";
  485. String field2Name = "field2";
  486. FieldType field1 = FieldType.newBuilder()
  487. .withPrimaryKey(true)
  488. .withAutoID(true)
  489. .withDataType(DataType.Int64)
  490. .withName(field1Name)
  491. .withDescription("hello")
  492. .build();
  493. FieldType field2 = FieldType.newBuilder()
  494. .withDataType(DataType.BinaryVector)
  495. .withName(field2Name)
  496. .withDescription("world")
  497. .withDimension(dimension)
  498. .build();
  499. // create collection
  500. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  501. .withCollectionName(randomCollectionName)
  502. .withDescription("test")
  503. .addFieldType(field1)
  504. .addFieldType(field2)
  505. .build();
  506. R<RpcStatus> createR = client.createCollection(createParam);
  507. assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  508. // insert data
  509. int rowCount = 10000;
  510. List<ByteBuffer> vectors = generateBinaryVectors(rowCount);
  511. List<InsertParam.Field> fields = new ArrayList<>();
  512. // no need to provide id here since this field is auto_id
  513. fields.add(new InsertParam.Field(field2Name, vectors));
  514. InsertParam insertParam = InsertParam.newBuilder()
  515. .withCollectionName(randomCollectionName)
  516. .withFields(fields)
  517. .build();
  518. R<MutationResult> insertR = client.insert(insertParam);
  519. assertEquals(R.Status.Success.getCode(), insertR.getStatus().intValue());
  520. // System.out.println(insertR.getData());
  521. MutationResultWrapper insertResultWrapper = new MutationResultWrapper(insertR.getData());
  522. System.out.println(insertResultWrapper.getInsertCount() + " rows inserted");
  523. List<Long> ids = insertResultWrapper.getLongIDs();
  524. // System.out.println("Auto-generated ids: " + ids);
  525. // get collection statistics
  526. R<GetCollectionStatisticsResponse> statR = client.getCollectionStatistics(GetCollectionStatisticsParam
  527. .newBuilder()
  528. .withCollectionName(randomCollectionName)
  529. .withFlush(true)
  530. .build());
  531. assertEquals(R.Status.Success.getCode(), statR.getStatus().intValue());
  532. GetCollStatResponseWrapper stat = new GetCollStatResponseWrapper(statR.getData());
  533. System.out.println("Collection row count: " + stat.getRowCount());
  534. // create index
  535. CreateIndexParam indexParam2 = CreateIndexParam.newBuilder()
  536. .withCollectionName(randomCollectionName)
  537. .withFieldName(field2Name)
  538. .withIndexType(IndexType.BIN_IVF_FLAT)
  539. .withExtraParam("{\"nlist\":64}")
  540. .withMetricType(MetricType.JACCARD)
  541. .withSyncMode(Boolean.TRUE)
  542. .withSyncWaitingInterval(500L)
  543. .withSyncWaitingTimeout(30L)
  544. .build();
  545. R<RpcStatus> createIndexR2 = client.createIndex(indexParam2);
  546. assertEquals(R.Status.Success.getCode(), createIndexR2.getStatus().intValue());
  547. // load collection
  548. R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
  549. .withCollectionName(randomCollectionName)
  550. .build());
  551. assertEquals(R.Status.Success.getCode(), loadR.getStatus().intValue());
  552. // search with BIN_FLAT index
  553. int searchTarget = 99;
  554. List<ByteBuffer> oneVector = new ArrayList<>();
  555. oneVector.add(vectors.get(searchTarget));
  556. SearchParam searchOneParam = SearchParam.newBuilder()
  557. .withCollectionName(randomCollectionName)
  558. .withMetricType(MetricType.JACCARD)
  559. .withTopK(5)
  560. .withVectors(oneVector)
  561. .withVectorFieldName(field2Name)
  562. .addOutField(field2Name)
  563. .build();
  564. R<SearchResults> searchOne = client.search(searchOneParam);
  565. assertEquals(R.Status.Success.getCode(), searchOne.getStatus().intValue());
  566. SearchResultsWrapper oneResult = new SearchResultsWrapper(searchOne.getData().getResults());
  567. List<SearchResultsWrapper.IDScore> oneScores = oneResult.getIDScore(0);
  568. System.out.println("The search result of id " + ids.get(searchTarget) + " with SUPERSTRUCTURE metric:");
  569. System.out.println(oneScores);
  570. // verify the output vector, the top1 item is equal to the target vector
  571. List<?> items = oneResult.getFieldData(field2Name, 0);
  572. Assertions.assertEquals(items.size(), 5);
  573. ByteBuffer firstItem = (ByteBuffer) items.get(0);
  574. for (int i = 0; i < firstItem.limit(); ++i) {
  575. Assertions.assertEquals(firstItem.get(i), vectors.get(searchTarget).get(i));
  576. }
  577. // release collection
  578. ReleaseCollectionParam releaseCollectionParam = ReleaseCollectionParam.newBuilder()
  579. .withCollectionName(randomCollectionName).build();
  580. R<RpcStatus> releaseCollectionR = client.releaseCollection(releaseCollectionParam);
  581. assertEquals(R.Status.Success.getCode(), releaseCollectionR.getStatus().intValue());
  582. DropIndexParam dropIndexParam = DropIndexParam.newBuilder()
  583. .withCollectionName(randomCollectionName).build();
  584. R<RpcStatus> dropIndexR = client.dropIndex(dropIndexParam);
  585. assertEquals(R.Status.Success.getCode(), dropIndexR.getStatus().intValue());
  586. // create index
  587. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  588. .withCollectionName(randomCollectionName)
  589. .withFieldName(field2Name)
  590. .withIndexName("abv")
  591. .withIndexType(IndexType.BIN_IVF_FLAT)
  592. .withMetricType(MetricType.HAMMING)
  593. .withExtraParam("{\"nlist\":64}")
  594. .build();
  595. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  596. assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  597. // get index description
  598. DescribeIndexParam descIndexParam = DescribeIndexParam.newBuilder()
  599. .withCollectionName(randomCollectionName)
  600. .withIndexName(indexParam.getIndexName())
  601. .build();
  602. R<DescribeIndexResponse> descIndexR = client.describeIndex(descIndexParam);
  603. assertEquals(R.Status.Success.getCode(), descIndexR.getStatus().intValue());
  604. // load collection
  605. R<RpcStatus> loadR2 = client.loadCollection(LoadCollectionParam.newBuilder()
  606. .withCollectionName(randomCollectionName)
  607. .build());
  608. assertEquals(R.Status.Success.getCode(), loadR2.getStatus().intValue());
  609. // pick some vectors to search with index
  610. int nq = 5;
  611. List<Long> targetVectorIDs = new ArrayList<>();
  612. List<ByteBuffer> targetVectors = new ArrayList<>();
  613. Random ran = new Random();
  614. int randomIndex = ran.nextInt(rowCount - nq);
  615. for (int i = randomIndex; i < randomIndex + nq; ++i) {
  616. targetVectorIDs.add(ids.get(i));
  617. targetVectors.add(vectors.get(i));
  618. }
  619. int topK = 5;
  620. SearchParam searchParam = SearchParam.newBuilder()
  621. .withCollectionName(randomCollectionName)
  622. .withMetricType(MetricType.HAMMING)
  623. .withTopK(topK)
  624. .withVectors(targetVectors)
  625. .withVectorFieldName(field2Name)
  626. .withParams("{\"nprobe\":8}")
  627. .build();
  628. R<SearchResults> searchR = client.search(searchParam);
  629. // System.out.println(searchR);
  630. assertEquals(R.Status.Success.getCode(), searchR.getStatus().intValue());
  631. // verify the search result
  632. SearchResultsWrapper results = new SearchResultsWrapper(searchR.getData().getResults());
  633. for (int i = 0; i < targetVectors.size(); ++i) {
  634. List<SearchResultsWrapper.IDScore> scores = results.getIDScore(i);
  635. System.out.println("The result of No." + i + " target vector(ID = " + targetVectorIDs.get(i) + "):");
  636. System.out.println(scores);
  637. assertEquals(targetVectorIDs.get(i).longValue(), scores.get(0).getLongID());
  638. }
  639. // drop collection
  640. DropCollectionParam dropParam = DropCollectionParam.newBuilder()
  641. .withCollectionName(randomCollectionName)
  642. .build();
  643. R<RpcStatus> dropR = client.dropCollection(dropParam);
  644. assertEquals(R.Status.Success.getCode(), dropR.getStatus().intValue());
  645. }
  646. @Test
  647. void testAsyncMethods() {
  648. String randomCollectionName = generator.generate(10);
  649. // collection schema
  650. String field1Name = "long_field";
  651. String field2Name = "vec_field";
  652. List<FieldType> fieldsSchema = new ArrayList<>();
  653. fieldsSchema.add(FieldType.newBuilder()
  654. .withPrimaryKey(true)
  655. .withAutoID(true)
  656. .withDataType(DataType.Int64)
  657. .withName(field1Name)
  658. .withDescription("identity")
  659. .build());
  660. fieldsSchema.add(FieldType.newBuilder()
  661. .withDataType(DataType.FloatVector)
  662. .withName(field2Name)
  663. .withDescription("face")
  664. .withDimension(dimension)
  665. .build());
  666. // create collection
  667. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  668. .withCollectionName(randomCollectionName)
  669. .withDescription("test")
  670. .withFieldTypes(fieldsSchema)
  671. .build();
  672. R<RpcStatus> createR = client.createCollection(createParam);
  673. assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  674. // insert async
  675. List<ListenableFuture<R<MutationResult>>> futureResponses = new ArrayList<>();
  676. int rowCount = 1000;
  677. for (long i = 0L; i < 10; ++i) {
  678. List<List<Float>> vectors = normalizeFloatVectors(generateFloatVectors(rowCount));
  679. List<InsertParam.Field> fieldsInsert = new ArrayList<>();
  680. fieldsInsert.add(new InsertParam.Field(field2Name, vectors));
  681. InsertParam insertParam = InsertParam.newBuilder()
  682. .withCollectionName(randomCollectionName)
  683. .withFields(fieldsInsert)
  684. .build();
  685. ListenableFuture<R<MutationResult>> insertFuture = client.insertAsync(insertParam);
  686. futureResponses.add(insertFuture);
  687. }
  688. // get insert result
  689. List<Long> queryIDs = new ArrayList<>();
  690. for (ListenableFuture<R<MutationResult>> response : futureResponses) {
  691. try {
  692. R<MutationResult> insertR = response.get();
  693. assertEquals(R.Status.Success.getCode(), insertR.getStatus().intValue());
  694. MutationResultWrapper wrapper = new MutationResultWrapper(insertR.getData());
  695. queryIDs.add(wrapper.getLongIDs().get(0));
  696. } catch (ExecutionException | InterruptedException e) {
  697. System.out.println("failed to insert:" + e.getMessage());
  698. return;
  699. }
  700. }
  701. // get collection statistics
  702. R<GetCollectionStatisticsResponse> statR = client.getCollectionStatistics(GetCollectionStatisticsParam
  703. .newBuilder()
  704. .withCollectionName(randomCollectionName)
  705. .withFlush(true)
  706. .build());
  707. assertEquals(R.Status.Success.getCode(), statR.getStatus().intValue());
  708. GetCollStatResponseWrapper stat = new GetCollStatResponseWrapper(statR.getData());
  709. System.out.println("Collection row count: " + stat.getRowCount());
  710. // create index
  711. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  712. .withCollectionName(randomCollectionName)
  713. .withFieldName(field2Name)
  714. .withIndexType(IndexType.FLAT)
  715. .withMetricType(MetricType.IP)
  716. .withSyncMode(Boolean.TRUE)
  717. .withSyncWaitingInterval(500L)
  718. .withSyncWaitingTimeout(30L)
  719. .build();
  720. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  721. assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  722. // load collection
  723. R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
  724. .withCollectionName(randomCollectionName)
  725. .build());
  726. assertEquals(R.Status.Success.getCode(), loadR.getStatus().intValue());
  727. // search async
  728. List<List<Float>> targetVectors = normalizeFloatVectors(generateFloatVectors(2));
  729. int topK = 5;
  730. SearchParam searchParam = SearchParam.newBuilder()
  731. .withCollectionName(randomCollectionName)
  732. .withMetricType(MetricType.IP)
  733. .withTopK(topK)
  734. .withVectors(targetVectors)
  735. .withVectorFieldName(field2Name)
  736. .build();
  737. ListenableFuture<R<SearchResults>> searchFuture = client.searchAsync(searchParam);
  738. // query async
  739. String expr = field1Name + " in " + queryIDs.toString();
  740. List<String> outputFields = Arrays.asList(field1Name, field2Name);
  741. QueryParam queryParam = QueryParam.newBuilder()
  742. .withCollectionName(randomCollectionName)
  743. .withExpr(expr)
  744. .withOutFields(outputFields)
  745. .build();
  746. ListenableFuture<R<QueryResults>> queryFuture = client.queryAsync(queryParam);
  747. try {
  748. // get search results
  749. R<SearchResults> searchR = searchFuture.get();
  750. assertEquals(R.Status.Success.getCode(), searchR.getStatus().intValue());
  751. // verify search result
  752. SearchResultsWrapper results = new SearchResultsWrapper(searchR.getData().getResults());
  753. System.out.println("Search results:");
  754. for (int i = 0; i < targetVectors.size(); ++i) {
  755. List<SearchResultsWrapper.IDScore> scores = results.getIDScore(i);
  756. assertEquals(topK, scores.size());
  757. System.out.println(scores.toString());
  758. }
  759. // get query results
  760. R<QueryResults> queryR = queryFuture.get();
  761. assertEquals(R.Status.Success.getCode(), queryR.getStatus().intValue());
  762. // verify query result
  763. QueryResultsWrapper queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
  764. for (String fieldName : outputFields) {
  765. FieldDataWrapper wrapper = queryResultsWrapper.getFieldWrapper(fieldName);
  766. System.out.println("Query data of " + fieldName + ", row count: " + wrapper.getRowCount());
  767. System.out.println(wrapper.getFieldData());
  768. assertEquals(queryIDs.size(), wrapper.getFieldData().size());
  769. }
  770. } catch (InterruptedException | ExecutionException e) {
  771. e.printStackTrace();
  772. }
  773. // drop collection
  774. DropCollectionParam dropParam = DropCollectionParam.newBuilder()
  775. .withCollectionName(randomCollectionName)
  776. .build();
  777. R<RpcStatus> dropR = client.dropCollection(dropParam);
  778. assertEquals(R.Status.Success.getCode(), dropR.getStatus().intValue());
  779. }
  780. // this case can be executed when the milvus image of version 2.1 is published.
  781. @Test
  782. void testCredential() {
  783. String randomCollectionName = generator.generate(10);
  784. // collection schema
  785. String field1Name = "long_field";
  786. String field2Name = "vec_field";
  787. List<FieldType> fieldsSchema = new ArrayList<>();
  788. fieldsSchema.add(FieldType.newBuilder()
  789. .withPrimaryKey(true)
  790. .withAutoID(false)
  791. .withDataType(DataType.Int64)
  792. .withName(field1Name)
  793. .withDescription("identity")
  794. .build());
  795. fieldsSchema.add(FieldType.newBuilder()
  796. .withDataType(DataType.FloatVector)
  797. .withName(field2Name)
  798. .withDescription("face")
  799. .withDimension(dimension)
  800. .build());
  801. // create collection
  802. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  803. .withCollectionName(randomCollectionName)
  804. .withDescription("test")
  805. .withFieldTypes(fieldsSchema)
  806. .build();
  807. R<RpcStatus> createR = client.createCollection(createParam);
  808. // create index
  809. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  810. .withCollectionName(randomCollectionName)
  811. .withFieldName(field2Name)
  812. .withIndexType(IndexType.IVF_FLAT)
  813. .withIndexName("xxx")
  814. .withMetricType(MetricType.L2)
  815. .withExtraParam("{\"nlist\":256}")
  816. .withSyncMode(Boolean.TRUE)
  817. .withSyncWaitingInterval(500L)
  818. .withSyncWaitingTimeout(30L)
  819. .build();
  820. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  821. client.getIndexState(GetIndexStateParam.newBuilder()
  822. .withCollectionName(randomCollectionName)
  823. .withIndexName(indexParam.getIndexName())
  824. .build());
  825. R<RpcStatus> dropIndexR = client.dropIndex(DropIndexParam.newBuilder()
  826. .withCollectionName(randomCollectionName)
  827. .withIndexName(indexParam.getIndexName())
  828. .build());
  829. client.dropCollection(DropCollectionParam.newBuilder().withCollectionName(randomCollectionName).build());
  830. }
  831. @Test
  832. void testStringField() {
  833. String randomCollectionName = generator.generate(10);
  834. // collection schema
  835. String field1Name = "str_id";
  836. String field2Name = "vec_field";
  837. String field3Name = "str_field";
  838. String field4Name = "int_field";
  839. List<FieldType> fieldsSchema = new ArrayList<>();
  840. fieldsSchema.add(FieldType.newBuilder()
  841. .withPrimaryKey(true)
  842. .withAutoID(false)
  843. .withDataType(DataType.VarChar)
  844. .withName(field1Name)
  845. .withMaxLength(32)
  846. .withDescription("string identity")
  847. .build());
  848. fieldsSchema.add(FieldType.newBuilder()
  849. .withDataType(DataType.FloatVector)
  850. .withName(field2Name)
  851. .withDescription("face")
  852. .withDimension(dimension)
  853. .build());
  854. fieldsSchema.add(FieldType.newBuilder()
  855. .withDataType(DataType.VarChar)
  856. .withName(field3Name)
  857. .withMaxLength(32)
  858. .withDescription("comment")
  859. .build());
  860. fieldsSchema.add(FieldType.newBuilder()
  861. .withDataType(DataType.Int64)
  862. .withName(field4Name)
  863. .withDescription("sequence")
  864. .build());
  865. // create collection
  866. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  867. .withCollectionName(randomCollectionName)
  868. .withDescription("test")
  869. .withFieldTypes(fieldsSchema)
  870. .build();
  871. R<RpcStatus> createR = client.createCollection(createParam);
  872. assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  873. R<DescribeCollectionResponse> response = client.describeCollection(DescribeCollectionParam.newBuilder()
  874. .withCollectionName(randomCollectionName)
  875. .build());
  876. DescCollResponseWrapper desc = new DescCollResponseWrapper(response.getData());
  877. System.out.println(desc.toString());
  878. // insert data
  879. int rowCount = 10000;
  880. List<String> ids = new ArrayList<>();
  881. List<String> comments = new ArrayList<>();
  882. List<Long> sequences = new ArrayList<>();
  883. for (long i = 0L; i < rowCount; ++i) {
  884. ids.add(generator.generate(8));
  885. comments.add(generator.generate(8));
  886. sequences.add(i);
  887. }
  888. List<List<Float>> vectors = generateFloatVectors(rowCount);
  889. List<InsertParam.Field> fieldsInsert = new ArrayList<>();
  890. fieldsInsert.add(new InsertParam.Field(field1Name, ids));
  891. fieldsInsert.add(new InsertParam.Field(field3Name, comments));
  892. fieldsInsert.add(new InsertParam.Field(field2Name, vectors));
  893. fieldsInsert.add(new InsertParam.Field(field4Name, sequences));
  894. InsertParam insertParam = InsertParam.newBuilder()
  895. .withCollectionName(randomCollectionName)
  896. .withFields(fieldsInsert)
  897. .build();
  898. R<MutationResult> insertR = client.withTimeout(10, TimeUnit.SECONDS).insert(insertParam);
  899. assertEquals(R.Status.Success.getCode(), insertR.getStatus().intValue());
  900. MutationResultWrapper insertResultWrapper = new MutationResultWrapper(insertR.getData());
  901. System.out.println(insertResultWrapper.getInsertCount() + " rows inserted");
  902. // get collection statistics
  903. R<GetCollectionStatisticsResponse> statR = client.getCollectionStatistics(GetCollectionStatisticsParam
  904. .newBuilder()
  905. .withCollectionName(randomCollectionName)
  906. .withFlush(true)
  907. .build());
  908. assertEquals(R.Status.Success.getCode(), statR.getStatus().intValue());
  909. GetCollStatResponseWrapper stat = new GetCollStatResponseWrapper(statR.getData());
  910. System.out.println("Collection row count: " + stat.getRowCount());
  911. // create index
  912. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  913. .withCollectionName(randomCollectionName)
  914. .withFieldName(field3Name)
  915. .withIndexName("stridx")
  916. .withIndexType(IndexType.TRIE)
  917. .withSyncMode(Boolean.TRUE)
  918. .withSyncWaitingInterval(500L)
  919. .withSyncWaitingTimeout(30L)
  920. .build();
  921. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  922. assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  923. // get index description
  924. DescribeIndexParam descIndexParam = DescribeIndexParam.newBuilder()
  925. .withCollectionName(randomCollectionName)
  926. .withIndexName(indexParam.getIndexName())
  927. .build();
  928. R<DescribeIndexResponse> descIndexR = client.describeIndex(descIndexParam);
  929. assertEquals(R.Status.Success.getCode(), descIndexR.getStatus().intValue());
  930. // create index
  931. CreateIndexParam indexParam2 = CreateIndexParam.newBuilder()
  932. .withCollectionName(randomCollectionName)
  933. .withFieldName(field2Name)
  934. .withIndexType(IndexType.FLAT)
  935. .withMetricType(MetricType.IP)
  936. .withSyncMode(Boolean.TRUE)
  937. .withSyncWaitingInterval(500L)
  938. .withSyncWaitingTimeout(30L)
  939. .build();
  940. R<RpcStatus> createIndexR2 = client.createIndex(indexParam2);
  941. assertEquals(R.Status.Success.getCode(), createIndexR2.getStatus().intValue());
  942. // load collection
  943. R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
  944. .withCollectionName(randomCollectionName)
  945. .build());
  946. assertEquals(R.Status.Success.getCode(), loadR.getStatus().intValue());
  947. // query vectors to verify
  948. List<Long> queryItems = new ArrayList<>();
  949. List<String> queryIds = new ArrayList<>();
  950. int nq = 5;
  951. Random ran = new Random();
  952. int randomIndex = ran.nextInt(rowCount - nq);
  953. for (int i = randomIndex; i < randomIndex + nq; ++i) {
  954. queryIds.add(ids.get(i));
  955. queryItems.add(sequences.get(i));
  956. }
  957. String expr = field4Name + " in " + queryItems.toString();
  958. List<String> outputFields = Arrays.asList(field1Name, field3Name);
  959. QueryParam queryParam = QueryParam.newBuilder()
  960. .withCollectionName(randomCollectionName)
  961. .withExpr(expr)
  962. .withOutFields(outputFields)
  963. .build();
  964. R<QueryResults> queryR = client.query(queryParam);
  965. assertEquals(R.Status.Success.getCode(), queryR.getStatus().intValue());
  966. // verify query result
  967. QueryResultsWrapper queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
  968. for (String fieldName : outputFields) {
  969. FieldDataWrapper wrapper = queryResultsWrapper.getFieldWrapper(fieldName);
  970. System.out.println("Query data of " + fieldName + ", row count: " + wrapper.getRowCount());
  971. System.out.println(wrapper.getFieldData());
  972. assertEquals(nq, wrapper.getFieldData().size());
  973. if (fieldName.compareTo(field1Name) == 0) {
  974. List<?> out = queryResultsWrapper.getFieldWrapper(field1Name).getFieldData();
  975. assertEquals(nq, out.size());
  976. for (Object o : out) {
  977. String id = (String) o;
  978. assertTrue(queryIds.contains(id));
  979. }
  980. }
  981. }
  982. // search
  983. int topK = 5;
  984. List<List<Float>> targetVectors = new ArrayList<>();
  985. for (Long seq : queryItems) {
  986. targetVectors.add(vectors.get(seq.intValue()));
  987. }
  988. SearchParam searchParam = SearchParam.newBuilder()
  989. .withCollectionName(randomCollectionName)
  990. .withMetricType(MetricType.IP)
  991. .withTopK(topK)
  992. .withVectors(targetVectors)
  993. .withVectorFieldName(field2Name)
  994. .addOutField(field4Name)
  995. .build();
  996. R<SearchResults> searchR = client.search(searchParam);
  997. assertEquals(R.Status.Success.getCode(), searchR.getStatus().intValue());
  998. // verify the search result
  999. SearchResultsWrapper results = new SearchResultsWrapper(searchR.getData().getResults());
  1000. for (int i = 0; i < targetVectors.size(); ++i) {
  1001. List<SearchResultsWrapper.IDScore> scores = results.getIDScore(i);
  1002. System.out.println("The result of No." + i + " target vector(ID = " + queryIds.get(i) + "):");
  1003. System.out.println(scores);
  1004. }
  1005. client.dropCollection(DropCollectionParam.newBuilder().withCollectionName(randomCollectionName).build());
  1006. }
  1007. private static void testIndex(String collectionName, String fieldName,
  1008. IndexType type, MetricType metric,
  1009. String params, Boolean syncMode) {
  1010. // create index
  1011. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  1012. .withCollectionName(collectionName)
  1013. .withFieldName(fieldName)
  1014. .withIndexName("index")
  1015. .withIndexType(type)
  1016. .withMetricType(metric)
  1017. .withExtraParam(params)
  1018. .withSyncMode(syncMode)
  1019. .build();
  1020. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  1021. assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  1022. // drop index
  1023. DropIndexParam dropIndexParam = DropIndexParam.newBuilder()
  1024. .withCollectionName(collectionName)
  1025. .withIndexName(indexParam.getIndexName())
  1026. .build();
  1027. R<RpcStatus> dropIndexR = client.dropIndex(dropIndexParam);
  1028. assertEquals(R.Status.Success.getCode(), dropIndexR.getStatus().intValue());
  1029. }
  1030. @Test
  1031. void testFloatVectorIndex() {
  1032. String randomCollectionName = generator.generate(10);
  1033. // collection schema
  1034. String field1Name = "idg_field";
  1035. String field2Name = "vec_field";
  1036. List<FieldType> fieldsSchema = new ArrayList<>();
  1037. fieldsSchema.add(FieldType.newBuilder()
  1038. .withPrimaryKey(true)
  1039. .withAutoID(false)
  1040. .withDataType(DataType.Int64)
  1041. .withName(field1Name)
  1042. .build());
  1043. fieldsSchema.add(FieldType.newBuilder()
  1044. .withDataType(DataType.FloatVector)
  1045. .withName(field2Name)
  1046. .withDimension(dimension)
  1047. .build());
  1048. // create collection
  1049. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  1050. .withCollectionName(randomCollectionName)
  1051. .withDescription("test")
  1052. .withFieldTypes(fieldsSchema)
  1053. .build();
  1054. R<RpcStatus> createR = client.createCollection(createParam);
  1055. assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  1056. // test all supported indexes
  1057. Map<IndexType, String> indexTypes = new HashMap<>();
  1058. indexTypes.put(IndexType.FLAT, "{}");
  1059. indexTypes.put(IndexType.IVF_FLAT, "{\"nlist\":128}");
  1060. indexTypes.put(IndexType.IVF_SQ8, "{\"nlist\":128}");
  1061. indexTypes.put(IndexType.IVF_PQ, "{\"nlist\":128, \"m\":16, \"nbits\":8}");
  1062. indexTypes.put(IndexType.ANNOY, "{\"n_trees\":16}");
  1063. indexTypes.put(IndexType.HNSW, "{\"M\":16,\"efConstruction\":64}");
  1064. indexTypes.put(IndexType.RHNSW_FLAT, "{\"M\":16,\"efConstruction\":64}");
  1065. indexTypes.put(IndexType.RHNSW_PQ, "{\"M\":16,\"efConstruction\":64, \"PQM\":16}");
  1066. indexTypes.put(IndexType.RHNSW_SQ, "{\"M\":16,\"efConstruction\":64}");
  1067. List<MetricType> metricTypes = new ArrayList<>();
  1068. metricTypes.add(MetricType.L2);
  1069. metricTypes.add(MetricType.IP);
  1070. for (IndexType type : indexTypes.keySet()) {
  1071. for (MetricType metric : metricTypes) {
  1072. testIndex(randomCollectionName, field2Name, type, metric, indexTypes.get(type), Boolean.TRUE);
  1073. testIndex(randomCollectionName, field2Name, type, metric, indexTypes.get(type), Boolean.FALSE);
  1074. }
  1075. }
  1076. client.dropCollection(DropCollectionParam.newBuilder().withCollectionName(randomCollectionName).build());
  1077. }
  1078. @Test
  1079. void testBinaryVectorIndex() {
  1080. String randomCollectionName = generator.generate(10);
  1081. // collection schema
  1082. String field1Name = "id_field";
  1083. String field2Name = "vector_field";
  1084. FieldType field1 = FieldType.newBuilder()
  1085. .withPrimaryKey(true)
  1086. .withAutoID(true)
  1087. .withDataType(DataType.Int64)
  1088. .withName(field1Name)
  1089. .build();
  1090. FieldType field2 = FieldType.newBuilder()
  1091. .withDataType(DataType.BinaryVector)
  1092. .withName(field2Name)
  1093. .withDimension(dimension)
  1094. .build();
  1095. // create collection
  1096. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  1097. .withCollectionName(randomCollectionName)
  1098. .withDescription("test")
  1099. .addFieldType(field1)
  1100. .addFieldType(field2)
  1101. .build();
  1102. R<RpcStatus> createR = client.createCollection(createParam);
  1103. assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  1104. // test all supported indexes
  1105. List<MetricType> flatMetricTypes = new ArrayList<>();
  1106. flatMetricTypes.add(MetricType.SUBSTRUCTURE);
  1107. flatMetricTypes.add(MetricType.SUPERSTRUCTURE);
  1108. for (MetricType metric : flatMetricTypes) {
  1109. testIndex(randomCollectionName, field2Name, IndexType.BIN_FLAT, metric, "{}", Boolean.TRUE);
  1110. testIndex(randomCollectionName, field2Name, IndexType.BIN_FLAT, metric, "{}", Boolean.FALSE);
  1111. }
  1112. List<MetricType> ivfMetricTypes = new ArrayList<>();
  1113. ivfMetricTypes.add(MetricType.HAMMING);
  1114. ivfMetricTypes.add(MetricType.JACCARD);
  1115. ivfMetricTypes.add(MetricType.TANIMOTO);
  1116. for (MetricType metric : ivfMetricTypes) {
  1117. testIndex(randomCollectionName, field2Name, IndexType.BIN_IVF_FLAT, metric, "{\"nlist\":128}", Boolean.TRUE);
  1118. testIndex(randomCollectionName, field2Name, IndexType.BIN_IVF_FLAT, metric, "{\"nlist\":128}", Boolean.FALSE);
  1119. }
  1120. client.dropCollection(DropCollectionParam.newBuilder().withCollectionName(randomCollectionName).build());
  1121. }
  1122. @Test
  1123. void testDynamicField() {
  1124. String randomCollectionName = generator.generate(10);
  1125. // collection schema
  1126. String field1Name = "id_field";
  1127. String field2Name = "vec_field";
  1128. String field3Name = "json_field";
  1129. List<FieldType> fieldsSchema = new ArrayList<>();
  1130. fieldsSchema.add(FieldType.newBuilder()
  1131. .withPrimaryKey(true)
  1132. .withAutoID(false)
  1133. .withDataType(DataType.Int64)
  1134. .withName(field1Name)
  1135. .withDescription("identity")
  1136. .build());
  1137. fieldsSchema.add(FieldType.newBuilder()
  1138. .withDataType(DataType.FloatVector)
  1139. .withName(field2Name)
  1140. .withDescription("face")
  1141. .withDimension(dimension)
  1142. .build());
  1143. fieldsSchema.add(FieldType.newBuilder()
  1144. .withDataType(DataType.JSON)
  1145. .withName(field3Name)
  1146. .withDescription("info")
  1147. .build());
  1148. // create collection
  1149. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  1150. .withCollectionName(randomCollectionName)
  1151. .withFieldTypes(fieldsSchema)
  1152. .withEnableDynamicField(true)
  1153. .build();
  1154. R<RpcStatus> createR = client.createCollection(createParam);
  1155. assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  1156. R<DescribeCollectionResponse> response = client.describeCollection(DescribeCollectionParam.newBuilder()
  1157. .withCollectionName(randomCollectionName)
  1158. .build());
  1159. DescCollResponseWrapper desc = new DescCollResponseWrapper(response.getData());
  1160. System.out.println(desc.toString());
  1161. // create index
  1162. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  1163. .withCollectionName(randomCollectionName)
  1164. .withFieldName(field2Name)
  1165. .withIndexName("abv")
  1166. .withIndexType(IndexType.FLAT)
  1167. .withMetricType(MetricType.L2)
  1168. .withExtraParam("{}")
  1169. .build();
  1170. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  1171. assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  1172. // load collection
  1173. R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
  1174. .withCollectionName(randomCollectionName)
  1175. .build());
  1176. assertEquals(R.Status.Success.getCode(), loadR.getStatus().intValue());
  1177. int rowCount = 10;
  1178. // insert data by row-based
  1179. List<JSONObject> rows = new ArrayList<>();
  1180. for (long i = 0L; i < rowCount; ++i) {
  1181. JSONObject row = new JSONObject();
  1182. row.put(field1Name, i);
  1183. row.put(field2Name, generateFloatVectors(1).get(0));
  1184. // JSON field
  1185. JSONObject info = new JSONObject();
  1186. info.put("row-based-info", i);
  1187. row.put(field3Name, info);
  1188. // extra meta is automatically stored in dynamic field
  1189. row.put("extra_meta", i % 3 == 0);
  1190. row.put(generator.generate(5), 100);
  1191. rows.add(row);
  1192. }
  1193. InsertParam insertRowParam = InsertParam.newBuilder()
  1194. .withCollectionName(randomCollectionName)
  1195. .withRows(rows)
  1196. .build();
  1197. R<MutationResult> insertRowResp = client.insert(insertRowParam);
  1198. assertEquals(R.Status.Success.getCode(), insertRowResp.getStatus().intValue());
  1199. System.out.println(rowCount + " rows inserted");
  1200. // insert data by column-based
  1201. List<Long> ids = new ArrayList<>();
  1202. List<JSONObject> infos = new ArrayList<>();
  1203. for (long i = 0L; i < rowCount; ++i) {
  1204. ids.add(rowCount + i);
  1205. JSONObject obj = new JSONObject();
  1206. obj.put("column-based-info", i);
  1207. obj.put(generator.generate(5), i);
  1208. infos.add(obj);
  1209. }
  1210. List<List<Float>> vectors = generateFloatVectors(rowCount);
  1211. List<InsertParam.Field> fieldsInsert = new ArrayList<>();
  1212. fieldsInsert.add(new InsertParam.Field(field1Name, ids));
  1213. fieldsInsert.add(new InsertParam.Field(field2Name, vectors));
  1214. fieldsInsert.add(new InsertParam.Field(field3Name, infos));
  1215. InsertParam insertColumnsParam = InsertParam.newBuilder()
  1216. .withCollectionName(randomCollectionName)
  1217. .withFields(fieldsInsert)
  1218. .build();
  1219. R<MutationResult> insertColumnResp = client.insert(insertColumnsParam);
  1220. assertEquals(R.Status.Success.getCode(), insertColumnResp.getStatus().intValue());
  1221. System.out.println(rowCount + " rows inserted");
  1222. // get collection statistics
  1223. R<GetCollectionStatisticsResponse> statR = client.getCollectionStatistics(GetCollectionStatisticsParam
  1224. .newBuilder()
  1225. .withCollectionName(randomCollectionName)
  1226. .withFlush(true)
  1227. .build());
  1228. assertEquals(R.Status.Success.getCode(), statR.getStatus().intValue());
  1229. GetCollStatResponseWrapper stat = new GetCollStatResponseWrapper(statR.getData());
  1230. System.out.println("Collection row count: " + stat.getRowCount());
  1231. // retrieve rows
  1232. String expr = "extra_meta == true";
  1233. List<String> outputFields = Arrays.asList(field3Name, "extra_meta");
  1234. QueryParam queryParam = QueryParam.newBuilder()
  1235. .withCollectionName(randomCollectionName)
  1236. .withExpr(expr)
  1237. .withOutFields(outputFields)
  1238. .build();
  1239. R<QueryResults> queryR = client.query(queryParam);
  1240. assertEquals(R.Status.Success.getCode(), queryR.getStatus().intValue());
  1241. QueryResultsWrapper queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
  1242. List<QueryResultsWrapper.RowRecord> records = queryResultsWrapper.getRowRecords();
  1243. System.out.println("Query results:");
  1244. for (QueryResultsWrapper.RowRecord record:records) {
  1245. System.out.println(record);
  1246. Object extraMeta = record.get("extra_meta");
  1247. if (extraMeta != null) {
  1248. System.out.println("'extra_meta' is from dynamic field, value: " + extraMeta);
  1249. }
  1250. }
  1251. // search
  1252. List<List<Float>> targetVectors = generateFloatVectors(2);
  1253. int topK = 5;
  1254. SearchParam searchParam = SearchParam.newBuilder()
  1255. .withCollectionName(randomCollectionName)
  1256. .withMetricType(MetricType.L2)
  1257. .withTopK(topK)
  1258. .withVectors(targetVectors)
  1259. .withVectorFieldName(field2Name)
  1260. .withParams("{}")
  1261. .withOutFields(outputFields)
  1262. .build();
  1263. R<SearchResults> searchR = client.search(searchParam);
  1264. assertEquals(R.Status.Success.getCode(), searchR.getStatus().intValue());
  1265. // verify the search result
  1266. SearchResultsWrapper results = new SearchResultsWrapper(searchR.getData().getResults());
  1267. for (int i = 0; i < targetVectors.size(); ++i) {
  1268. List<SearchResultsWrapper.IDScore> scores = results.getIDScore(i);
  1269. System.out.println("The result of No." + i + " target vector:");
  1270. for (SearchResultsWrapper.IDScore score:scores) {
  1271. System.out.println(score);
  1272. Object extraMeta = score.get("extra_meta");
  1273. if (extraMeta != null) {
  1274. System.out.println("'extra_meta' is from dynamic field, value: " + extraMeta);
  1275. }
  1276. }
  1277. }
  1278. // drop collection
  1279. R<RpcStatus> dropR = client.dropCollection(DropCollectionParam.newBuilder()
  1280. .withCollectionName(randomCollectionName)
  1281. .build());
  1282. assertEquals(R.Status.Success.getCode(), dropR.getStatus().intValue());
  1283. }
  1284. @Test
  1285. void testHighLevelGet() {
  1286. // collection schema
  1287. String field1Name = "id_field";
  1288. String field2Name = "vector_field";
  1289. FieldType int64PrimaryField = FieldType.newBuilder()
  1290. .withPrimaryKey(true)
  1291. .withAutoID(false)
  1292. .withDataType(DataType.Int64)
  1293. .withName(field1Name)
  1294. .build();
  1295. FieldType varcharPrimaryField = FieldType.newBuilder()
  1296. .withPrimaryKey(true)
  1297. .withDataType(DataType.VarChar)
  1298. .withName(field1Name)
  1299. .withMaxLength(128)
  1300. .build();
  1301. FieldType vectorField = FieldType.newBuilder()
  1302. .withDataType(DataType.FloatVector)
  1303. .withName(field2Name)
  1304. .withDimension(dimension)
  1305. .build();
  1306. testCollectionHighLevelGet(int64PrimaryField, vectorField);
  1307. testCollectionHighLevelGet(varcharPrimaryField, vectorField);
  1308. }
  1309. @Test
  1310. void testHighLevelDelete() {
  1311. // collection schema
  1312. String field1Name = "id_field";
  1313. String field2Name = "vector_field";
  1314. FieldType int64PrimaryField = FieldType.newBuilder()
  1315. .withPrimaryKey(true)
  1316. .withAutoID(false)
  1317. .withDataType(DataType.Int64)
  1318. .withName(field1Name)
  1319. .build();
  1320. FieldType varcharPrimaryField = FieldType.newBuilder()
  1321. .withPrimaryKey(true)
  1322. .withDataType(DataType.VarChar)
  1323. .withName(field1Name)
  1324. .withMaxLength(128)
  1325. .build();
  1326. FieldType vectorField = FieldType.newBuilder()
  1327. .withDataType(DataType.FloatVector)
  1328. .withName(field2Name)
  1329. .withDimension(dimension)
  1330. .build();
  1331. testCollectionHighLevelDelete(int64PrimaryField, vectorField);
  1332. testCollectionHighLevelDelete(varcharPrimaryField, vectorField);
  1333. }
  1334. void testCollectionHighLevelGet(FieldType primaryField, FieldType vectorField) {
  1335. // create collection
  1336. String randomCollectionName = generator.generate(10);
  1337. highLevelCreateCollection(primaryField, vectorField, randomCollectionName);
  1338. // insert data
  1339. List<String> primaryIds = new ArrayList<>();
  1340. int rowCount = 10;
  1341. List<JSONObject> rows = new ArrayList<>();
  1342. for (long i = 0L; i < rowCount; ++i) {
  1343. JSONObject row = new JSONObject();
  1344. row.put(primaryField.getName(), primaryField.getDataType() == DataType.Int64 ? i : String.valueOf(i));
  1345. row.put(vectorField.getName(), generateFloatVectors(1).get(0));
  1346. rows.add(row);
  1347. primaryIds.add(String.valueOf(i));
  1348. }
  1349. InsertParam insertRowParam = InsertParam.newBuilder()
  1350. .withCollectionName(randomCollectionName)
  1351. .withRows(rows)
  1352. .build();
  1353. R<MutationResult> insertRowResp = client.insert(insertRowParam);
  1354. Assertions.assertEquals(R.Status.Success.getCode(), insertRowResp.getStatus().intValue());
  1355. testHighLevelGet(randomCollectionName, primaryIds);
  1356. client.dropCollection(DropCollectionParam.newBuilder().withCollectionName(randomCollectionName).build());
  1357. }
  1358. private static void highLevelCreateCollection(FieldType primaryField, FieldType vectorField, String randomCollectionName) {
  1359. CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
  1360. .withCollectionName(randomCollectionName)
  1361. .withDescription("test")
  1362. .addFieldType(primaryField)
  1363. .addFieldType(vectorField)
  1364. .build();
  1365. R<RpcStatus> createR = client.createCollection(createParam);
  1366. Assertions.assertEquals(R.Status.Success.getCode(), createR.getStatus().intValue());
  1367. // create index
  1368. CreateIndexParam indexParam = CreateIndexParam.newBuilder()
  1369. .withCollectionName(randomCollectionName)
  1370. .withFieldName(vectorField.getName())
  1371. .withIndexName("abv")
  1372. .withIndexType(IndexType.FLAT)
  1373. .withMetricType(MetricType.L2)
  1374. .withExtraParam("{}")
  1375. .build();
  1376. R<RpcStatus> createIndexR = client.createIndex(indexParam);
  1377. Assertions.assertEquals(R.Status.Success.getCode(), createIndexR.getStatus().intValue());
  1378. // load collection
  1379. R<RpcStatus> loadR = client.loadCollection(LoadCollectionParam.newBuilder()
  1380. .withCollectionName(randomCollectionName)
  1381. .build());
  1382. Assertions.assertEquals(R.Status.Success.getCode(), loadR.getStatus().intValue());
  1383. }
  1384. void testCollectionHighLevelDelete(FieldType primaryField, FieldType vectorField) {
  1385. // create collection & buildIndex & loadCollection
  1386. String randomCollectionName = generator.generate(10);
  1387. highLevelCreateCollection(primaryField, vectorField, randomCollectionName);
  1388. // insert data
  1389. List<String> primaryIds = new ArrayList<>();
  1390. int rowCount = 10;
  1391. List<JSONObject> rows = new ArrayList<>();
  1392. for (long i = 0L; i < rowCount; ++i) {
  1393. JSONObject row = new JSONObject();
  1394. row.put(primaryField.getName(), primaryField.getDataType() == DataType.Int64 ? i : String.valueOf(i));
  1395. row.put(vectorField.getName(), generateFloatVectors(1).get(0));
  1396. rows.add(row);
  1397. primaryIds.add(String.valueOf(i));
  1398. }
  1399. InsertParam insertRowParam = InsertParam.newBuilder()
  1400. .withCollectionName(randomCollectionName)
  1401. .withRows(rows)
  1402. .build();
  1403. R<MutationResult> insertRowResp = client.insert(insertRowParam);
  1404. Assertions.assertEquals(R.Status.Success.getCode(), insertRowResp.getStatus().intValue());
  1405. // high level delete
  1406. testHighLevelDelete(randomCollectionName, primaryIds);
  1407. client.dropCollection(DropCollectionParam.newBuilder().withCollectionName(randomCollectionName).build());
  1408. }
  1409. private static void testHighLevelGet(String collectionName, List primaryIds) {
  1410. GetIdsParam getIdsParam = GetIdsParam.newBuilder()
  1411. .withCollectionName(collectionName)
  1412. .withPrimaryIds(primaryIds)
  1413. .build();
  1414. R<GetResponse> getResponseR = client.get(getIdsParam);
  1415. String outPutStr = String.format("collectionName:%s, primaryIds:%s, getResponseR:%s", collectionName, primaryIds, getResponseR.getData());
  1416. System.out.println(outPutStr);
  1417. Assertions.assertEquals(R.Status.Success.getCode(), getResponseR.getStatus().intValue());
  1418. }
  1419. private static void testHighLevelDelete(String collectionName, List primaryIds) {
  1420. DeleteIdsParam deleteIdsParam = DeleteIdsParam.newBuilder()
  1421. .withCollectionName(collectionName)
  1422. .withPrimaryIds(primaryIds)
  1423. .build();
  1424. R<DeleteResponse> deleteResponseR = client.delete(deleteIdsParam);
  1425. String outPutStr = String.format("collectionName:%s, primaryIds:%s, deleteResponseR:%s", collectionName, primaryIds, deleteResponseR);
  1426. System.out.println(outPutStr);
  1427. Assertions.assertEquals(R.Status.Success.getCode(), deleteResponseR.getStatus().intValue());
  1428. }
  1429. }