diff --git a/client/index/auto.go b/client/index/auto.go new file mode 100644 index 0000000000000..8490ffa8d4d13 --- /dev/null +++ b/client/index/auto.go @@ -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, + }, + } +} diff --git a/client/index/common.go b/client/index/common.go index 58e4da922f39f..db3b3610c827b 100644 --- a/client/index/common.go +++ b/client/index/common.go @@ -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" ) diff --git a/client/index/flat.go b/client/index/flat.go index 205fb4df2f1eb..cc336c23d5d27 100644 --- a/client/index/flat.go +++ b/client/index/flat.go @@ -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, + }, + } +} diff --git a/client/index/ivf.go b/client/index/ivf.go index e22d243459bdb..fda76c6866739 100644 --- a/client/index/ivf.go +++ b/client/index/ivf.go @@ -51,6 +51,8 @@ func NewIvfFlatIndex(metricType MetricType, nlist int) Index { } } +var _ Index = ivfPQIndex{} + type ivfPQIndex struct { baseIndex @@ -82,6 +84,8 @@ func NewIvfPQIndex(metricType MetricType, nlist int, m int, nbits int) Index { } } +var _ Index = ivfSQ8Index{} + type ivfSQ8Index struct { baseIndex @@ -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 NewBinIvfFlat(metricType MetricType, nlist int) Index { + return ivfPQIndex{ + baseIndex: baseIndex{ + metricType: metricType, + indexType: BinIvfFlat, + }, + + nlist: nlist, + } +} diff --git a/client/index/scalar.go b/client/index/scalar.go new file mode 100644 index 0000000000000..88433e1eeece2 --- /dev/null +++ b/client/index/scalar.go @@ -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, + } +} diff --git a/client/index/scann.go b/client/index/scann.go index e0f866684942f..196145be4a6f3 100644 --- a/client/index/scann.go +++ b/client/index/scann.go @@ -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,