mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
enhance: Support fp32 to fp16/bf16 auto-conversion in search (#47133)
Add automatic conversion of fp32 search vectors to fp16/bf16 when the collection field uses lower precision types. This allows users to search fp16/bf16 collections with fp32 query vectors (which is what most embedding models output). Changes: - Add vector type matching utility (isVectorTypeMatch) - Add float16/bfloat16 range validation before conversion - Add ConvertPlaceholderGroup function for type conversion - Integrate conversion into normal search and hybrid search - Add GetFieldByID helper function in typeutil The conversion happens in the Proxy layer before requests are sent to QueryNode. Out-of-range values are rejected with clear error messages. issue: #47207 Signed-off-by: Li Liu <li.liu@zilliz.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.5
parent
92c2f83328
commit
8fb3fc927e
@@ -526,6 +526,11 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
|
||||
if typeutil.IsFieldSparseFloatVector(t.schema.CollectionSchema, internalSubReq.FieldId) {
|
||||
metrics.ProxySearchSparseNumNonZeros.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), t.collectionName, metrics.HybridSearchLabel, strconv.FormatInt(internalSubReq.FieldId, 10)).Observe(float64(typeutil.EstimateSparseVectorNNZFromPlaceholderGroup(internalSubReq.PlaceholderGroup, int(internalSubReq.GetNq()))))
|
||||
}
|
||||
// Convert placeholder group vector type if needed (e.g., fp32 -> fp16/bf16)
|
||||
internalSubReq.PlaceholderGroup, err = t.convertPlaceholderIfNeeded(subReq.GetPlaceholderGroup(), internalSubReq.FieldId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
t.SearchRequest.SubReqs[index] = internalSubReq
|
||||
t.queryInfos[index] = queryInfo
|
||||
log.Debug("proxy init search request",
|
||||
@@ -718,7 +723,11 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error {
|
||||
if typeutil.IsFieldSparseFloatVector(t.schema.CollectionSchema, t.SearchRequest.FieldId) {
|
||||
metrics.ProxySearchSparseNumNonZeros.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), t.collectionName, metrics.SearchLabel, strconv.FormatInt(t.SearchRequest.FieldId, 10)).Observe(float64(typeutil.EstimateSparseVectorNNZFromPlaceholderGroup(t.request.GetPlaceholderGroup(), int(t.request.GetNq()))))
|
||||
}
|
||||
t.SearchRequest.PlaceholderGroup = t.request.GetPlaceholderGroup()
|
||||
// Convert placeholder group vector type if needed (e.g., fp32 -> fp16/bf16)
|
||||
t.SearchRequest.PlaceholderGroup, err = t.convertPlaceholderIfNeeded(t.request.GetPlaceholderGroup(), t.SearchRequest.FieldId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
t.SearchRequest.Topk = queryInfo.GetTopk()
|
||||
t.SearchRequest.MetricType = queryInfo.GetMetricType()
|
||||
t.queryInfos = append(t.queryInfos, queryInfo)
|
||||
@@ -747,6 +756,15 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// convertPlaceholderIfNeeded converts fp32 vectors to fp16/bf16 if the target field uses lower precision.
|
||||
func (t *searchTask) convertPlaceholderIfNeeded(phgBytes []byte, fieldID int64) ([]byte, error) {
|
||||
field := typeutil.GetFieldByID(t.schema.CollectionSchema, fieldID)
|
||||
if field == nil {
|
||||
return phgBytes, nil
|
||||
}
|
||||
return ConvertPlaceholderGroup(phgBytes, field)
|
||||
}
|
||||
|
||||
func (t *searchTask) tryGeneratePlan(params []*commonpb.KeyValuePair, dsl string, exprTemplateValues map[string]*schemapb.TemplateValue) (*planpb.PlanNode, *planpb.QueryInfo, int64, bool, error) {
|
||||
annsFieldName, err := funcutil.GetAttrByKeyFromRepeatedKV(AnnsFieldKey, params)
|
||||
if err != nil || len(annsFieldName) == 0 {
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
// 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 proxy
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
|
||||
)
|
||||
|
||||
// Vector type compatibility matrix:
|
||||
//
|
||||
// Query Type | Field Type | Action
|
||||
// ----------------|------------------|------------------
|
||||
// FloatVector | FloatVector | pass through
|
||||
// FloatVector | Float16Vector | convert fp32->fp16
|
||||
// FloatVector | BFloat16Vector | convert fp32->bf16
|
||||
// Float16Vector | Float16Vector | pass through
|
||||
// BFloat16Vector | BFloat16Vector | pass through
|
||||
// BinaryVector | BinaryVector | pass through
|
||||
// Int8Vector | Int8Vector | pass through
|
||||
// SparseFloat | SparseFloat | pass through
|
||||
// * | * | error (incompatible)
|
||||
|
||||
// placeholderTypeToDataType maps PlaceholderType to corresponding DataType
|
||||
var placeholderTypeToDataType = map[commonpb.PlaceholderType]schemapb.DataType{
|
||||
commonpb.PlaceholderType_FloatVector: schemapb.DataType_FloatVector,
|
||||
commonpb.PlaceholderType_Float16Vector: schemapb.DataType_Float16Vector,
|
||||
commonpb.PlaceholderType_BFloat16Vector: schemapb.DataType_BFloat16Vector,
|
||||
commonpb.PlaceholderType_BinaryVector: schemapb.DataType_BinaryVector,
|
||||
commonpb.PlaceholderType_Int8Vector: schemapb.DataType_Int8Vector,
|
||||
commonpb.PlaceholderType_SparseFloatVector: schemapb.DataType_SparseFloatVector,
|
||||
}
|
||||
|
||||
// isVectorTypeMatch checks if the placeholder type matches the field data type exactly.
|
||||
func isVectorTypeMatch(placeholderType commonpb.PlaceholderType, fieldType schemapb.DataType) bool {
|
||||
expectedDataType, ok := placeholderTypeToDataType[placeholderType]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return expectedDataType == fieldType
|
||||
}
|
||||
|
||||
const (
|
||||
// float16MaxValue is the maximum representable value in float16
|
||||
float16MaxValue = 65504.0
|
||||
// float16MinPositive is the minimum positive normal value in float16
|
||||
float16MinPositive = 6.103515625e-5
|
||||
)
|
||||
|
||||
// validateFloat32ForFloat16 checks if all float32 values can be safely converted to float16.
|
||||
// Returns error if any value exceeds float16 range or underflows.
|
||||
func validateFloat32ForFloat16(values []float32) error {
|
||||
for i, v := range values {
|
||||
absV := math.Abs(float64(v))
|
||||
|
||||
// Check overflow
|
||||
if absV > float16MaxValue {
|
||||
return fmt.Errorf("value at dimension %d (%v) exceeds float16 range [-65504, 65504]", i, v)
|
||||
}
|
||||
|
||||
// Check underflow (non-zero values smaller than min positive)
|
||||
if v != 0 && absV < float16MinPositive {
|
||||
return fmt.Errorf("value at dimension %d (%v) underflows float16 precision (min abs value: %v)", i, v, float16MinPositive)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateFloat32ForBFloat16 checks if all float32 values can be safely converted to bfloat16.
|
||||
// BFloat16 has the same exponent range as float32, so only need to check for Inf/NaN.
|
||||
func validateFloat32ForBFloat16(values []float32) error {
|
||||
for i, v := range values {
|
||||
if math.IsInf(float64(v), 0) {
|
||||
return fmt.Errorf("value at dimension %d is infinity, cannot convert to bfloat16", i)
|
||||
}
|
||||
if math.IsNaN(float64(v)) {
|
||||
return fmt.Errorf("value at dimension %d is NaN, cannot convert to bfloat16", i)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ConvertPlaceholderGroup checks and converts placeholder group vector types if needed.
|
||||
// If the placeholder type matches the field type, returns the original bytes unchanged.
|
||||
// If the placeholder is fp32 and field is fp16/bf16, converts the vectors.
|
||||
// Otherwise returns an error for incompatible types.
|
||||
func ConvertPlaceholderGroup(phgBytes []byte, fieldSchema *schemapb.FieldSchema) ([]byte, error) {
|
||||
var phg commonpb.PlaceholderGroup
|
||||
if err := proto.Unmarshal(phgBytes, &phg); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal placeholder group: %w", err)
|
||||
}
|
||||
|
||||
if len(phg.Placeholders) == 0 {
|
||||
return phgBytes, nil
|
||||
}
|
||||
|
||||
placeholder := phg.Placeholders[0]
|
||||
fieldType := fieldSchema.GetDataType()
|
||||
|
||||
// Check if types already match
|
||||
if isVectorTypeMatch(placeholder.Type, fieldType) {
|
||||
return phgBytes, nil
|
||||
}
|
||||
|
||||
// If placeholder is not a vector type (e.g., VarChar for text embedding), pass through.
|
||||
// Let downstream logic handle non-vector placeholders.
|
||||
if _, isVectorType := placeholderTypeToDataType[placeholder.Type]; !isVectorType {
|
||||
return phgBytes, nil
|
||||
}
|
||||
|
||||
// Only handle fp32 -> fp16/bf16 conversion.
|
||||
// For other field types (e.g., SparseFloatVector, BinaryVector), pass through
|
||||
// and let downstream logic handle the type mismatch with appropriate error messages.
|
||||
if fieldType != schemapb.DataType_Float16Vector && fieldType != schemapb.DataType_BFloat16Vector {
|
||||
return phgBytes, nil
|
||||
}
|
||||
|
||||
// Check if conversion is supported (fp32 -> fp16/bf16)
|
||||
if placeholder.Type != commonpb.PlaceholderType_FloatVector {
|
||||
return nil, fmt.Errorf("vector type must be the same: field type %s, search type %s",
|
||||
fieldType.String(), placeholder.Type.String())
|
||||
}
|
||||
|
||||
switch fieldType {
|
||||
case schemapb.DataType_Float16Vector:
|
||||
return convertPlaceholder(&phg, validateFloat32ForFloat16, typeutil.Float32ArrayToFloat16Bytes, commonpb.PlaceholderType_Float16Vector)
|
||||
case schemapb.DataType_BFloat16Vector:
|
||||
return convertPlaceholder(&phg, validateFloat32ForBFloat16, typeutil.Float32ArrayToBFloat16Bytes, commonpb.PlaceholderType_BFloat16Vector)
|
||||
default:
|
||||
// This should never be reached due to the check above, but keep for safety
|
||||
return phgBytes, nil
|
||||
}
|
||||
}
|
||||
|
||||
// convertPlaceholder converts fp32 vectors in placeholder to the target type.
|
||||
func convertPlaceholder(
|
||||
phg *commonpb.PlaceholderGroup,
|
||||
validateFn func([]float32) error,
|
||||
convertFn func([]float32) []byte,
|
||||
targetType commonpb.PlaceholderType,
|
||||
) ([]byte, error) {
|
||||
placeholder := phg.Placeholders[0]
|
||||
convertedValues := make([][]byte, len(placeholder.Values))
|
||||
|
||||
for i, valueBytes := range placeholder.Values {
|
||||
floats, err := bytesToFloat32Array(valueBytes)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse float32 vector at index %d: %w", i, err)
|
||||
}
|
||||
|
||||
if err := validateFn(floats); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
convertedValues[i] = convertFn(floats)
|
||||
}
|
||||
|
||||
placeholder.Type = targetType
|
||||
placeholder.Values = convertedValues
|
||||
|
||||
return proto.Marshal(phg)
|
||||
}
|
||||
|
||||
// bytesToFloat32Array converts byte slice to float32 array.
|
||||
func bytesToFloat32Array(data []byte) ([]float32, error) {
|
||||
if len(data)%4 != 0 {
|
||||
return nil, fmt.Errorf("invalid float32 vector data length: %d", len(data))
|
||||
}
|
||||
|
||||
dim := len(data) / 4
|
||||
result := make([]float32, dim)
|
||||
for i := 0; i < dim; i++ {
|
||||
bits := binary.LittleEndian.Uint32(data[i*4 : (i+1)*4])
|
||||
result[i] = math.Float32frombits(bits)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,368 @@
|
||||
// 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 proxy
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
|
||||
)
|
||||
|
||||
func TestIsVectorTypeMatch(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
placeholderType commonpb.PlaceholderType
|
||||
fieldType schemapb.DataType
|
||||
expected bool
|
||||
}{
|
||||
{"fp32 to fp32", commonpb.PlaceholderType_FloatVector, schemapb.DataType_FloatVector, true},
|
||||
{"fp16 to fp16", commonpb.PlaceholderType_Float16Vector, schemapb.DataType_Float16Vector, true},
|
||||
{"bf16 to bf16", commonpb.PlaceholderType_BFloat16Vector, schemapb.DataType_BFloat16Vector, true},
|
||||
{"binary to binary", commonpb.PlaceholderType_BinaryVector, schemapb.DataType_BinaryVector, true},
|
||||
{"int8 to int8", commonpb.PlaceholderType_Int8Vector, schemapb.DataType_Int8Vector, true},
|
||||
{"sparse to sparse", commonpb.PlaceholderType_SparseFloatVector, schemapb.DataType_SparseFloatVector, true},
|
||||
{"fp32 to fp16", commonpb.PlaceholderType_FloatVector, schemapb.DataType_Float16Vector, false},
|
||||
{"fp32 to bf16", commonpb.PlaceholderType_FloatVector, schemapb.DataType_BFloat16Vector, false},
|
||||
{"fp16 to fp32", commonpb.PlaceholderType_Float16Vector, schemapb.DataType_FloatVector, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := isVectorTypeMatch(tt.placeholderType, tt.fieldType)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateFloat32ForFloat16(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
values []float32
|
||||
expectErr bool
|
||||
errMsg string
|
||||
}{
|
||||
{
|
||||
name: "normal values",
|
||||
values: []float32{0.0, 1.0, -1.0, 0.5, -0.5},
|
||||
expectErr: false,
|
||||
},
|
||||
{
|
||||
name: "max boundary",
|
||||
values: []float32{65504.0, -65504.0},
|
||||
expectErr: false,
|
||||
},
|
||||
{
|
||||
name: "overflow positive",
|
||||
values: []float32{0.0, 65505.0},
|
||||
expectErr: true,
|
||||
errMsg: "exceeds float16 range",
|
||||
},
|
||||
{
|
||||
name: "overflow negative",
|
||||
values: []float32{-65505.0, 0.0},
|
||||
expectErr: true,
|
||||
errMsg: "exceeds float16 range",
|
||||
},
|
||||
{
|
||||
name: "underflow",
|
||||
values: []float32{0.0, 1e-9},
|
||||
expectErr: true,
|
||||
errMsg: "underflows float16",
|
||||
},
|
||||
{
|
||||
name: "zero is valid",
|
||||
values: []float32{0.0, 0.0, 0.0},
|
||||
expectErr: false,
|
||||
},
|
||||
{
|
||||
name: "min positive normal",
|
||||
values: []float32{6.104e-5},
|
||||
expectErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := validateFloat32ForFloat16(tt.values)
|
||||
if tt.expectErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), tt.errMsg)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateFloat32ForBFloat16(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
values []float32
|
||||
expectErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal values",
|
||||
values: []float32{0.0, 1.0, -1.0, 0.5, -0.5},
|
||||
expectErr: false,
|
||||
},
|
||||
{
|
||||
name: "large values within float32 range",
|
||||
values: []float32{1e38, -1e38},
|
||||
expectErr: false,
|
||||
},
|
||||
{
|
||||
name: "very small values",
|
||||
values: []float32{1e-38, -1e-38},
|
||||
expectErr: false,
|
||||
},
|
||||
{
|
||||
name: "infinity",
|
||||
values: []float32{float32(math.Inf(1))},
|
||||
expectErr: true,
|
||||
},
|
||||
{
|
||||
name: "negative infinity",
|
||||
values: []float32{float32(math.Inf(-1))},
|
||||
expectErr: true,
|
||||
},
|
||||
{
|
||||
name: "NaN",
|
||||
values: []float32{float32(math.NaN())},
|
||||
expectErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := validateFloat32ForBFloat16(tt.values)
|
||||
if tt.expectErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func createFloat32PlaceholderGroup(vectors [][]float32) []byte {
|
||||
values := make([][]byte, len(vectors))
|
||||
for i, vec := range vectors {
|
||||
values[i] = typeutil.Float32ArrayToBytes(vec)
|
||||
}
|
||||
|
||||
phg := &commonpb.PlaceholderGroup{
|
||||
Placeholders: []*commonpb.PlaceholderValue{{
|
||||
Tag: "$0",
|
||||
Type: commonpb.PlaceholderType_FloatVector,
|
||||
Values: values,
|
||||
}},
|
||||
}
|
||||
|
||||
bytes, _ := proto.Marshal(phg)
|
||||
return bytes
|
||||
}
|
||||
|
||||
func TestConvertPlaceholderGroupToFloat16(t *testing.T) {
|
||||
vectors := [][]float32{
|
||||
{0.1, 0.2, 0.3, 0.4},
|
||||
{0.5, 0.6, 0.7, 0.8},
|
||||
}
|
||||
phgBytes := createFloat32PlaceholderGroup(vectors)
|
||||
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_Float16Vector,
|
||||
}
|
||||
|
||||
convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.NoError(t, err)
|
||||
|
||||
var resultPhg commonpb.PlaceholderGroup
|
||||
err = proto.Unmarshal(convertedBytes, &resultPhg)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, commonpb.PlaceholderType_Float16Vector, resultPhg.Placeholders[0].Type)
|
||||
assert.Equal(t, 4*2, len(resultPhg.Placeholders[0].Values[0])) // 4 dimensions * 2 bytes
|
||||
}
|
||||
|
||||
func TestConvertPlaceholderGroupToBFloat16(t *testing.T) {
|
||||
vectors := [][]float32{
|
||||
{0.1, 0.2, 0.3, 0.4},
|
||||
}
|
||||
phgBytes := createFloat32PlaceholderGroup(vectors)
|
||||
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_BFloat16Vector,
|
||||
}
|
||||
|
||||
convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.NoError(t, err)
|
||||
|
||||
var resultPhg commonpb.PlaceholderGroup
|
||||
err = proto.Unmarshal(convertedBytes, &resultPhg)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, commonpb.PlaceholderType_BFloat16Vector, resultPhg.Placeholders[0].Type)
|
||||
assert.Equal(t, 4*2, len(resultPhg.Placeholders[0].Values[0]))
|
||||
}
|
||||
|
||||
func TestConvertPlaceholderGroupTypeMismatch(t *testing.T) {
|
||||
// Test: Float16Vector searching BFloat16Vector field -> should error
|
||||
// (Only fp32 -> fp16/bf16 conversion is supported)
|
||||
phg := &commonpb.PlaceholderGroup{
|
||||
Placeholders: []*commonpb.PlaceholderValue{{
|
||||
Tag: "$0",
|
||||
Type: commonpb.PlaceholderType_Float16Vector,
|
||||
Values: [][]byte{{0, 0, 0, 0}},
|
||||
}},
|
||||
}
|
||||
phgBytes, _ := proto.Marshal(phg)
|
||||
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_BFloat16Vector,
|
||||
}
|
||||
|
||||
_, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "vector type must be the same")
|
||||
}
|
||||
|
||||
func TestConvertPlaceholderGroupPassThrough(t *testing.T) {
|
||||
// Test: Float16Vector searching FloatVector field -> pass through
|
||||
// Let downstream handle the type mismatch
|
||||
phg := &commonpb.PlaceholderGroup{
|
||||
Placeholders: []*commonpb.PlaceholderValue{{
|
||||
Tag: "$0",
|
||||
Type: commonpb.PlaceholderType_Float16Vector,
|
||||
Values: [][]byte{{0, 0, 0, 0}},
|
||||
}},
|
||||
}
|
||||
phgBytes, _ := proto.Marshal(phg)
|
||||
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_FloatVector,
|
||||
}
|
||||
|
||||
result, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, phgBytes, result)
|
||||
}
|
||||
|
||||
func TestConvertPlaceholderGroupNoConversionNeeded(t *testing.T) {
|
||||
vectors := [][]float32{{0.1, 0.2, 0.3, 0.4}}
|
||||
phgBytes := createFloat32PlaceholderGroup(vectors)
|
||||
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_FloatVector,
|
||||
}
|
||||
|
||||
convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, phgBytes, convertedBytes)
|
||||
}
|
||||
|
||||
func TestConvertPlaceholderGroupOutOfRange(t *testing.T) {
|
||||
vectors := [][]float32{{0.1, 100000.0, 0.3, 0.4}}
|
||||
phgBytes := createFloat32PlaceholderGroup(vectors)
|
||||
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_Float16Vector,
|
||||
}
|
||||
|
||||
_, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "exceeds float16 range")
|
||||
}
|
||||
|
||||
func TestConvertPlaceholderGroupIntegration(t *testing.T) {
|
||||
dim := int64(128)
|
||||
nq := 10
|
||||
|
||||
// Create fp32 vectors (simulating typical embedding model output)
|
||||
vectors := make([][]float32, nq)
|
||||
for i := 0; i < nq; i++ {
|
||||
vectors[i] = make([]float32, dim)
|
||||
for j := int64(0); j < dim; j++ {
|
||||
vectors[i][j] = float32(i+1) * 0.01 * float32(j+1) / float32(dim)
|
||||
}
|
||||
}
|
||||
|
||||
// Create placeholder group
|
||||
values := make([][]byte, nq)
|
||||
for i, vec := range vectors {
|
||||
values[i] = typeutil.Float32ArrayToBytes(vec)
|
||||
}
|
||||
|
||||
phg := &commonpb.PlaceholderGroup{
|
||||
Placeholders: []*commonpb.PlaceholderValue{{
|
||||
Tag: "$0",
|
||||
Type: commonpb.PlaceholderType_FloatVector,
|
||||
Values: values,
|
||||
}},
|
||||
}
|
||||
phgBytes, err := proto.Marshal(phg)
|
||||
assert.NoError(t, err)
|
||||
|
||||
t.Run("convert to float16", func(t *testing.T) {
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_Float16Vector,
|
||||
}
|
||||
|
||||
convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.NoError(t, err)
|
||||
|
||||
var resultPhg commonpb.PlaceholderGroup
|
||||
err = proto.Unmarshal(convertedBytes, &resultPhg)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, commonpb.PlaceholderType_Float16Vector, resultPhg.Placeholders[0].Type)
|
||||
assert.Equal(t, nq, len(resultPhg.Placeholders[0].Values))
|
||||
assert.Equal(t, int(dim)*2, len(resultPhg.Placeholders[0].Values[0])) // 2 bytes per float16
|
||||
})
|
||||
|
||||
t.Run("convert to bfloat16", func(t *testing.T) {
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_BFloat16Vector,
|
||||
}
|
||||
|
||||
convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.NoError(t, err)
|
||||
|
||||
var resultPhg commonpb.PlaceholderGroup
|
||||
err = proto.Unmarshal(convertedBytes, &resultPhg)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, commonpb.PlaceholderType_BFloat16Vector, resultPhg.Placeholders[0].Type)
|
||||
assert.Equal(t, nq, len(resultPhg.Placeholders[0].Values))
|
||||
assert.Equal(t, int(dim)*2, len(resultPhg.Placeholders[0].Values[0])) // 2 bytes per bfloat16
|
||||
})
|
||||
|
||||
t.Run("no conversion for matching type", func(t *testing.T) {
|
||||
fieldSchema := &schemapb.FieldSchema{
|
||||
DataType: schemapb.DataType_FloatVector,
|
||||
}
|
||||
|
||||
convertedBytes, err := ConvertPlaceholderGroup(phgBytes, fieldSchema)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, phgBytes, convertedBytes) // Should be unchanged
|
||||
})
|
||||
}
|
||||
@@ -1917,6 +1917,16 @@ func GetFieldByName(schema *schemapb.CollectionSchema, fieldName string) *schema
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetFieldByID returns the field schema with the given field ID, or nil if not found.
|
||||
func GetFieldByID(schema *schemapb.CollectionSchema, fieldID int64) *schemapb.FieldSchema {
|
||||
for _, field := range schema.GetFields() {
|
||||
if field.GetFieldID() == fieldID {
|
||||
return field
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func IsPrimaryFieldDataExist(datas []*schemapb.FieldData, primaryFieldSchema *schemapb.FieldSchema) bool {
|
||||
primaryFieldID := primaryFieldSchema.FieldID
|
||||
primaryFieldName := primaryFieldSchema.Name
|
||||
|
||||
Reference in New Issue
Block a user