Skip to content

Commit

Permalink
enhance: [GoSDK] Sync IndexType supported from v1 (#34310)
Browse files Browse the repository at this point in the history
See also #31293

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
  • Loading branch information
congqixia committed Jul 1, 2024
1 parent 4cf1a35 commit 87bccb1
Show file tree
Hide file tree
Showing 6 changed files with 156 additions and 6 deletions.
39 changes: 39 additions & 0 deletions client/index/auto.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
// 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

var _ Index = autoIndex{}

type autoIndex struct {
baseIndex
}

func (idx autoIndex) Params() map[string]string {
return map[string]string{
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(AUTOINDEX),
}
}

func NewAutoIndex(metricType MetricType) Index {
return autoIndex{
baseIndex: baseIndex{
indexType: AUTOINDEX,
metricType: metricType,
},
}
}
4 changes: 3 additions & 1 deletion client/index/common.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,5 +57,7 @@ const (
GPUCagra IndexType = "GPU_CAGRA"
GPUBruteForce IndexType = "GPU_BRUTE_FORCE"

Scalar IndexType = "SCALAR"
Trie IndexType = "Trie"
Sorted IndexType = "STL_SORT"
Inverted IndexType = "INVERTED"
)
21 changes: 21 additions & 0 deletions client/index/flat.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,24 @@ func NewFlatIndex(metricType MetricType) Index {
},
}
}

var _ Index = binFlatIndex{}

type binFlatIndex struct {
baseIndex
}

func (idx binFlatIndex) Params() map[string]string {
return map[string]string{
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(BinFlat),
}
}

func NewBinFlatIndex(metricType MetricType) Index {
return binFlatIndex{
baseIndex: baseIndex{
metricType: metricType,
},
}
}
31 changes: 31 additions & 0 deletions client/index/ivf.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ func NewIvfFlatIndex(metricType MetricType, nlist int) Index {
}
}

var _ Index = ivfPQIndex{}

type ivfPQIndex struct {
baseIndex

Expand Down Expand Up @@ -82,6 +84,8 @@ func NewIvfPQIndex(metricType MetricType, nlist int, m int, nbits int) Index {
}
}

var _ Index = ivfSQ8Index{}

type ivfSQ8Index struct {
baseIndex

Expand All @@ -106,3 +110,30 @@ func NewIvfSQ8Index(metricType MetricType, nlist int) Index {
nlist: nlist,
}
}

var _ Index = binIvfFlat{}

type binIvfFlat struct {
baseIndex

nlist int
}

func (idx binIvfFlat) Params() map[string]string {
return map[string]string{
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(IvfSQ8),
ivfNlistKey: strconv.Itoa(idx.nlist),
}
}

func NewBinIvfFlatIndex(metricType MetricType, nlist int) Index {
return binIvfFlat{
baseIndex: baseIndex{
metricType: metricType,
indexType: BinIvfFlat,
},

nlist: nlist,
}
}
56 changes: 56 additions & 0 deletions client/index/scalar.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
// 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

type scalarIndex struct {
name string
indexType IndexType
}

func (idx scalarIndex) Name() string {
return idx.name
}

func (idx scalarIndex) IndexType() IndexType {
return idx.indexType
}

func (idx scalarIndex) Params() map[string]string {
return map[string]string{
IndexTypeKey: string(idx.indexType),
}
}

var _ Index = scalarIndex{}

func NewTrieIndex() Index {
return scalarIndex{
indexType: Trie,
}
}

func NewInvertedIndex() Index {
return scalarIndex{
indexType: Inverted,
}
}

func NewSortedIndex() Index {
return scalarIndex{
indexType: Sorted,
}
}
11 changes: 6 additions & 5 deletions client/index/scann.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,17 +32,18 @@ type scannIndex struct {

func (idx scannIndex) Params() map[string]string {
return map[string]string{
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(IvfFlat),
ivfNlistKey: strconv.Itoa(idx.nlist),
MetricTypeKey: string(idx.metricType),
IndexTypeKey: string(SCANN),
scannNlistKey: strconv.Itoa(idx.nlist),
scannWithRawDataKey: strconv.FormatBool(idx.withRawData),
}
}

func NewSCANNIndex(metricType MetricType, nlist int) Index {
func NewSCANNIndex(metricType MetricType, nlist int, withRawData bool) Index {
return ivfFlatIndex{
baseIndex: baseIndex{
metricType: metricType,
indexType: IvfFlat,
indexType: SCANN,
},

nlist: nlist,
Expand Down

0 comments on commit 87bccb1

Please sign in to comment.