io.milvus.client.SearchParam 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 restore the calcDistance interface that is removed in 2.1.0-beta4
There is a newer version: 2.2.2.1
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.client;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import javax.annotation.Nonnull;

/** Contains parameters for search */
public class SearchParam {

  private final String collectionName;
  private final List> floatVectors;
  private final List binaryVectors;
  private final List partitionTags;
  private final long topK;
  private final String paramsInJson;

  private SearchParam(@Nonnull Builder builder) {
    this.collectionName = builder.collectionName;
    this.floatVectors = builder.floatVectors;
    this.binaryVectors = builder.binaryVectors;
    this.partitionTags = builder.partitionTags;
    this.topK = builder.topK;
    this.paramsInJson = builder.paramsInJson;
  }

  public String getCollectionName() {
    return collectionName;
  }

  public List> getFloatVectors() {
    return floatVectors;
  }

  public List getBinaryVectors() {
    return binaryVectors;
  }

  public List getPartitionTags() {
    return partitionTags;
  }

  public long getTopK() {
    return topK;
  }

  public String getParamsInJson() {
    return paramsInJson;
  }

  /** Builder for SearchParam */
  public static class Builder {
    // Required parameters
    private final String collectionName;

    // Optional parameters - initialized to default values
    private List> floatVectors = new ArrayList<>();
    private List binaryVectors = new ArrayList<>();
    private List partitionTags = new ArrayList<>();
    private long topK = 1024;
    private String paramsInJson;

    /** @param collectionName collection to search from */
    public Builder(@Nonnull String collectionName) {
      this.collectionName = collectionName;
    }

    /**
     * Default to an empty ArrayList. You can search either float or binary vectors,
     * not both.
     *
     * @param floatVectors a List of float vectors to be queries. Each inner List
     *      represents a float vector.
     * @return Builder
     */
    public SearchParam.Builder withFloatVectors(@Nonnull List> floatVectors) {
      this.floatVectors = floatVectors;
      return this;
    }

    /**
     * Default to an empty ArrayList. You can search either float or binary vectors,
     * not both.
     *
     * @param binaryVectors a List of binary vectors to be queried. Each 
     *     ByteBuffer object represents a binary vector, with every 8 bits constituting a
     *     byte.
     * @return Builder
     * @see ByteBuffer
     */
    public SearchParam.Builder withBinaryVectors(@Nonnull List binaryVectors) {
      this.binaryVectors = binaryVectors;
      return this;
    }

    /**
     * Optional. Search vectors with corresponding partitionTags. Default to an empty
     * List
     *
     * @param partitionTags a List of partition tags
     * @return Builder
     */
    public Builder withPartitionTags(@Nonnull List partitionTags) {
      this.partitionTags = partitionTags;
      return this;
    }

    /**
     * Optional. Limits search result to topK. Default to 1024.
     *
     * @param topK a topK number
     * @return Builder
     */
    public Builder withTopK(long topK) {
      this.topK = topK;
      return this;
    }

    /**
     * Optional. Default to empty String. Search parameters are different for different
     * index types. Refer to https://milvus.io/docs/milvus_operation.md
     * for more information.
     *
     *      *   FLAT/IVFLAT/SQ8/IVFPQ: {"nprobe": 32}
     *   nprobe range:[1,999999]
     *
     *   NSG: {"search_length": 100}
     *   search_length range:[10, 300]
     *
     *   HNSW: {"ef": 64}
     *   ef range:[topk, 4096]
     *
     *   ANNOY: {search_k", 0.05 * totalDataCount}
     *   search_k range: none
     * 
     *
     * @param paramsInJson extra parameters in JSON format
     * @return Builder
     */
    public SearchParam.Builder withParamsInJson(@Nonnull String paramsInJson) {
      this.paramsInJson = paramsInJson;
      return this;
    }

    public SearchParam build() {
      return new SearchParam(this);
    }
  }
}