io.milvus.response.GetBulkloadStateWrapper Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of milvus-sdk-java Show documentation
Java SDK for Milvus, a distributed high-performance vector search engine. update grpc to 1.42.1 update protobuf to 3.19.1
There is a newer version: 2.2.2.1
Show newest version
package io.milvus.response;

import io.milvus.exception.IllegalResponseException;
import io.milvus.grpc.GetImportStateResponse;
import io.milvus.grpc.ImportState;
import io.milvus.grpc.KeyValuePair;
import io.milvus.param.Constant;
import lombok.NonNull;

import java.util.ArrayList;
import java.util.List;

/**
 * Util class to wrap response of getBulkloadState interface.
 */
public class GetBulkloadStateWrapper {
    private final GetImportStateResponse response;

    public GetBulkloadStateWrapper(@NonNull GetImportStateResponse response) {
        this.response = response;
    }

    /**
     * Gets ID of the bulk load task.
     *
     * @return Long ID of the bulk load task
     */
    public long getTaskID() {
        return response.getId();
    }

    /**
     * Gets the long ID array for auto-id primary key, generated by bulk load task.
     *
     * @return List of Long, ID array returned by bulk load task
     */
    public List getAutoGeneratedIDs() {
        // the id list of response is id ranges
        // for example, if the response return [1, 100, 200, 250]
        // the full id list should be [1, 2, 3 ... , 99, 100, 200, 201, 202 ... , 249, 250]
        List ranges = response.getIdListList();
        if (ranges.size()%2 != 0) {
            throw new IllegalResponseException("The bulkload state response id range list is illegal");
        }
        List ids = new ArrayList<>();
        for (int i = 0; i < ranges.size()/2; i++) {
            Long begin = ranges.get(i*2);
            Long end = ranges.get(i*2+1);
            for (long j = begin; j <= end; j++) {
                ids.add(j);
            }
        }
        return ids;
    }

    /**
     * Gets state of the bulk load task.
     *
     * @return ImportState state of the bulk load task
     */
    public ImportState getState() {
        return response.getState();
    }

    /**
     * Gets how many rows were imported by the bulk load task.
     *
     * @return Long how many rows were imported by the bulk load task
     */
    public long getImportedCount() {
        return response.getRowCount();
    }

    /**
     * Gets failed reason of the bulk load task.
     *
     * @return String failed reason of the bulk load task
     */
    public String getFailedReason() {
        return getInfo(Constant.FAILED_REASON);
    }

    /**
     * Gets target files of the bulk load task.
     *
     * @return String target files of the bulk load task
     */
    public String getFiles() {
        return getInfo(Constant.IMPORT_FILES);
    }

    /**
     * Gets target collection name of the bulk load task.
     *
     * @return String target collection name
     */
    public String getCollectionName() {
        return getInfo(Constant.IMPORT_COLLECTION);
    }

    /**
     * Gets target partition name of the bulk load task.
     *
     * @return String target partition name
     */
    public String getPartitionName() {
        return getInfo(Constant.IMPORT_PARTITION);
    }

    /**
     * A flag indicating whether import data are queryable (i.e. loaded in query nodes).
     *
     * @return boolean whether import data are queryable
     */
    public boolean queryable() {
        return response.getDataQueryable();
    }

    /**
     * A flag indicating whether import data are indexed.
     *
     * @return boolean whether import data are queryable
     */
    public boolean indexed() {
        return response.getDataIndexed();
    }

    private String getInfo(@NonNull String key) {
        List infos = response.getInfosList();
        for (KeyValuePair kv : infos) {
            if (kv.getKey().compareTo(key) == 0) {
                return kv.getValue();
            }
        }

        return "";
    }

    /**
     * Construct a String by {@link DescCollResponseWrapper} instance.
     *
     * @return String
     */
    @Override
    public String toString() {
        return "bulk load task state{" +
                ", autoGeneratedIDs:" + getAutoGeneratedIDs() +
                ", state:" + getState().name() +
                ", failed reason:" + getFailedReason() +
                ", files:" + getFiles() +
                '}';
    }
}