Skip to content

Commit

Permalink
Support scalar index field
Browse files Browse the repository at this point in the history
  • Loading branch information
@wayblink
wayblink committed May 25, 2023
1 parent 4206e01 commit fa0cdc0
Show file tree
Hide file tree
Showing 8 changed files with 426 additions and 241 deletions.
4 changes: 2 additions & 2 deletions client/data.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -362,10 +362,10 @@ func (c *GrpcClient) CalcDistance(ctx context.Context, collName string, partitio
return nil, err
}
}
if err := c.checkCollField(ctx, collName, opLeft.Name()); err != nil {
if err := c.checkCollVecField(ctx, collName, opLeft.Name()); err != nil {
return nil, err
}
if err := c.checkCollField(ctx, collName, opRight.Name()); err != nil {
if err := c.checkCollVecField(ctx, collName, opRight.Name()); err != nil {
return nil, err
}

Expand Down
37 changes: 30 additions & 7 deletions client/index.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import (
"github.com/milvus-io/milvus-sdk-go/v2/entity"
)

func (c *GrpcClient) checkCollField(ctx context.Context, collName string, fieldName string) error {
func (c *GrpcClient) checkCollVecField(ctx context.Context, collName string, fieldName string) error {
if err := c.checkCollectionExists(ctx, collName); err != nil {
return err
}
Expand All @@ -45,6 +45,28 @@ func (c *GrpcClient) checkCollField(ctx context.Context, collName string, fieldN
return nil
}

// check if the collection and field exist
func (c *GrpcClient) checkCollField(ctx context.Context, collName string, fieldName string) error {
if err := c.checkCollectionExists(ctx, collName); err != nil {
return err
}
coll, err := c.DescribeCollection(ctx, collName)
if err != nil {
return err
}
var f *entity.Field
for _, field := range coll.Schema.Fields {
if field.Name == fieldName {
f = field
break
}
}
if f == nil {
return fmt.Errorf("field %s of collection %s does not exist", fieldName, collName)
}
return nil
}

type indexDef struct {
name string
fieldName string
Expand Down Expand Up @@ -77,7 +99,7 @@ func (c *GrpcClient) CreateIndex(ctx context.Context, collName string, fieldName
if c.Service == nil {
return ErrClientNotReady
}
if err := c.checkCollField(ctx, collName, fieldName); err != nil {
if err := c.checkCollVecField(ctx, collName, fieldName); err != nil {
return err
}

Expand Down Expand Up @@ -132,7 +154,7 @@ func (c *GrpcClient) DescribeIndex(ctx context.Context, collName string, fieldNa
if c.Service == nil {
return []entity.Index{}, ErrClientNotReady
}
if err := c.checkCollField(ctx, collName, fieldName); err != nil {
if err := c.checkCollVecField(ctx, collName, fieldName); err != nil {
return []entity.Index{}, err
}

Expand All @@ -146,8 +168,9 @@ func (c *GrpcClient) DescribeIndex(ctx context.Context, collName string, fieldNa
params := entity.KvPairsMap(info.Params)
it := params["index_type"] // TODO change to const
idx := entity.NewGenericIndex(
info.IndexName,
info.GetIndexName(),
entity.IndexType(it),
info.GetFieldName(),
params,
)
indexes = append(indexes, idx)
Expand All @@ -161,7 +184,7 @@ func (c *GrpcClient) DropIndex(ctx context.Context, collName string, fieldName s
if c.Service == nil {
return ErrClientNotReady
}
if err := c.checkCollField(ctx, collName, fieldName); err != nil {
if err := c.checkCollVecField(ctx, collName, fieldName); err != nil {
return err
}

Expand Down Expand Up @@ -189,7 +212,7 @@ func (c *GrpcClient) GetIndexState(ctx context.Context, collName string, fieldNa
if c.Service == nil {
return entity.IndexState(common.IndexState_Failed), ErrClientNotReady
}
if err := c.checkCollField(ctx, collName, fieldName); err != nil {
if err := c.checkCollVecField(ctx, collName, fieldName); err != nil {
return entity.IndexState(common.IndexState_IndexStateNone), err
}

Expand Down Expand Up @@ -217,7 +240,7 @@ func (c *GrpcClient) GetIndexBuildProgress(ctx context.Context, collName string,
if c.Service == nil {
return 0, 0, ErrClientNotReady
}
if err := c.checkCollField(ctx, collName, fieldName); err != nil {
if err := c.checkCollVecField(ctx, collName, fieldName); err != nil {
return 0, 0, err
}

Expand Down
47 changes: 35 additions & 12 deletions entity/genidx/genidx.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,11 @@ func(i *Index{{.IdxName}}) IndexType() IndexType {
return IndexType("{{.IdxType}}")
}
// FieldName returns FieldName, implementing Index interface
func(i *Index{{.IdxName}}) FieldName() string {
return "{{.FieldName}}"
}
// SupportBinary returns whether index type support binary vector
func(i *Index{{.IdxName}}) SupportBinary() bool {
return {{.VectorSupport}} & 2 > 0
Expand Down Expand Up @@ -225,6 +230,7 @@ func TestIndex{{.IdxName}}SearchParam(t *testing.T) {
type idxDef struct {
IdxName string
IdxType entity.IndexType
FieldName string
VectorSupport int8
ConstructParams []idxParam
SearchParams []idxParam
Expand Down Expand Up @@ -333,6 +339,7 @@ func main() {
{
IdxName: "Flat",
IdxType: entity.Flat,
FieldName: "vec_field",
ConstructParams: []idxParam{},
SearchParams: []idxParam{},
ValidExamples: []string{
Expand All @@ -348,6 +355,7 @@ func main() {
{
IdxName: "BinFlat",
IdxType: entity.BinFlat,
FieldName: "vec_field",
VectorSupport: int8(binaryVectorSupport),
ConstructParams: []idxParam{
{
Expand Down Expand Up @@ -378,8 +386,9 @@ func main() {
},
// IVF_FLAT
{
IdxName: "IvfFlat",
IdxType: entity.IvfFlat,
IdxName: "IvfFlat",
IdxType: entity.IvfFlat,
FieldName: "vec_field",
ConstructParams: []idxParam{
{
Name: "nlist",
Expand Down Expand Up @@ -411,6 +420,7 @@ func main() {
{
IdxName: "BinIvfFlat",
IdxType: entity.BinIvfFlat,
FieldName: "vec_field",
VectorSupport: int8(binaryVectorSupport),
ConstructParams: []idxParam{
{
Expand Down Expand Up @@ -441,8 +451,9 @@ func main() {
},
// IVF_SQ8
{
IdxName: "IvfSQ8",
IdxType: entity.IvfSQ8,
IdxName: "IvfSQ8",
IdxType: entity.IvfSQ8,
FieldName: "vec_field",
ConstructParams: []idxParam{
{
Name: "nlist",
Expand Down Expand Up @@ -472,8 +483,9 @@ func main() {
},
// IVF_PQ
{
IdxName: "IvfPQ",
IdxType: entity.IvfPQ,
IdxName: "IvfPQ",
IdxType: entity.IvfPQ,
FieldName: "vec_field",
ConstructParams: []idxParam{
{
Name: "nlist",
Expand Down Expand Up @@ -513,8 +525,9 @@ func main() {
},
// HNSW
{
IdxName: "HNSW",
IdxType: entity.HNSW,
IdxName: "HNSW",
IdxType: entity.HNSW,
FieldName: "vec_field",
ConstructParams: []idxParam{
{
Name: "M",
Expand Down Expand Up @@ -550,8 +563,9 @@ func main() {
},
// IVF_HNSW
{
IdxName: "IvfHNSW",
IdxType: entity.IvfHNSW,
IdxName: "IvfHNSW",
IdxType: entity.IvfHNSW,
FieldName: "vec_field",
ConstructParams: []idxParam{
{
Name: "nlist",
Expand Down Expand Up @@ -599,8 +613,9 @@ func main() {
},
// ANNOY
{
IdxName: "ANNOY",
IdxType: entity.ANNOY,
IdxName: "ANNOY",
IdxType: entity.ANNOY,
FieldName: "vec_field",
ConstructParams: []idxParam{
{
Name: "n_trees",
Expand Down Expand Up @@ -628,6 +643,7 @@ func main() {
{
IdxName: "DISKANN",
IdxType: entity.DISKANN,
FieldName: "vec_field",
ConstructParams: []idxParam{},
SearchParams: []idxParam{
{
Expand All @@ -650,6 +666,7 @@ func main() {
{
IdxName: "AUTOINDEX",
IdxType: entity.AUTOINDEX,
FieldName: "vec_field",
ConstructParams: []idxParam{},
SearchParams: []idxParam{
{
Expand All @@ -670,6 +687,12 @@ func main() {
"level = -1",
},
},
{
IdxName: "TRIE",
IdxType: entity.TRIE,
FieldName: "varchar_field",
ConstructParams: []idxParam{},
},
},
}

Expand Down
19 changes: 14 additions & 5 deletions entity/index.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ const (
ANNOY IndexType = "ANNOY"
AUTOINDEX IndexType = "AUTOINDEX"
DISKANN IndexType = "DISKANN"
TRIE IndexType = "TRIE"
)

// Metric Constants
Expand All @@ -61,6 +62,7 @@ type Index interface {
Name() string
IndexType() IndexType
Params() map[string]string
FieldName() string
}

// SearchParam interface for index related search param
Expand All @@ -70,8 +72,9 @@ type SearchParam interface {
}

type baseIndex struct {
it IndexType
name string
it IndexType
name string
fieldName string
}

// Name implements Index
Expand Down Expand Up @@ -102,12 +105,18 @@ func (gi GenericIndex) Params() map[string]string {
return m
}

// FieldName implements Index
func (b baseIndex) FieldName() string {
return b.fieldName
}

// NewGenericIndex create generic index instance
func NewGenericIndex(name string, it IndexType, params map[string]string) Index {
func NewGenericIndex(name string, it IndexType, fieldName string, params map[string]string) Index {
return GenericIndex{
baseIndex: baseIndex{
it: it,
name: name,
it: it,
name: name,
fieldName: fieldName,
},
params: params,
}
Expand Down
Loading

0 comments on commit fa0cdc0

Please sign in to comment.