コード|milvusライブラリに1億のデータを挿入
4830 ワード
公式テストに必要なメモリサイズツール:https://milvus.io/tools/sizing/
package com.darwin.milvusdemo.test;
import com.google.gson.JsonObject;
import io.milvus.client.*;
import java.time.LocalDate;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.DoubleStream;
public class MilvusClientExample {
// Helper function that generates random vectors
static List> generateVectors(long vectorCount, long dimension) {
SplittableRandom splitcollectionRandom = new SplittableRandom();
List> vectors = new ArrayList<>();
for (long i = 0; i < vectorCount; ++i) {
splitcollectionRandom = splitcollectionRandom.split();
DoubleStream doubleStream = splitcollectionRandom.doubles(dimension);
List vector =
doubleStream.boxed().map(Double::floatValue).collect(Collectors.toList());
vectors.add(vector);
}
return vectors;
}
// Helper function that normalizes a vector if you are using IP (Inner Product) as your metric
// type
static List normalizeVector(List vector) {
float squareSum = vector.stream().map(x -> x * x).reduce((float) 0, Float::sum);
final float norm = (float) Math.sqrt(squareSum);
vector = vector.stream().map(x -> x / norm).collect(Collectors.toList());
return vector;
}
public static void main(String[] args) throws InterruptedException, ConnectFailedException {
// You may need to change the following to the host and port of your Milvus server
String host = "10.10.8.74";
int port = 19530;
// Create Milvus client
MilvusClient client = new MilvusGrpcClient();
// Connect to Milvus server
ConnectParam connectParam = new ConnectParam.Builder().withHost(host).withPort(port).build();
try {
Response connectResponse = client.connect(connectParam);
} catch (ConnectFailedException e) {
System.out.println("Failed to connect to Milvus server: " + e.toString());
throw e;
}
// client.dropCollection("snap_face");
// Create a collection with the following collection mapping
final String collectionName = "snap_face"; // collection name
final int dimension = 512; // dimension of each vector
final int indexFileSize = 1024; // maximum size (in MB) of each index file
final MetricType metricType = MetricType.IP; // we choose IP (Inner Product) as our metric type
CollectionMapping collectionMapping =
new CollectionMapping.Builder(collectionName, dimension)
.withIndexFileSize(indexFileSize)
.withMetricType(metricType)
.build();
//
client.createCollection(collectionMapping);
// Check whether the collection exists
HasCollectionResponse hasCollectionResponse = client.hasCollection(collectionName);
if (!hasCollectionResponse.ok()) {
System.err.println("1 wrong");
return;
}
final IndexType indexType = IndexType.IVF_SQ8_H;
// Each index type has its optional parameters you can set. Refer to the Milvus documentation
// for how to set the optimal parameters based on your needs.
JsonObject indexParamsJson = new JsonObject();
indexParamsJson.addProperty("nlist", 65536);
Index index =
new Index.Builder(collectionName, indexType)
.withParamsInJson(indexParamsJson.toString())
.build();
Response createIndexResponse = client.createIndex(index);
if (!createIndexResponse.ok()) {
System.err.println("2 wrong");
return;
}
Response response = client.createPartition("snap_face", LocalDate.now().toString());
if (!response.ok()) {
System.err.println("3 wrong");
return;
}
for (int i = 0; i < 1000; i++) {
System.out.println(i + " successful");
final int vectorCount = 100000;
List> vectors = generateVectors(vectorCount, dimension);
vectors =
vectors.stream().map(MilvusClientExample::normalizeVector).collect(Collectors.toList());
InsertParam insertParam =
new InsertParam.Builder(collectionName).withPartitionTag(LocalDate.now().toString()).withFloatVectors(vectors).build();
InsertResponse insertResponse = client.insert(insertParam);
}
}
}