mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-28 22:45:26 +08:00
175f49ea60
Related to #31293 Previous column is written using go before 1.18, which does not support generic type. After migrating client pkg to main repo, generic base could be used to reduce lots of duplicated code. --------- Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
411 lines
12 KiB
Go
411 lines
12 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 column
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
"github.com/milvus-io/milvus/client/v2/entity"
|
|
)
|
|
|
|
//go:generate go run gen/gen.go
|
|
|
|
// Column interface field type for column-based data frame
|
|
type Column interface {
|
|
Name() string
|
|
Type() entity.FieldType
|
|
Len() int
|
|
Slice(int, int) Column
|
|
FieldData() *schemapb.FieldData
|
|
AppendValue(interface{}) error
|
|
Get(int) (interface{}, error)
|
|
GetAsInt64(int) (int64, error)
|
|
GetAsString(int) (string, error)
|
|
GetAsDouble(int) (float64, error)
|
|
GetAsBool(int) (bool, error)
|
|
}
|
|
|
|
var errFieldDataTypeNotMatch = errors.New("FieldData type not matched")
|
|
|
|
// IDColumns converts schemapb.IDs to corresponding column
|
|
// currently Int64 / string may be in IDs
|
|
func IDColumns(schema *entity.Schema, ids *schemapb.IDs, begin, end int) (Column, error) {
|
|
var idColumn Column
|
|
pkField := schema.PKField()
|
|
if pkField == nil {
|
|
return nil, errors.New("PK Field not found")
|
|
}
|
|
switch pkField.DataType {
|
|
case entity.FieldTypeInt64:
|
|
data := ids.GetIntId().GetData()
|
|
if data == nil {
|
|
return NewColumnInt64(pkField.Name, nil), nil
|
|
}
|
|
if end >= 0 {
|
|
idColumn = NewColumnInt64(pkField.Name, data[begin:end])
|
|
} else {
|
|
idColumn = NewColumnInt64(pkField.Name, data[begin:])
|
|
}
|
|
case entity.FieldTypeVarChar, entity.FieldTypeString:
|
|
data := ids.GetStrId().GetData()
|
|
if data == nil {
|
|
return NewColumnVarChar(pkField.Name, nil), nil
|
|
}
|
|
if end >= 0 {
|
|
idColumn = NewColumnVarChar(pkField.Name, data[begin:end])
|
|
} else {
|
|
idColumn = NewColumnVarChar(pkField.Name, data[begin:])
|
|
}
|
|
default:
|
|
return nil, fmt.Errorf("unsupported id type %v", pkField.DataType)
|
|
}
|
|
return idColumn, nil
|
|
}
|
|
|
|
// FieldDataColumn converts schemapb.FieldData to Column, used int search result conversion logic
|
|
// begin, end specifies the start and end positions
|
|
func FieldDataColumn(fd *schemapb.FieldData, begin, end int) (Column, error) {
|
|
switch fd.GetType() {
|
|
case schemapb.DataType_Bool:
|
|
data, ok := fd.GetScalars().GetData().(*schemapb.ScalarField_BoolData)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnBool(fd.GetFieldName(), data.BoolData.GetData()[begin:]), nil
|
|
}
|
|
return NewColumnBool(fd.GetFieldName(), data.BoolData.GetData()[begin:end]), nil
|
|
|
|
case schemapb.DataType_Int8:
|
|
data, ok := getIntData(fd)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
values := make([]int8, 0, len(data.IntData.GetData()))
|
|
for _, v := range data.IntData.GetData() {
|
|
values = append(values, int8(v))
|
|
}
|
|
|
|
if end < 0 {
|
|
return NewColumnInt8(fd.GetFieldName(), values[begin:]), nil
|
|
}
|
|
|
|
return NewColumnInt8(fd.GetFieldName(), values[begin:end]), nil
|
|
|
|
case schemapb.DataType_Int16:
|
|
data, ok := getIntData(fd)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
values := make([]int16, 0, len(data.IntData.GetData()))
|
|
for _, v := range data.IntData.GetData() {
|
|
values = append(values, int16(v))
|
|
}
|
|
if end < 0 {
|
|
return NewColumnInt16(fd.GetFieldName(), values[begin:]), nil
|
|
}
|
|
|
|
return NewColumnInt16(fd.GetFieldName(), values[begin:end]), nil
|
|
|
|
case schemapb.DataType_Int32:
|
|
data, ok := getIntData(fd)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnInt32(fd.GetFieldName(), data.IntData.GetData()[begin:]), nil
|
|
}
|
|
return NewColumnInt32(fd.GetFieldName(), data.IntData.GetData()[begin:end]), nil
|
|
|
|
case schemapb.DataType_Int64:
|
|
data, ok := fd.GetScalars().GetData().(*schemapb.ScalarField_LongData)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnInt64(fd.GetFieldName(), data.LongData.GetData()[begin:]), nil
|
|
}
|
|
return NewColumnInt64(fd.GetFieldName(), data.LongData.GetData()[begin:end]), nil
|
|
|
|
case schemapb.DataType_Float:
|
|
data, ok := fd.GetScalars().GetData().(*schemapb.ScalarField_FloatData)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnFloat(fd.GetFieldName(), data.FloatData.GetData()[begin:]), nil
|
|
}
|
|
return NewColumnFloat(fd.GetFieldName(), data.FloatData.GetData()[begin:end]), nil
|
|
|
|
case schemapb.DataType_Double:
|
|
data, ok := fd.GetScalars().GetData().(*schemapb.ScalarField_DoubleData)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnDouble(fd.GetFieldName(), data.DoubleData.GetData()[begin:]), nil
|
|
}
|
|
return NewColumnDouble(fd.GetFieldName(), data.DoubleData.GetData()[begin:end]), nil
|
|
|
|
case schemapb.DataType_String:
|
|
data, ok := fd.GetScalars().GetData().(*schemapb.ScalarField_StringData)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnString(fd.GetFieldName(), data.StringData.GetData()[begin:]), nil
|
|
}
|
|
return NewColumnString(fd.GetFieldName(), data.StringData.GetData()[begin:end]), nil
|
|
|
|
case schemapb.DataType_VarChar:
|
|
data, ok := fd.GetScalars().GetData().(*schemapb.ScalarField_StringData)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnVarChar(fd.GetFieldName(), data.StringData.GetData()[begin:]), nil
|
|
}
|
|
return NewColumnVarChar(fd.GetFieldName(), data.StringData.GetData()[begin:end]), nil
|
|
|
|
case schemapb.DataType_Array:
|
|
data := fd.GetScalars().GetArrayData()
|
|
if data == nil {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
var arrayData []*schemapb.ScalarField
|
|
if end < 0 {
|
|
arrayData = data.GetData()[begin:]
|
|
} else {
|
|
arrayData = data.GetData()[begin:end]
|
|
}
|
|
|
|
return parseArrayData(fd.GetFieldName(), data.GetElementType(), arrayData)
|
|
|
|
case schemapb.DataType_JSON:
|
|
data, ok := fd.GetScalars().GetData().(*schemapb.ScalarField_JsonData)
|
|
isDynamic := fd.GetIsDynamic()
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
if end < 0 {
|
|
return NewColumnJSONBytes(fd.GetFieldName(), data.JsonData.GetData()[begin:]).WithIsDynamic(isDynamic), nil
|
|
}
|
|
return NewColumnJSONBytes(fd.GetFieldName(), data.JsonData.GetData()[begin:end]).WithIsDynamic(isDynamic), nil
|
|
|
|
case schemapb.DataType_FloatVector:
|
|
vectors := fd.GetVectors()
|
|
x, ok := vectors.GetData().(*schemapb.VectorField_FloatVector)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
data := x.FloatVector.GetData()
|
|
dim := int(vectors.GetDim())
|
|
if end < 0 {
|
|
end = len(data) / dim
|
|
}
|
|
vector := make([][]float32, 0, end-begin) // shall not have remanunt
|
|
for i := begin; i < end; i++ {
|
|
v := make([]float32, dim)
|
|
copy(v, data[i*dim:(i+1)*dim])
|
|
vector = append(vector, v)
|
|
}
|
|
return NewColumnFloatVector(fd.GetFieldName(), dim, vector), nil
|
|
|
|
case schemapb.DataType_BinaryVector:
|
|
vectors := fd.GetVectors()
|
|
x, ok := vectors.GetData().(*schemapb.VectorField_BinaryVector)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
data := x.BinaryVector
|
|
if data == nil {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
dim := int(vectors.GetDim())
|
|
blen := dim / 8
|
|
if end < 0 {
|
|
end = len(data) / blen
|
|
}
|
|
vector := make([][]byte, 0, end-begin)
|
|
for i := begin; i < end; i++ {
|
|
v := make([]byte, blen)
|
|
copy(v, data[i*blen:(i+1)*blen])
|
|
vector = append(vector, v)
|
|
}
|
|
return NewColumnBinaryVector(fd.GetFieldName(), dim, vector), nil
|
|
|
|
case schemapb.DataType_Float16Vector:
|
|
vectors := fd.GetVectors()
|
|
x, ok := vectors.GetData().(*schemapb.VectorField_Float16Vector)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
data := x.Float16Vector
|
|
dim := int(vectors.GetDim())
|
|
if end < 0 {
|
|
end = len(data) / dim / 2
|
|
}
|
|
vector := make([][]byte, 0, end-begin)
|
|
for i := begin; i < end; i++ {
|
|
v := make([]byte, dim*2)
|
|
copy(v, data[i*dim*2:(i+1)*dim*2])
|
|
vector = append(vector, v)
|
|
}
|
|
return NewColumnFloat16Vector(fd.GetFieldName(), dim, vector), nil
|
|
|
|
case schemapb.DataType_BFloat16Vector:
|
|
vectors := fd.GetVectors()
|
|
x, ok := vectors.GetData().(*schemapb.VectorField_Bfloat16Vector)
|
|
if !ok {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
data := x.Bfloat16Vector
|
|
dim := int(vectors.GetDim())
|
|
if end < 0 {
|
|
end = len(data) / dim / 2
|
|
}
|
|
vector := make([][]byte, 0, end-begin) // shall not have remanunt
|
|
for i := begin; i < end; i++ {
|
|
v := make([]byte, dim*2)
|
|
copy(v, data[i*dim*2:(i+1)*dim*2])
|
|
vector = append(vector, v)
|
|
}
|
|
return NewColumnBFloat16Vector(fd.GetFieldName(), dim, vector), nil
|
|
case schemapb.DataType_SparseFloatVector:
|
|
sparseVectors := fd.GetVectors().GetSparseFloatVector()
|
|
if sparseVectors == nil {
|
|
return nil, errFieldDataTypeNotMatch
|
|
}
|
|
data := sparseVectors.Contents
|
|
if end < 0 {
|
|
end = len(data)
|
|
}
|
|
data = data[begin:end]
|
|
vectors := make([]entity.SparseEmbedding, 0, len(data))
|
|
for _, bs := range data {
|
|
vector, err := entity.DeserializeSliceSparseEmbedding(bs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
vectors = append(vectors, vector)
|
|
}
|
|
return NewColumnSparseVectors(fd.GetFieldName(), vectors), nil
|
|
default:
|
|
return nil, fmt.Errorf("unsupported data type %s", fd.GetType())
|
|
}
|
|
}
|
|
|
|
func parseArrayData(fieldName string, elementType schemapb.DataType, fieldDataList []*schemapb.ScalarField) (Column, error) {
|
|
switch elementType {
|
|
case schemapb.DataType_Bool:
|
|
var data [][]bool
|
|
for _, fd := range fieldDataList {
|
|
data = append(data, fd.GetBoolData().GetData())
|
|
}
|
|
return NewColumnBoolArray(fieldName, data), nil
|
|
|
|
case schemapb.DataType_Int8:
|
|
var data [][]int8
|
|
for _, fd := range fieldDataList {
|
|
raw := fd.GetIntData().GetData()
|
|
row := make([]int8, 0, len(raw))
|
|
for _, item := range raw {
|
|
row = append(row, int8(item))
|
|
}
|
|
data = append(data, row)
|
|
}
|
|
return NewColumnInt8Array(fieldName, data), nil
|
|
|
|
case schemapb.DataType_Int16:
|
|
var data [][]int16
|
|
for _, fd := range fieldDataList {
|
|
raw := fd.GetIntData().GetData()
|
|
row := make([]int16, 0, len(raw))
|
|
for _, item := range raw {
|
|
row = append(row, int16(item))
|
|
}
|
|
data = append(data, row)
|
|
}
|
|
return NewColumnInt16Array(fieldName, data), nil
|
|
|
|
case schemapb.DataType_Int32:
|
|
var data [][]int32
|
|
for _, fd := range fieldDataList {
|
|
data = append(data, fd.GetIntData().GetData())
|
|
}
|
|
return NewColumnInt32Array(fieldName, data), nil
|
|
|
|
case schemapb.DataType_Int64:
|
|
var data [][]int64
|
|
for _, fd := range fieldDataList {
|
|
data = append(data, fd.GetLongData().GetData())
|
|
}
|
|
return NewColumnInt64Array(fieldName, data), nil
|
|
|
|
case schemapb.DataType_Float:
|
|
var data [][]float32
|
|
for _, fd := range fieldDataList {
|
|
data = append(data, fd.GetFloatData().GetData())
|
|
}
|
|
return NewColumnFloatArray(fieldName, data), nil
|
|
|
|
case schemapb.DataType_Double:
|
|
var data [][]float64
|
|
for _, fd := range fieldDataList {
|
|
data = append(data, fd.GetDoubleData().GetData())
|
|
}
|
|
return NewColumnDoubleArray(fieldName, data), nil
|
|
|
|
case schemapb.DataType_VarChar, schemapb.DataType_String:
|
|
var data [][]string
|
|
for _, fd := range fieldDataList {
|
|
strs := fd.GetStringData().GetData()
|
|
bytesData := make([]string, 0, len(strs))
|
|
bytesData = append(bytesData, strs...)
|
|
|
|
data = append(data, bytesData)
|
|
}
|
|
|
|
return NewColumnVarCharArray(fieldName, data), nil
|
|
|
|
default:
|
|
return nil, fmt.Errorf("unsupported element type %s", elementType)
|
|
}
|
|
}
|
|
|
|
// getIntData get int32 slice from result field data
|
|
// also handles LongData bug (see also https://github.com/milvus-io/milvus/issues/23850)
|
|
func getIntData(fd *schemapb.FieldData) (*schemapb.ScalarField_IntData, bool) {
|
|
switch data := fd.GetScalars().GetData().(type) {
|
|
case *schemapb.ScalarField_IntData:
|
|
return data, true
|
|
case *schemapb.ScalarField_LongData:
|
|
// only alway empty LongData for backward compatibility
|
|
if len(data.LongData.GetData()) == 0 {
|
|
return &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{},
|
|
}, true
|
|
}
|
|
return nil, false
|
|
default:
|
|
return nil, false
|
|
}
|
|
}
|