All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.milvus.param.dml.AnnSearchParam Maven / Gradle / Ivy

There is a newer version: 2.4.8
Show newest version
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package io.milvus.param.dml;

import io.milvus.exception.ParamException;
import io.milvus.grpc.PlaceholderType;
import io.milvus.param.MetricType;
import io.milvus.param.ParamUtils;

import lombok.Getter;
import lombok.NonNull;
import lombok.ToString;

import java.nio.ByteBuffer;
import java.util.List;
import java.util.SortedMap;

/**
 * Parameters for hybridSearch interface.
 */
@Getter
public class AnnSearchParam {

    private final String metricType;
    private final String vectorFieldName;
    private final int topK;
    private final String expr;
    private final List vectors;
    private final Long NQ;
    private final String params;
    private final PlaceholderType plType;

    private AnnSearchParam(@NonNull Builder builder) {
        this.metricType = builder.metricType.name();
        this.vectorFieldName = builder.vectorFieldName;
        this.topK = builder.topK;
        this.expr = builder.expr;
        this.vectors = builder.vectors;
        this.NQ = builder.NQ;
        this.params = builder.params;
        this.plType = builder.plType;
    }

    public static Builder newBuilder() {
        return new Builder();
    }

    /**
     * Builder for {@link AnnSearchParam} class.
     */
    public static class Builder {
        private MetricType metricType = MetricType.None;
        private String vectorFieldName;
        private Integer topK;
        private String expr = "";
        private List vectors;
        private Long NQ;
        private String params = "{}";

        // plType is used to distinct vector type
        // for Float16Vector/BFloat16Vector and BinaryVector, user inputs ByteBuffer
        // the sdk cannot distinct a ByteBuffer is a BinarVector or a Float16Vector
        private PlaceholderType plType = PlaceholderType.None;

        Builder() {
        }

        /**
         * Sets metric type of ANN searching.
         *
         * @param metricType metric type
         * @return Builder
         */
        public Builder withMetricType(@NonNull MetricType metricType) {
            this.metricType = metricType;
            return this;
        }

        /**
         * Sets target vector field by name. Field name cannot be empty or null.
         *
         * @param vectorFieldName vector field name
         * @return Builder
         */
        public Builder withVectorFieldName(@NonNull String vectorFieldName) {
            this.vectorFieldName = vectorFieldName;
            return this;
        }

        /**
         * Sets topK value of ANN search.
         *
         * @param topK topK value
         * @return Builder
         */
        public Builder withTopK(@NonNull Integer topK) {
            this.topK = topK;
            return this;
        }

        /**
         * Sets expression to filter out entities before searching (Optional).
         * @see Boolean Expression Rules
         *
         * @param expr filtering expression
         * @return Builder
         */
        public Builder withExpr(@NonNull String expr) {
            this.expr = expr;
            return this;
        }

        /**
         * Sets the target vectors to search on FloatVector field.
         *
         * @param vectors target vectors to search
         * @return Builder
         */
        public Builder withFloatVectors(@NonNull List> vectors) {
            this.vectors = vectors;
            this.NQ = (long) vectors.size();
            this.plType = PlaceholderType.FloatVector;
            return this;
        }

        /**
         * Sets the target vectors to search on BinaryVector field.
         *
         * @param vectors target vectors to search
         * @return Builder
         */
        public Builder withBinaryVectors(@NonNull List vectors) {
            this.vectors = vectors;
            this.NQ = (long) vectors.size();
            this.plType = PlaceholderType.BinaryVector;
            return this;
        }

        /**
         * Sets the target vectors to search on Float16Vector field.
         *
         * @param vectors target vectors to search
         * @return Builder
         */
        public Builder withFloat16Vectors(@NonNull List vectors) {
            this.vectors = vectors;
            this.NQ = (long) vectors.size();
            this.plType = PlaceholderType.Float16Vector;
            return this;
        }

        /**
         * Sets the target vectors to search on BFloat16Vector field.
         *
         * @param vectors target vectors to search
         * @return Builder
         */
        public Builder withBFloat16Vectors(@NonNull List vectors) {
            this.vectors = vectors;
            this.NQ = (long) vectors.size();
            this.plType = PlaceholderType.BFloat16Vector;
            return this;
        }

        /**
         * Sets the target vectors to search on SparseFloatVector field.
         *
         * @param vectors target vectors to search
         * @return Builder
         */
        public Builder withSparseFloatVectors(@NonNull List> vectors) {
            this.vectors = vectors;
            this.NQ = (long) vectors.size();
            this.plType = PlaceholderType.SparseFloatVector;
            return this;
        }


        /**
         * Sets the search parameters specific to the index type.
         *
         * For example: IVF index, the search parameters can be "{\"nprobe\":10}"
         * For more information: @see Index Selection
         *
         * @param params extra parameters in json format
         * @return Builder
         */
        public Builder withParams(@NonNull String params) {
            this.params = params;
            return this;
        }


        /**
         * Verifies parameters and creates a new {@link AnnSearchParam} instance.
         *
         * @return {@link AnnSearchParam}
         */
        public AnnSearchParam build() throws ParamException {
            ParamUtils.CheckNullEmptyString(vectorFieldName, "Target field name");

            if (topK <= 0) {
                throw new ParamException("TopK value is illegal");
            }

            if (vectors.isEmpty()) {
                throw new ParamException("At lease a vector is required for AnnSearchParam");
            }

            SearchParam.verifyVectors(vectors);

            return new AnnSearchParam(this);
        }
    }

    /**
     *
     * Warning: don't use lombok@ToString to annotate this class
     * because large number of vectors will waste time in toString() method.
     *
     */
    @Override
    public String toString() {
        return "AnnSearchParam{" +
                "metricType=" + metricType +
                ", vectorFieldName='" + vectorFieldName + '\'' +
                ", expr='" + expr + '\'' +
                ", topK=" + topK +
                ", nq=" + NQ +
                ", expr='" + expr + '\'' +
                ", params='" + params + '\'' +
                '}';
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy