milvus/client/index/ivf.go
congqixia 244d2c04f6
feat: Add milvusclient package and migrate GoSDK (#32907)
Related to #31293

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2024-05-10 18:01:30 +08:00

109 lines
2.3 KiB
Go

// Licensed to the LF AI & Data foundation 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 index
import "strconv"
const (
ivfNlistKey = `nlist`
ivfPQMKey = `m`
ivfPQNbits = `nbits`
)
var _ Index = ivfFlatIndex{}
type ivfFlatIndex struct {
baseIndex
nlist int
}
func (idx ivfFlatIndex) Params() map[string]string {
return map[string]string{
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(IvfFlat),
ivfNlistKey: strconv.Itoa(idx.nlist),
}
}
func NewIvfFlatIndex(metricType MetricType, nlist int) Index {
return ivfFlatIndex{
baseIndex: baseIndex{
metricType: metricType,
indexType: IvfFlat,
},
nlist: nlist,
}
}
type ivfPQIndex struct {
baseIndex
nlist int
m int
nbits int
}
func (idx ivfPQIndex) Params() map[string]string {
return map[string]string{
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(IvfPQ),
ivfNlistKey: strconv.Itoa(idx.nlist),
ivfPQMKey: strconv.Itoa(idx.m),
ivfPQNbits: strconv.Itoa(idx.nbits),
}
}
func NewIvfPQIndex(metricType MetricType, nlist int, m int, nbits int) Index {
return ivfPQIndex{
baseIndex: baseIndex{
metricType: metricType,
indexType: IvfPQ,
},
nlist: nlist,
m: m,
nbits: nbits,
}
}
type ivfSQ8Index struct {
baseIndex
nlist int
}
func (idx ivfSQ8Index) Params() map[string]string {
return map[string]string{
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(IvfSQ8),
ivfNlistKey: strconv.Itoa(idx.nlist),
}
}
func NewIvfSQ8Index(metricType MetricType, nlist int) Index {
return ivfPQIndex{
baseIndex: baseIndex{
metricType: metricType,
indexType: IvfSQ8,
},
nlist: nlist,
}
}