github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/entity/index_cagra.go (about)

     1  package entity
     2  
     3  import (
     4  	"encoding/json"
     5  	"strconv"
     6  
     7  	"github.com/cockroachdb/errors"
     8  )
     9  
    10  var _ Index = &IndexGPUCagra{}
    11  
    12  // IndexGPUCagra index type for GPU Cagra index.
    13  type IndexGPUCagra struct {
    14  	metricType              MetricType
    15  	intermediateGraphDegree int
    16  	graphDegree             int
    17  }
    18  
    19  // NewIndexGPUCagra returns an Index with GPU_CAGRA type.
    20  //   - intermediate_graph_degree:  The number of k-nearest neighbors (k) of this intermediate k-NN graph, trade off the quality of the final searchable CAGRA graph;
    21  //   - graph_degree: CAGRA's optimized graph fixed-degree number;
    22  func NewIndexGPUCagra(metricType MetricType,
    23  	intermediateGraphDegree, graphDegree int) (*IndexGPUCagra, error) {
    24  
    25  	if intermediateGraphDegree < graphDegree {
    26  		return nil, errors.New("Graph degree cannot be larger than intermediate graph degree")
    27  	}
    28  
    29  	return &IndexGPUCagra{
    30  		metricType:              metricType,
    31  		intermediateGraphDegree: intermediateGraphDegree,
    32  		graphDegree:             graphDegree,
    33  	}, nil
    34  }
    35  
    36  // Name returns index type name, implementing Index interface
    37  func (i *IndexGPUCagra) Name() string {
    38  	return "GPUCagra"
    39  }
    40  
    41  // IndexType returns IndexType, implementing Index interface
    42  func (i *IndexGPUCagra) IndexType() IndexType {
    43  	return GPUCagra
    44  }
    45  
    46  // SupportBinary returns whether index type support binary vector
    47  func (i *IndexGPUCagra) SupportBinary() bool {
    48  	return false
    49  }
    50  
    51  // Params returns index construction params, implementing Index interface
    52  func (i *IndexGPUCagra) Params() map[string]string {
    53  	params := map[string]string{ //auto generated mapping
    54  		"intermediate_graph_degree": strconv.FormatInt(int64(i.intermediateGraphDegree), 10),
    55  		"graph_degree":              strconv.FormatInt(int64(i.graphDegree), 10),
    56  	}
    57  	bs, _ := json.Marshal(params)
    58  	return map[string]string{
    59  		"params":      string(bs),
    60  		"index_type":  string(i.IndexType()),
    61  		"metric_type": string(i.metricType),
    62  	}
    63  }
    64  
    65  type IndexGPUCagraSearchParam struct {
    66  	baseSearchParams
    67  }
    68  
    69  // - itopk_size: the main parameter that can be adjusted to trade off search speed, which specifies the size of an internal sorted list that stores the nodes that can be explored in the next iteration;
    70  // - search_width: the number of the closest parent vertices that are traversed to expand their children in each search iteration;
    71  // - min_iterations: Lower limit of search iterations;
    72  // - max_iterations: Upper limit of search iterations. Auto select when 0;
    73  // - team_size: Number of threads used to calculate a single distance.
    74  func NewIndexGPUCagraSearchParam(
    75  	itopkSize int,
    76  	searchWidth int,
    77  	minIterations int,
    78  	maxIterations int,
    79  	teamSize int,
    80  ) (*IndexGPUCagraSearchParam, error) {
    81  
    82  	if !(teamSize == 0 || teamSize == 4 || teamSize == 8 || teamSize == 16 || teamSize == 32) {
    83  		return nil, errors.New("teamSize shall be 0, 4, 8, 16 or 32")
    84  	}
    85  
    86  	sp := &IndexGPUCagraSearchParam{
    87  		baseSearchParams: newBaseSearchParams(),
    88  	}
    89  
    90  	sp.params["itopk_size"] = itopkSize
    91  	sp.params["search_width"] = searchWidth
    92  	sp.params["min_iterations"] = minIterations
    93  	sp.params["max_iterations"] = maxIterations
    94  	sp.params["team_size"] = teamSize
    95  
    96  	return sp, nil
    97  }
    98  
    99  // IndexGPUBruteForce index type for GPU brute force search.
   100  type IndexGPUBruteForce struct {
   101  	metricType MetricType
   102  }
   103  
   104  func NewIndexGPUBruteForce(metricType MetricType) (*IndexGPUBruteForce, error) {
   105  	return &IndexGPUBruteForce{
   106  		metricType: metricType,
   107  	}, nil
   108  }
   109  
   110  // Name returns index type name, implementing Index interface
   111  func (i *IndexGPUBruteForce) Name() string {
   112  	return "GPUBruteForce"
   113  }
   114  
   115  // IndexType returns IndexType, implementing Index interface
   116  func (i *IndexGPUBruteForce) IndexType() IndexType {
   117  	return GPUBruteForce
   118  }
   119  
   120  // SupportBinary returns whether index type support binary vector
   121  func (i *IndexGPUBruteForce) SupportBinary() bool {
   122  	return false
   123  }
   124  
   125  // Params returns index construction params, implementing Index interface
   126  func (i *IndexGPUBruteForce) Params() map[string]string {
   127  	return map[string]string{
   128  		"params":      "{}",
   129  		"index_type":  string(i.IndexType()),
   130  		"metric_type": string(i.metricType),
   131  	}
   132  }
   133  
   134  type IndexGPUBruteForceSearchParam struct {
   135  	baseSearchParams
   136  }
   137  
   138  func NewIndexGPUBruteForceSearchParam() (*IndexGPUBruteForceSearchParam, error) {
   139  	return &IndexGPUBruteForceSearchParam{
   140  		baseSearchParams: newBaseSearchParams(),
   141  	}, nil
   142  }