io.milvus.response.DescCollResponseWrapper 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 database.
There is a newer version: 2.4.8
package io.milvus.response;

import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.exception.ParamException;
import io.milvus.grpc.*;
import io.milvus.param.ParamUtils;
import io.milvus.param.collection.CollectionSchemaParam;
import io.milvus.param.collection.FieldType;
import lombok.NonNull;

import java.util.*;

/**
 * Util class to wrap response of describeCollection interface.
 */
public class DescCollResponseWrapper {
    private final DescribeCollectionResponse response;

    public DescCollResponseWrapper(@NonNull DescribeCollectionResponse response) {
        this.response = response;
    }

    /**
     * Get name of the collection.
     *
     * @return String name of the collection
     */
    public String getCollectionName() {
        CollectionSchema schema = response.getSchema();
        return schema.getName();
    }

    /**
     * Get database name of the collection.
     *
     * @return String name of the database
     */
    public String getDatabaseName() {
        return response.getDbName();
    }

    /**
     * Get description of the collection.
     *
     * @return String description of the collection
     */
    public String getCollectionDescription() {
        CollectionSchema schema = response.getSchema();
        return schema.getDescription();
    }

    /**
     * Get internal id of the collection.
     *
     * @return long internal id of the collection
     */
    public long getCollectionID() {
        return response.getCollectionID();
    }

    /**
     * Get shard number of the collection.
     *
     * @return int shard number of the collection
     */
    public int getShardNumber() {
        return response.getShardsNum();
    }

    /**
     * Get consistency level of the collection.
     *
     * @return ConsistencyLevelEnum consistency level of the collection
     */
    public ConsistencyLevelEnum getConsistencyLevel() {
        // may throw IllegalArgumentException
        return ConsistencyLevelEnum.valueOf(response.getConsistencyLevel().name().toUpperCase());
    }

    /**
     * Get utc timestamp when collection created.
     *
     * @return long utc timestamp when collection created
     */
    public long getCreatedUtcTimestamp() {
        return response.getCreatedUtcTimestamp();
    }

    /**
     * Get aliases of the collection.
     *
     * @return List of String, aliases of the collection
     */
    public List getAliases() {
        List aliases = new ArrayList<>();
        for (int i = 0; i < response.getAliasesCount(); ++i) {
            aliases.add(response.getAliases(i));
        }

        return aliases;
    }

    /**
     * Get schema of the collection's fields.
     *
     * @return List of FieldType, schema of the collection's fields
     */
    public List getFields() {
        List results = new ArrayList<>();
        CollectionSchema schema = response.getSchema();
        List fields = schema.getFieldsList();
        fields.forEach((field) -> results.add(ParamUtils.ConvertField(field)));

        return results;
    }

    /**
     * Get schema of a field by name.
     * Return null if the field doesn't exist
     *
     * @param fieldName field name to get field description
     * @return {@link FieldType} schema of the field
     */
    public FieldType getFieldByName(@NonNull String fieldName) {
        CollectionSchema schema = response.getSchema();
        for (int i = 0; i < schema.getFieldsCount(); ++i) {
            FieldSchema field = schema.getFields(i);
            if (fieldName.compareTo(field.getName()) == 0) {
                return ParamUtils.ConvertField(field);
            }
        }

        return null;
    }

    // duplicated with isDynamicFieldEnabled()
    @Deprecated
    public boolean getEnableDynamicField() {
        CollectionSchema schema = response.getSchema();
        return schema.getEnableDynamicField();
    }

    /**
     * Get whether the collection dynamic field is enabled
     *
     * @return boolean
     */
    public boolean isDynamicFieldEnabled() {
        CollectionSchema schema = response.getSchema();
        return schema.getEnableDynamicField();
    }

    /**
     * Get the partition key field.
     * Return null if the partition key field doesn't exist.
     *
     * @return {@link FieldType} schema of the partition key field
     */
    public FieldType getPartitionKeyField() {
        CollectionSchema schema = response.getSchema();
        for (int i = 0; i < schema.getFieldsCount(); ++i) {
            FieldSchema field = schema.getFields(i);
            if (field.getIsPartitionKey()) {
                return ParamUtils.ConvertField(field);
            }
        }

        return null;
    }

    /**
     * Get the primary key field.
     * throw ParamException if the primary key field doesn't exist.
     *
     * @return {@link FieldType} schema of the primary key field
     */
    public FieldType getPrimaryField() {
        CollectionSchema schema = response.getSchema();
        for (int i = 0; i < schema.getFieldsCount(); ++i) {
            FieldSchema field = schema.getFields(i);
            if (field.getIsPrimaryKey()) {
                return ParamUtils.ConvertField(field);
            }
        }

        throw new ParamException("No primary key found.");
    }

    /**
     * Get the vector key field.
     * throw ParamException if the vector key field doesn't exist.
     *
     * @return {@link FieldType} schema of the vector key field
     */
    public FieldType getVectorField() {
        CollectionSchema schema = response.getSchema();
        for (int i = 0; i < schema.getFieldsCount(); ++i) {
            FieldSchema field = schema.getFields(i);
            if (field.getDataType() == DataType.FloatVector || field.getDataType() == DataType.BinaryVector) {
                return ParamUtils.ConvertField(field);
            }
        }

        throw new ParamException("No vector key found.");
    }

    /**
     * Get properties of the collection.
     *
     * @return List of String, aliases of the collection
     */
    public Map getProperties() {
        Map pairs = new HashMap<>();
        List props = response.getPropertiesList();
        props.forEach((prop) -> pairs.put(prop.getKey(), prop.getValue()));

        return pairs;
    }

    /**
     * Get the collection schema of collection
     *
     * @return {@link CollectionSchemaParam} schema of the collection
     */
    public CollectionSchemaParam getSchema() {
        return CollectionSchemaParam.newBuilder()
                .withFieldTypes(getFields())
                .withEnableDynamicField(isDynamicFieldEnabled())
                .build();
    }

    /**
     * Construct a String by {@link DescCollResponseWrapper} instance.
     *
     * @return String
     */
    @Override
    public String toString() {
        return "Collection Description{" +
                "name:'" + getCollectionName() + '\'' +
                ", databaseName:'" + getDatabaseName() + '\'' +
                ", description:'" + getCollectionDescription() + '\'' +
                ", id:" + getCollectionID() +
                ", shardNumber:" + getShardNumber() +
                ", createdUtcTimestamp:" + getCreatedUtcTimestamp() +
                ", aliases:" + getAliases().toString() +
                ", fields:" + getFields().toString() +
                ", isDynamicFieldEnabled:" + isDynamicFieldEnabled() +
                ", consistencyLevel:" + getConsistencyLevel().name() +
                ", properties:" + getProperties() +
                '}';
    }
}