mirror of https://github.com/milvus-io/milvus.git
120 lines
3.6 KiB
Go
120 lines
3.6 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 milvusclient
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/milvus-io/milvus/client/v2/entity"
|
|
"github.com/milvus-io/milvus/client/v2/index"
|
|
)
|
|
|
|
type SearchIteratorOption interface {
|
|
SearchOption() *searchOption
|
|
}
|
|
|
|
type searchIteratorOption struct {
|
|
*searchOption
|
|
batchSize int
|
|
}
|
|
|
|
func (opt *searchIteratorOption) SearchOption() *searchOption {
|
|
opt.annRequest.topK = opt.batchSize
|
|
opt.WithSearchParam(IteratorSearchBatchSizeKey, fmt.Sprintf("%d", opt.batchSize))
|
|
return opt.searchOption
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithBatchSize(batchSize int) *searchIteratorOption {
|
|
opt.batchSize = batchSize
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithPartitions(partitionNames ...string) *searchIteratorOption {
|
|
opt.partitionNames = partitionNames
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithFilter(expr string) *searchIteratorOption {
|
|
opt.annRequest.WithFilter(expr)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithTemplateParam(key string, val any) *searchIteratorOption {
|
|
opt.annRequest.WithTemplateParam(key, val)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithOffset(offset int) *searchIteratorOption {
|
|
opt.annRequest.WithOffset(offset)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithOutputFields(fieldNames ...string) *searchIteratorOption {
|
|
opt.outputFields = fieldNames
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithConsistencyLevel(consistencyLevel entity.ConsistencyLevel) *searchIteratorOption {
|
|
opt.consistencyLevel = consistencyLevel
|
|
opt.useDefaultConsistencyLevel = false
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithANNSField(annsField string) *searchIteratorOption {
|
|
opt.annRequest.WithANNSField(annsField)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithGroupByField(groupByField string) *searchIteratorOption {
|
|
opt.annRequest.WithGroupByField(groupByField)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithGroupSize(groupSize int) *searchIteratorOption {
|
|
opt.annRequest.WithGroupSize(groupSize)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithStrictGroupSize(strictGroupSize bool) *searchIteratorOption {
|
|
opt.annRequest.WithStrictGroupSize(strictGroupSize)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithIgnoreGrowing(ignoreGrowing bool) *searchIteratorOption {
|
|
opt.annRequest.WithIgnoreGrowing(ignoreGrowing)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithAnnParam(ap index.AnnParam) *searchIteratorOption {
|
|
opt.annRequest.WithAnnParam(ap)
|
|
return opt
|
|
}
|
|
|
|
func (opt *searchIteratorOption) WithSearchParam(key, value string) *searchIteratorOption {
|
|
opt.annRequest.WithSearchParam(key, value)
|
|
return opt
|
|
}
|
|
|
|
func NewSearchIteratorOption(collectionName string, vector entity.Vector) *searchIteratorOption {
|
|
return &searchIteratorOption{
|
|
searchOption: NewSearchOption(collectionName, 100, []entity.Vector{vector}).
|
|
WithSearchParam(IteratorKey, "true").
|
|
WithSearchParam(IteratorSearchV2Key, "true"),
|
|
batchSize: 1000,
|
|
}
|
|
}
|