diff --git a/entity/columns.go b/entity/columns.go index cf6e3a90..cb13572e 100644 --- a/entity/columns.go +++ b/entity/columns.go @@ -36,6 +36,11 @@ type Column interface { GetAsString(int) (string, error) GetAsDouble(int) (float64, error) GetAsBool(int) (bool, error) + GetAsBFloat16Vector(int) (BFloat16Vector, error) + GetAsFloatVector(int) (FloatVector, error) + GetAsFloat16Vector(int) (Float16Vector, error) + GetAsBinaryVector(int) (BinaryVector, error) + GetAsSparseFloatVector(int) (SparseEmbedding, error) } // ColumnBase adds conversion methods support for fixed-type columns. @@ -57,6 +62,26 @@ func (b ColumnBase) GetAsBool(_ int) (bool, error) { return false, errors.New("conversion between fixed-type column not support") } +func (b ColumnBase) GetAsBFloat16Vector(int) (BFloat16Vector, error) { + return nil, errors.New("conversion between fixed-type column not support") +} + +func (b ColumnBase) GetAsFloatVector(int) (FloatVector, error) { + return nil, errors.New("conversion between fixed-type column not support") +} + +func (b ColumnBase) GetAsFloat16Vector(int) (Float16Vector, error) { + return nil, errors.New("conversion between fixed-type column not support") +} + +func (b ColumnBase) GetAsBinaryVector(int) (BinaryVector, error) { + return nil, errors.New("conversion between fixed-type column not support") +} + +func (b ColumnBase) GetAsSparseFloatVector(int) (SparseEmbedding, error) { + return nil, errors.New("conversion between fixed-type column not support") +} + // Vector interface vector used int search type Vector interface { Dim() int diff --git a/entity/columns_sparse.go b/entity/columns_sparse.go index 86693f0d..ee822c90 100644 --- a/entity/columns_sparse.go +++ b/entity/columns_sparse.go @@ -25,6 +25,7 @@ type SparseEmbedding interface { Dim() int // the dimension Len() int // the actual items in this vector Get(idx int) (pos uint32, value float32, ok bool) + Map() map[uint32]float32 Serialize() []byte FieldType() FieldType } @@ -68,6 +69,15 @@ func (e sliceSparseEmbedding) Serialize() []byte { return row } +// Return sparse embedding as map +func (e sliceSparseEmbedding) Map() map[uint32]float32 { + result := make(map[uint32]float32, 0) + for i := 0; i < e.len; i++ { + result[e.positions[i]] = e.values[i] + } + return result +} + // Less implements sort.Interce func (e sliceSparseEmbedding) Less(i, j int) bool { return e.positions[i] < e.positions[j] @@ -166,6 +176,10 @@ func (c *ColumnSparseFloatVector) Get(idx int) (interface{}, error) { return c.vectors[idx], nil } +func (c *ColumnSparseFloatVector) GetAsSparseFloatVector(idx int) (SparseEmbedding, error) { + return c.ValueByIdx(idx) +} + // ValueByIdx returns value of the provided index // error occurs when index out of range func (c *ColumnSparseFloatVector) ValueByIdx(idx int) (SparseEmbedding, error) { diff --git a/entity/columns_vector_gen.go b/entity/columns_vector_gen.go index f690d36d..890cf202 100755 --- a/entity/columns_vector_gen.go +++ b/entity/columns_vector_gen.go @@ -66,6 +66,14 @@ func (c *ColumnBinaryVector) Get(idx int) (interface{}, error) { return c.values[idx], nil } +// Get returns values at index as interface{}. +func (c *ColumnBinaryVector) GetAsBinaryVector(idx int) (BinaryVector, error) { + if idx < 0 || idx >= c.Len() { + return nil, errors.New("index out of range") + } + return c.values[idx], nil +} + // AppendValue append value into column func(c *ColumnBinaryVector) AppendValue(i interface{}) error { v, ok := i.([]byte) @@ -153,6 +161,14 @@ func (c *ColumnFloatVector) Get(idx int) (interface{}, error) { return c.values[idx], nil } +// Get returns values at index as FloatVector. +func (c *ColumnFloatVector) GetAsFloatVector(idx int) (FloatVector, error) { + if idx < 0 || idx >= c.Len() { + return nil, errors.New("index out of range") + } + return c.values[idx], nil +} + func (c *ColumnFloatVector) Slice(start, end int) Column { l := c.Len() if start > l { @@ -280,6 +296,15 @@ func (c *ColumnFloat16Vector) Get(idx int) (interface{}, error) { return c.values[idx], nil } +// Get returns values at index as Float16Vector. +func (c *ColumnFloat16Vector) GetAsFloat16Vector(idx int) (Float16Vector, error) { + if idx < 0 || idx >= c.Len() { + return nil, errors.New("index out of range") + } + return c.values[idx], nil +} + + // AppendValue append value into column func(c *ColumnFloat16Vector) AppendValue(i interface{}) error { v, ok := i.([]byte) @@ -386,6 +411,14 @@ func (c *ColumnBFloat16Vector) Get(idx int) (interface{}, error) { return c.values[idx], nil } +// Get returns values at index as interface{}. +func (c *ColumnBFloat16Vector) GetAsBinaryVector(idx int) (BinaryVector, error) { + if idx < 0 || idx >= c.Len() { + return nil, errors.New("index out of range") + } + return c.values[idx], nil +} + // AppendValue append value into column func(c *ColumnBFloat16Vector) AppendValue(i interface{}) error { v, ok := i.([]byte)