コード|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);
        }

    }

}