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:
Li Liu
2026-01-21 14:45:30 +08:00
committed by GitHub
co-authored by Claude Opus 4.5
parent 92c2f83328
commit 8fb3fc927e
4 changed files with 595 additions and 1 deletions
+19 -1
View File
@@ -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 {
+198
View File
@@ -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
}
+368
View File
@@ -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
})
}
+10
View File
@@ -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