Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

926 lines
29 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 funcutil
import (
"bytes"
"context"
"encoding/binary"
"encoding/json"
"fmt"
"net"
"reflect"
"regexp"
"strconv"
"strings"
"time"
"github.com/cockroachdb/errors"
"go.uber.org/atomic"
"google.golang.org/grpc/codes"
grpcStatus "google.golang.org/grpc/status"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/util"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
const (
ControlChannelSuffix = "vcchan" // is the suffix of the virtual control channel
)
// PreferIPv6LocalIP controls whether IPv6 addresses are preferred when selecting local IPs.
var PreferIPv6LocalIP atomic.Bool
// CheckGrpcReady wait for context timeout, or wait 100ms then send nil to targetCh
func CheckGrpcReady(ctx context.Context, targetCh chan error) {
timer := time.NewTimer(100 * time.Millisecond)
defer timer.Stop()
select {
case <-timer.C:
targetCh <- nil
case <-ctx.Done():
return
}
}
// GetIP return the ip address
func GetIP(ip string) string {
if len(ip) == 0 {
return GetLocalIP()
}
// Support setting CIDR in the IP field to match interfaces based on CIDR. For example: 192.168.0.0/16
_, ipnet, err := net.ParseCIDR(ip)
if err == nil {
addrs, err := net.InterfaceAddrs()
if err == nil {
for _, addr := range addrs {
addrip, ok := addr.(*net.IPNet)
if ok && ipnet.Contains(addrip.IP) {
return addrip.IP.String()
}
}
}
panic(merr.WrapErrParameterInvalidMsg(`Network port does not have an IP address that falls within the given CIDR range`))
}
netIP := net.ParseIP(ip)
// not a valid ip addr
if netIP == nil {
mlog.Warn(context.TODO(), "cannot parse input ip, treat it as hostname/service name", mlog.String("ip", ip))
return ip
}
// only localhost or unicast is acceptable
if netIP.IsUnspecified() {
panic(merr.WrapErrParameterInvalidMsg(`"%s" in param table is Unspecified IP address and cannot be used`))
}
if netIP.IsMulticast() || netIP.IsLinkLocalMulticast() || netIP.IsInterfaceLocalMulticast() {
panic(merr.WrapErrParameterInvalidMsg(`"%s" in param table is Multicast IP address and cannot be used`))
}
return ip
}
// GetLocalIP return the local ip address
func GetLocalIP() string {
addrs, err := net.InterfaceAddrs()
if err != nil {
mlog.Warn(context.TODO(), "Failed to get interface addresses", mlog.Err(err))
return "127.0.0.1"
}
preferIPv6 := PreferIPv6LocalIP.Load()
ip := getValidLocalIP(addrs, preferIPv6)
if len(ip) != 0 {
return ip
}
mlog.Warn(context.TODO(), "No valid local IP found, falling back to loopback")
return "127.0.0.1"
}
// GetValidLocalIP return the first valid local ip address
func GetValidLocalIP(addrs []net.Addr) string {
return getValidLocalIP(addrs, PreferIPv6LocalIP.Load())
}
type ipCategory int
const (
ipCategoryIPv4Public ipCategory = iota
ipCategoryIPv4Private
ipCategoryIPv6Public
ipCategoryIPv6Private
ipCategoryIPv6LinkLocal
)
var (
// Default priority: private first, IPv4 first
defaultIPPriority = []ipCategory{
ipCategoryIPv4Private,
ipCategoryIPv4Public,
ipCategoryIPv6Private,
ipCategoryIPv6Public,
ipCategoryIPv6LinkLocal,
}
// When IPv6 is preferred: private first, IPv6 first
preferIPv6Priority = []ipCategory{
ipCategoryIPv6Private,
ipCategoryIPv6Public,
ipCategoryIPv4Private,
ipCategoryIPv4Public,
ipCategoryIPv6LinkLocal,
}
)
func getValidLocalIP(addrs []net.Addr, preferIPv6 bool) string {
candidates := make(map[ipCategory]net.IP, 5)
for _, addr := range addrs {
ipNet, ok := addr.(*net.IPNet)
if !ok {
continue
}
category, valid := categorizeLocalIP(ipNet.IP)
if !valid {
continue
}
if _, exists := candidates[category]; !exists {
ipCopy := make(net.IP, len(ipNet.IP))
copy(ipCopy, ipNet.IP)
candidates[category] = ipCopy
}
}
priorities := defaultIPPriority
if preferIPv6 {
priorities = preferIPv6Priority
}
for _, category := range priorities {
if ip, exists := candidates[category]; exists {
result := formatLocalIP(ip)
mlog.Debug(context.TODO(), "Selected IP by priority",
mlog.String("ip", result),
mlog.String("categoryName", getCategoryName(category)))
return result
}
}
mlog.Warn(context.TODO(), "No valid IP found in candidates")
return ""
}
// getCategoryName returns human-readable name for IP category (for debugging)
func getCategoryName(category ipCategory) string {
switch category {
case ipCategoryIPv4Private:
return "IPv4Private"
case ipCategoryIPv4Public:
return "IPv4Public"
case ipCategoryIPv6Private:
return "IPv6Private"
case ipCategoryIPv6Public:
return "IPv6Public"
case ipCategoryIPv6LinkLocal:
return "IPv6LinkLocal"
default:
return "Unknown"
}
}
// JSONToMap parse the jsonic index parameters to map
func JSONToMap(mStr string) (map[string]string, error) {
buffer := make(map[string]any)
err := json.Unmarshal([]byte(mStr), &buffer)
if err != nil {
return nil, merr.Wrap(err, "unmarshal params failed")
}
ret := make(map[string]string)
for key, value := range buffer {
valueStr := fmt.Sprintf("%v", value)
ret[key] = valueStr
}
return ret, nil
}
func MapToJSON(m map[string]string) (string, error) {
// error won't happen here.
bs, err := json.Marshal(m)
if err != nil {
return "", err
}
return string(bs), nil
}
func JSONToRoleDetails(mStr string) (map[string](map[string]([](map[string]string))), error) {
buffer := make(map[string](map[string]([](map[string]string))), 0)
err := json.Unmarshal([]byte(mStr), &buffer)
if err != nil {
return nil, merr.Wrap(err, "unmarshal `builtinRoles.Roles` failed")
}
ret := make(map[string](map[string]([](map[string]string))), 0)
for role, privilegesJSON := range buffer {
ret[role] = make(map[string]([](map[string]string)), 0)
privilegesArray := make([]map[string]string, 0)
for _, privileges := range privilegesJSON[util.RoleConfigPrivileges] {
privilegesArray = append(privilegesArray, map[string]string{
util.RoleConfigObjectType: privileges[util.RoleConfigObjectType],
util.RoleConfigObjectName: privileges[util.RoleConfigObjectName],
util.RoleConfigPrivilege: privileges[util.RoleConfigPrivilege],
util.RoleConfigDBName: privileges[util.RoleConfigDBName],
})
}
ret[role]["privileges"] = privilegesArray
}
return ret, nil
}
func RoleDetailsToJSON(m map[string](map[string]([](map[string]string)))) []byte {
bs, _ := json.Marshal(m)
return bs
}
const (
// PulsarMaxMessageSizeKey is the key of config item
PulsarMaxMessageSizeKey = "maxMessageSize"
)
// GetAttrByKeyFromRepeatedKV return the value corresponding to key in kv pair
func GetAttrByKeyFromRepeatedKV(key string, kvs []*commonpb.KeyValuePair) (string, error) {
for _, kv := range kvs {
if kv.Key == key {
return kv.Value, nil
}
}
return "", merr.WrapErrParameterInvalidMsg("key %s not found", key)
}
// TryGetAttrByKeyFromRepeatedKV return the value corresponding to key in kv pair
// return false if key not exist
func TryGetAttrByKeyFromRepeatedKV(key string, kvs []*commonpb.KeyValuePair) (string, bool) {
for _, kv := range kvs {
if kv.Key == key {
return kv.Value, true
}
}
return "", false
}
// CheckCtxValid check if the context is valid
func CheckCtxValid(ctx context.Context) bool {
return ctx.Err() != context.DeadlineExceeded && ctx.Err() != context.Canceled
}
func GetVecFieldIDs(schema *schemapb.CollectionSchema) []int64 {
var vecFieldIDs []int64
for _, field := range schema.Fields {
if typeutil.IsVectorType(field.DataType) {
vecFieldIDs = append(vecFieldIDs, field.FieldID)
}
}
return vecFieldIDs
}
func String2KeyValuePair(v string) ([]*commonpb.KeyValuePair, error) {
m := make(map[string]string)
err := json.Unmarshal([]byte(v), &m)
if err != nil {
return nil, err
}
return Map2KeyValuePair(m), nil
}
func Map2KeyValuePair(datas map[string]string) []*commonpb.KeyValuePair {
results := make([]*commonpb.KeyValuePair, len(datas))
offset := 0
for key, value := range datas {
results[offset] = &commonpb.KeyValuePair{
Key: key,
Value: value,
}
offset++
}
return results
}
func KeyValuePair2Map(datas []*commonpb.KeyValuePair) map[string]string {
results := make(map[string]string)
for _, pair := range datas {
results[pair.Key] = pair.Value
}
return results
}
func ConvertToKeyValuePairPointer(datas []commonpb.KeyValuePair) []*commonpb.KeyValuePair {
var kvs []*commonpb.KeyValuePair
for i := 0; i < len(datas); i++ {
kvs = append(kvs, &datas[i])
}
return kvs
}
// GenChannelSubName generate subName to watch channel
func GenChannelSubName(prefix string, collectionID int64, nodeID int64) string {
return fmt.Sprintf("%s-%d-%d", prefix, collectionID, nodeID)
}
// CheckPortAvailable check if a port is available to be listened on
func CheckPortAvailable(port int) bool {
addr := ":" + strconv.Itoa(port)
listener, err := net.Listen("tcp", addr)
if listener != nil {
listener.Close()
}
return err == nil
}
// GetAvailablePort return an available port that can be listened on
func GetAvailablePort() int {
listener, err := net.Listen("tcp", ":0")
if err != nil {
panic(err)
}
defer listener.Close()
return listener.Addr().(*net.TCPAddr).Port
}
// IsPhysicalChannel checks if the channel is a physical channel
func IsPhysicalChannel(channel string) bool {
i := strings.LastIndex(channel, "_")
if i == -1 {
return true
}
return !strings.Contains(channel[i+1:], "v")
}
// IsControlChannel checks if the channel is a control channel
func IsControlChannel(channel string) bool {
return strings.HasSuffix(channel, ControlChannelSuffix)
}
// IsOnPhysicalChannel checks if the channel is on the physical channel.
func IsOnPhysicalChannel(channel string, physicalChannel string) bool {
return ToPhysicalChannel(channel) == physicalChannel
}
// ToPhysicalChannel get physical channel name from virtual channel name
func ToPhysicalChannel(vchannel string) string {
if IsPhysicalChannel(vchannel) {
return vchannel
}
index := strings.LastIndex(vchannel, "_")
if index < 0 {
return vchannel
}
return vchannel[:index]
}
// GetControlChannel returns the control channel name of the pchannel.
func GetControlChannel(pchannel string) string {
return fmt.Sprintf("%s_%s", pchannel, ControlChannelSuffix)
}
func GetVirtualChannel(pchannel string, collectionID int64, idx int) string {
return fmt.Sprintf("%s_%dv%d", pchannel, collectionID, idx)
}
// ConvertChannelName assembles channel name according to parameters.
func ConvertChannelName(chanName string, tokenFrom string, tokenTo string) (string, error) {
if tokenFrom == "" {
return "", merr.WrapErrParameterInvalidMsg("the tokenFrom is empty")
}
if !strings.Contains(chanName, tokenFrom) {
return "", merr.WrapErrParameterInvalidMsg("cannot find token '%s' in '%s'", tokenFrom, chanName)
}
return strings.Replace(chanName, tokenFrom, tokenTo, 1), nil
}
func GetCollectionIDFromVChannel(vChannelName string) int64 {
re := regexp.MustCompile(`.*_(\d+)v\d+`)
matches := re.FindStringSubmatch(vChannelName)
if len(matches) > 1 {
number, err := strconv.ParseInt(matches[1], 0, 64)
if err == nil {
return number
}
}
return -1
}
func getNumRowsOfScalarField(datas interface{}) uint64 {
realTypeDatas := reflect.ValueOf(datas)
return uint64(realTypeDatas.Len())
}
func getNumRowsOfArrayVectorField(datas interface{}) uint64 {
realTypeDatas := reflect.ValueOf(datas)
return uint64(realTypeDatas.Len())
}
func GetNumRowsOfFloatVectorField(fDatas []float32, dim int64) (uint64, error) {
if dim <= 0 {
return 0, merr.WrapErrParameterInvalidMsg("dim(%d) should be greater than 0", dim)
}
l := len(fDatas)
if int64(l)%dim != 0 {
return 0, merr.WrapErrParameterInvalidMsg("the length(%d) of float data should divide the dim(%d)", l, dim)
}
return uint64(int64(l) / dim), nil
}
func GetNumRowsOfBinaryVectorField(bDatas []byte, dim int64) (uint64, error) {
if dim <= 0 {
return 0, merr.WrapErrParameterInvalidMsg("dim(%d) should be greater than 0", dim)
}
if dim%8 != 0 {
return 0, merr.WrapErrParameterInvalidMsg("dim(%d) should divide 8", dim)
}
l := len(bDatas)
if (8*int64(l))%dim != 0 {
return 0, merr.WrapErrParameterInvalidMsg("the num(%d) of all bits should divide the dim(%d)", 8*l, dim)
}
return uint64((8 * int64(l)) / dim), nil
}
func GetNumRowsOfFloat16VectorField(f16Datas []byte, dim int64) (uint64, error) {
if dim <= 0 {
return 0, merr.WrapErrParameterInvalidMsg("dim(%d) should be greater than 0", dim)
}
l := len(f16Datas)
rowWidth := dim * 2
if int64(l)%rowWidth != 0 {
return 0, merr.WrapErrParameterInvalidMsg("the length(%d) of float16 data should divide the row width(%d)", l, rowWidth)
}
return uint64(int64(l) / rowWidth), nil
}
func GetNumRowsOfBFloat16VectorField(bf16Datas []byte, dim int64) (uint64, error) {
if dim <= 0 {
return 0, merr.WrapErrParameterInvalidMsg("dim(%d) should be greater than 0", dim)
}
l := len(bf16Datas)
rowWidth := dim * 2
if int64(l)%rowWidth != 0 {
return 0, merr.WrapErrParameterInvalidMsg("the length(%d) of bfloat data should divide the row width(%d)", l, rowWidth)
}
return uint64(int64(l) / rowWidth), nil
}
func GetNumRowsOfInt8VectorField(iDatas []byte, dim int64) (uint64, error) {
if dim <= 0 {
return 0, merr.WrapErrParameterInvalidMsg("dim(%d) should be greater than 0", dim)
}
l := len(iDatas)
if int64(l)%dim != 0 {
return 0, merr.WrapErrParameterInvalidMsg("the length(%d) of int8 data should divide the dim(%d)", l, dim)
}
return uint64(int64(l) / dim), nil
}
func CountValidRows(validData []bool) uint64 {
validRows := uint64(0)
for _, valid := range validData {
if valid {
validRows++
}
}
return validRows
}
func GetVectorFieldPhysicalRows(fieldName string, dataType schemapb.DataType, vectors *schemapb.VectorField) (uint64, error) {
if vectors == nil {
return 0, merr.WrapErrParameterInvalidMsg("nullable vector field %s requires vector data", fieldName)
}
return getVectorFieldPhysicalRowsWithDim(fieldName, dataType, vectors, vectors.GetDim())
}
func getVectorFieldPhysicalRowsWithDim(fieldName string, dataType schemapb.DataType, vectors *schemapb.VectorField, dim int64) (uint64, error) {
switch dataType {
case schemapb.DataType_FloatVector:
return GetNumRowsOfFloatVectorField(vectors.GetFloatVector().GetData(), dim)
case schemapb.DataType_BinaryVector:
return GetNumRowsOfBinaryVectorField(vectors.GetBinaryVector(), dim)
case schemapb.DataType_Float16Vector:
return GetNumRowsOfFloat16VectorField(vectors.GetFloat16Vector(), dim)
case schemapb.DataType_BFloat16Vector:
return GetNumRowsOfBFloat16VectorField(vectors.GetBfloat16Vector(), dim)
case schemapb.DataType_SparseFloatVector:
if vectors.GetSparseFloatVector() == nil {
return 0, nil
}
return uint64(len(vectors.GetSparseFloatVector().GetContents())), nil
case schemapb.DataType_Int8Vector:
return GetNumRowsOfInt8VectorField(vectors.GetInt8Vector(), dim)
default:
return 0, merr.WrapErrParameterInvalidMsg("unsupported nullable vector type %s", dataType)
}
}
func ValidateNullableVectorCompactRows(fieldName string, validData []bool, physicalRows uint64, logicalRows uint64, requireValidData bool) error {
if len(validData) == 0 {
if requireValidData {
return merr.WrapErrParameterInvalidMsg("nullable vector field %s requires valid_data", fieldName)
}
return nil
}
if logicalRows > 0 && uint64(len(validData)) != logicalRows {
return merr.WrapErrParameterInvalidMsg("nullable vector field %s valid_data length mismatch: valid_data=%d, logical rows=%d", fieldName, len(validData), logicalRows)
}
validRows := CountValidRows(validData)
if physicalRows != validRows {
return merr.WrapErrParameterInvalidMsg("nullable vector field %s has %d valid rows, but compact physical payload rows is %d", fieldName, validRows, physicalRows)
}
return nil
}
func ValidateNullableVectorFieldDataCompact(fieldData *schemapb.FieldData, logicalRows uint64, requireValidData bool) error {
if fieldData == nil || !typeutil.IsSupportedNullableVectorType(fieldData.GetType()) {
return nil
}
if len(fieldData.GetValidData()) == 0 && !requireValidData {
return nil
}
physicalRows, err := GetVectorFieldPhysicalRows(fieldData.GetFieldName(), fieldData.GetType(), fieldData.GetVectors())
if err != nil {
return err
}
return ValidateNullableVectorCompactRows(fieldData.GetFieldName(), fieldData.GetValidData(), physicalRows, logicalRows, requireValidData)
}
func ValidateNullableVectorFieldDataCompactWithDim(fieldData *schemapb.FieldData, logicalRows uint64, requireValidData bool, dim int64) error {
if fieldData == nil || fieldData.GetVectors() == nil || fieldData.GetVectors().GetDim() != 0 || dim <= 0 {
return ValidateNullableVectorFieldDataCompact(fieldData, logicalRows, requireValidData)
}
if !typeutil.IsSupportedNullableVectorType(fieldData.GetType()) {
return nil
}
if len(fieldData.GetValidData()) == 0 && !requireValidData {
return nil
}
physicalRows, err := getVectorFieldPhysicalRowsWithDim(fieldData.GetFieldName(), fieldData.GetType(), fieldData.GetVectors(), dim)
if err != nil {
return err
}
return ValidateNullableVectorCompactRows(fieldData.GetFieldName(), fieldData.GetValidData(), physicalRows, logicalRows, requireValidData)
}
// GetNumRowOfFieldDataWithSchema returns num of rows with schema specification.
func GetNumRowOfFieldDataWithSchema(fieldData *schemapb.FieldData, helper *typeutil.SchemaHelper) (uint64, error) {
var fieldNumRows uint64
var err error
fieldSchema, err := helper.GetFieldFromName(fieldData.GetFieldName())
if err != nil {
return 0, err
}
if len(fieldData.GetValidData()) > 0 && typeutil.IsSupportedNullableVectorType(fieldSchema.GetDataType()) {
dim := fieldData.GetVectors().GetDim()
if dim == 0 && fieldSchema.GetDataType() != schemapb.DataType_SparseFloatVector {
dim, err = typeutil.GetDim(fieldSchema)
if err != nil {
return 0, err
}
}
if err := ValidateNullableVectorFieldDataCompactWithDim(fieldData, uint64(len(fieldData.GetValidData())), false, dim); err != nil {
return 0, err
}
return uint64(len(fieldData.GetValidData())), nil
}
switch fieldSchema.GetDataType() {
case schemapb.DataType_Bool:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetBoolData().GetData())
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetIntData().GetData())
case schemapb.DataType_Int64:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetLongData().GetData())
case schemapb.DataType_Float:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetFloatData().GetData())
case schemapb.DataType_Double:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetDoubleData().GetData())
case schemapb.DataType_Timestamptz:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetTimestamptzData().GetData())
if fieldNumRows == 0 {
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetStringData().GetData())
}
case schemapb.DataType_String, schemapb.DataType_VarChar, schemapb.DataType_Text:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetStringData().GetData())
case schemapb.DataType_Array:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetArrayData().GetData())
case schemapb.DataType_JSON:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetJsonData().GetData())
case schemapb.DataType_Geometry:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetGeometryData().GetData())
if fieldNumRows == 0 {
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetGeometryWktData().GetData())
}
case schemapb.DataType_FloatVector:
if len(fieldData.GetValidData()) > 0 {
fieldNumRows = uint64(len(fieldData.GetValidData()))
} else {
dim := fieldData.GetVectors().GetDim()
fieldNumRows, err = GetNumRowsOfFloatVectorField(fieldData.GetVectors().GetFloatVector().GetData(), dim)
if err != nil {
return 0, err
}
}
case schemapb.DataType_BinaryVector:
if len(fieldData.GetValidData()) > 0 {
fieldNumRows = uint64(len(fieldData.GetValidData()))
} else {
dim := fieldData.GetVectors().GetDim()
fieldNumRows, err = GetNumRowsOfBinaryVectorField(fieldData.GetVectors().GetBinaryVector(), dim)
if err != nil {
return 0, err
}
}
case schemapb.DataType_Float16Vector:
if len(fieldData.GetValidData()) > 0 {
fieldNumRows = uint64(len(fieldData.GetValidData()))
} else {
dim := fieldData.GetVectors().GetDim()
fieldNumRows, err = GetNumRowsOfFloat16VectorField(fieldData.GetVectors().GetFloat16Vector(), dim)
if err != nil {
return 0, err
}
}
case schemapb.DataType_BFloat16Vector:
if len(fieldData.GetValidData()) > 0 {
fieldNumRows = uint64(len(fieldData.GetValidData()))
} else {
dim := fieldData.GetVectors().GetDim()
fieldNumRows, err = GetNumRowsOfBFloat16VectorField(fieldData.GetVectors().GetBfloat16Vector(), dim)
if err != nil {
return 0, err
}
}
case schemapb.DataType_SparseFloatVector:
if len(fieldData.GetValidData()) > 0 {
fieldNumRows = uint64(len(fieldData.GetValidData()))
} else {
fieldNumRows = uint64(len(fieldData.GetVectors().GetSparseFloatVector().GetContents()))
}
case schemapb.DataType_Int8Vector:
if len(fieldData.GetValidData()) > 0 {
fieldNumRows = uint64(len(fieldData.GetValidData()))
} else {
dim := fieldData.GetVectors().GetDim()
fieldNumRows, err = GetNumRowsOfInt8VectorField(fieldData.GetVectors().GetInt8Vector(), dim)
if err != nil {
return 0, err
}
}
case schemapb.DataType_ArrayOfVector:
if len(fieldData.GetValidData()) > 0 {
fieldNumRows = uint64(len(fieldData.GetValidData()))
} else {
fieldNumRows = getNumRowsOfArrayVectorField(fieldData.GetVectors().GetVectorArray().GetData())
}
default:
return 0, merr.WrapErrParameterInvalidMsg("%s is not supported now", fieldSchema.GetDataType())
}
return fieldNumRows, nil
}
// GetNumRowOfFieldData returns num of rows from the field data type
func GetNumRowOfFieldData(fieldData *schemapb.FieldData) (uint64, error) {
var fieldNumRows uint64
var err error
switch fieldType := fieldData.Field.(type) {
case *schemapb.FieldData_Scalars:
scalarField := fieldData.GetScalars()
switch scalarType := scalarField.Data.(type) {
case *schemapb.ScalarField_BoolData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetBoolData().Data)
case *schemapb.ScalarField_IntData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetIntData().Data)
case *schemapb.ScalarField_LongData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetLongData().Data)
case *schemapb.ScalarField_FloatData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetFloatData().Data)
case *schemapb.ScalarField_DoubleData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetDoubleData().Data)
case *schemapb.ScalarField_TimestamptzData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetTimestamptzData().Data)
case *schemapb.ScalarField_StringData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetStringData().Data)
case *schemapb.ScalarField_ArrayData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetArrayData().Data)
case *schemapb.ScalarField_JsonData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetJsonData().Data)
case *schemapb.ScalarField_GeometryData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetGeometryData().Data)
default:
return 0, merr.WrapErrParameterInvalidMsg("%s is not supported now", scalarType)
}
case *schemapb.FieldData_Vectors:
vectorField := fieldData.GetVectors()
if len(fieldData.GetValidData()) > 0 {
if err := ValidateNullableVectorFieldDataCompact(fieldData, uint64(len(fieldData.GetValidData())), false); err != nil {
return 0, err
}
return uint64(len(fieldData.GetValidData())), nil
}
switch vectorFieldType := vectorField.Data.(type) {
case *schemapb.VectorField_FloatVector:
dim := vectorField.GetDim()
fieldNumRows, err = GetNumRowsOfFloatVectorField(vectorField.GetFloatVector().Data, dim)
if err != nil {
return 0, err
}
case *schemapb.VectorField_BinaryVector:
dim := vectorField.GetDim()
fieldNumRows, err = GetNumRowsOfBinaryVectorField(vectorField.GetBinaryVector(), dim)
if err != nil {
return 0, err
}
case *schemapb.VectorField_Float16Vector:
dim := vectorField.GetDim()
fieldNumRows, err = GetNumRowsOfFloat16VectorField(vectorField.GetFloat16Vector(), dim)
if err != nil {
return 0, err
}
case *schemapb.VectorField_Bfloat16Vector:
dim := vectorField.GetDim()
fieldNumRows, err = GetNumRowsOfBFloat16VectorField(vectorField.GetBfloat16Vector(), dim)
if err != nil {
return 0, err
}
case *schemapb.VectorField_SparseFloatVector:
fieldNumRows = uint64(len(vectorField.GetSparseFloatVector().GetContents()))
case *schemapb.VectorField_Int8Vector:
dim := vectorField.GetDim()
fieldNumRows, err = GetNumRowsOfInt8VectorField(vectorField.GetInt8Vector(), dim)
if err != nil {
return 0, err
}
case *schemapb.VectorField_VectorArray:
fieldNumRows = getNumRowsOfArrayVectorField(vectorField.GetVectorArray().Data)
default:
return 0, merr.WrapErrParameterInvalidMsg("%s is not supported now", vectorFieldType)
}
default:
return 0, merr.WrapErrParameterInvalidMsg("%s is not supported now", fieldType)
}
return fieldNumRows, nil
}
// ReadBinary read byte slice as receiver.
func ReadBinary(endian binary.ByteOrder, bs []byte, receiver interface{}) error {
buf := bytes.NewReader(bs)
return binary.Read(buf, endian, receiver)
}
// IsGrpcErr checks whether err is instance of grpc status error.
func IsGrpcErr(err error, targets ...codes.Code) bool {
set := typeutil.NewSet[codes.Code](targets...)
for {
if err == nil {
return false
}
s, ok := grpcStatus.FromError(err)
if ok {
return set.Len() == 0 || set.Contain(s.Code())
}
err = errors.Unwrap(err)
}
}
func IsEmptyString(str string) bool {
return strings.TrimSpace(str) == ""
}
// HandleTenantForEtcdPrefix builds an etcd prefix for range scans (always ends with /).
// Two layers: HandleTenantForEtcdPrefix("a", "b") => "a/b/"
// Three layers: HandleTenantForEtcdPrefix("a", "b", "c") => "a/b/c/"
func HandleTenantForEtcdPrefix(prefix string, tenant string, subPrefixes ...string) string {
res := prefix
if tenant != "" {
res += "/" + tenant
}
for _, sub := range subPrefixes {
res += "/" + sub
}
res += "/"
return res
}
func IsRevoke(operateType milvuspb.OperatePrivilegeType) bool {
return operateType == milvuspb.OperatePrivilegeType_Revoke
}
func IsGrant(operateType milvuspb.OperatePrivilegeType) bool {
return operateType == milvuspb.OperatePrivilegeType_Grant
}
func EncodeUserRoleCache(user string, role string) string {
return fmt.Sprintf("%s/%s", user, role)
}
func DecodeUserRoleCache(cache string) (string, string, error) {
index := strings.LastIndex(cache, "/")
if index == -1 {
return "", "", merr.WrapErrParameterInvalidMsg("invalid param, cache: [%s]", cache)
}
user := cache[:index]
role := cache[index+1:]
return user, role, nil
}
// isIPv4Private checks if an IPv4 address is in RFC 1918 private ranges
func isIPv4Private(ip net.IP) bool {
ipv4 := ip.To4()
if ipv4 == nil {
return false
}
// RFC 1918 private address ranges:
// 10.0.0.0/8 (10.0.0.0 to 10.255.255.255)
// 172.16.0.0/12 (172.16.0.0 to 172.31.255.255)
// 192.168.0.0/16 (192.168.0.0 to 192.168.255.255)
return ipv4[0] == 10 ||
(ipv4[0] == 172 && ipv4[1] >= 16 && ipv4[1] <= 31) ||
(ipv4[0] == 192 && ipv4[1] == 168)
}
func categorizeLocalIP(ip net.IP) (ipCategory, bool) {
if ip == nil {
return 0, false
}
if ip.IsLoopback() || ip.IsInterfaceLocalMulticast() || ip.IsLinkLocalMulticast() || ip.IsMulticast() || ip.IsUnspecified() {
return 0, false
}
if ipv4 := ip.To4(); ipv4 != nil {
if !ip.IsGlobalUnicast() {
return 0, false
}
if isIPv4Private(ipv4) {
return ipCategoryIPv4Private, true
}
return ipCategoryIPv4Public, true
}
ipv6 := ip.To16()
if ipv6 == nil {
return 0, false
}
if ip.IsLinkLocalUnicast() {
return ipCategoryIPv6LinkLocal, true
}
if isIPv6Private(ipv6) {
return ipCategoryIPv6Private, true
}
if ip.IsGlobalUnicast() {
return ipCategoryIPv6Public, true
}
mlog.Debug(context.TODO(), "IP categorization: uncategorized IPv6", mlog.String("ip", ip.String()))
return 0, false
}
// isIPv6Private checks if an IPv6 address is in private ranges
func isIPv6Private(ip net.IP) bool {
ip = ip.To16()
if len(ip) != net.IPv6len {
return false
}
// RFC 4193 Unique Local Addresses (ULA): fc00::/7
return ip[0]&0xfe == 0xfc
}
func formatLocalIP(ip net.IP) string {
if ip.To4() != nil {
return ip.String()
}
return "[" + ip.String() + "]"
}