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
3526 lines
127 KiB
Go
3526 lines
127 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 proxy
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"reflect"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
"unicode/utf8"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/samber/lo"
|
|
"go.opentelemetry.io/otel"
|
|
"golang.org/x/crypto/bcrypt"
|
|
"golang.org/x/time/rate"
|
|
"google.golang.org/grpc/metadata"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
"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/internal/agg"
|
|
"github.com/milvus-io/milvus/internal/json"
|
|
"github.com/milvus-io/milvus/internal/parser/planparserv2"
|
|
"github.com/milvus-io/milvus/internal/proxy/privilege"
|
|
"github.com/milvus-io/milvus/internal/types"
|
|
"github.com/milvus-io/milvus/internal/util/function/embedding"
|
|
"github.com/milvus-io/milvus/internal/util/function/models"
|
|
"github.com/milvus-io/milvus/internal/util/hookutil"
|
|
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
|
|
"github.com/milvus-io/milvus/internal/util/segcore"
|
|
typeutil2 "github.com/milvus-io/milvus/internal/util/typeutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/metrics"
|
|
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
|
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/contextutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/crypto"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/metric"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/rbacutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/requestutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/timestamptz"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
|
)
|
|
|
|
const (
|
|
strongTS = 0
|
|
boundedTS = 2
|
|
|
|
// enableMultipleVectorFields indicates whether to enable multiple vector fields.
|
|
enableMultipleVectorFields = true
|
|
|
|
defaultMaxSearchRequest = 1024
|
|
|
|
// DefaultArithmeticIndexType name of default index type for scalar field
|
|
DefaultArithmeticIndexType = indexparamcheck.IndexINVERTED
|
|
|
|
// DefaultStringIndexType name of default index type for varChar/string field
|
|
DefaultStringIndexType = indexparamcheck.IndexINVERTED
|
|
)
|
|
|
|
var logger = mlog.With(mlog.String("role", typeutil.ProxyRole))
|
|
|
|
// transformStructFieldNames transforms struct field names to structName[fieldName] format
|
|
// This ensures global uniqueness while allowing same field names across different structs
|
|
func transformStructFieldNames(schema *schemapb.CollectionSchema) error {
|
|
for _, structArrayField := range schema.StructArrayFields {
|
|
structName := structArrayField.Name
|
|
for _, field := range structArrayField.Fields {
|
|
// Create transformed name: structName[fieldName]
|
|
newName := typeutil.ConcatStructFieldName(structName, field.Name)
|
|
field.Name = newName
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// restoreStructFieldNames restores original field names from structName[fieldName] format
|
|
// This is used when returning schema information to users (e.g., in describe collection)
|
|
func restoreStructFieldNames(schema *schemapb.CollectionSchema) error {
|
|
for _, structArrayField := range schema.StructArrayFields {
|
|
structName := structArrayField.Name
|
|
expectedPrefix := structName + "["
|
|
|
|
for _, field := range structArrayField.Fields {
|
|
if strings.HasPrefix(field.Name, expectedPrefix) && strings.HasSuffix(field.Name, "]") {
|
|
// Extract fieldName: remove "structName[" prefix and "]" suffix
|
|
field.Name = field.Name[len(expectedPrefix) : len(field.Name)-1]
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// extractOriginalFieldName extracts the original field name from structName[fieldName] format
|
|
// This function should only be called on transformed struct field names
|
|
func extractOriginalFieldName(transformedName string) (string, error) {
|
|
idx := strings.Index(transformedName, "[")
|
|
if idx == -1 {
|
|
return "", merr.WrapErrParameterInvalidMsg("not a transformed struct field name: %s", transformedName)
|
|
}
|
|
|
|
if !strings.HasSuffix(transformedName, "]") {
|
|
return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, missing closing bracket", transformedName)
|
|
}
|
|
|
|
if idx == 0 {
|
|
return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, missing struct name", transformedName)
|
|
}
|
|
|
|
fieldName := transformedName[idx+1 : len(transformedName)-1]
|
|
if fieldName == "" {
|
|
return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, empty field name", transformedName)
|
|
}
|
|
|
|
return fieldName, nil
|
|
}
|
|
|
|
// isAlpha check if c is alpha.
|
|
func isAlpha(c uint8) bool {
|
|
if (c < 'A' || c > 'Z') && (c < 'a' || c > 'z') {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
// isNumber check if c is a number.
|
|
func isNumber(c uint8) bool {
|
|
if c < '0' || c > '9' {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
// check run analyzer params when collection name was set
|
|
func validateRunAnalyzer(req *milvuspb.RunAnalyzerRequest) error {
|
|
if req.GetAnalyzerParams() != "" {
|
|
return merr.WrapErrParameterInvalidMsg("run analyzer can't use analyzer params and (collection,field) in same time")
|
|
}
|
|
|
|
if req.GetFieldName() == "" {
|
|
return merr.WrapErrParameterInvalidMsg("must set field name when collection name was set")
|
|
}
|
|
|
|
if req.GetAnalyzerNames() != nil {
|
|
if len(req.GetAnalyzerNames()) != 1 && len(req.GetAnalyzerNames()) != len(req.GetPlaceholder()) {
|
|
return merr.WrapErrParameterInvalidMsg("only support set one analyzer name for all text or set analyzer name for each text, but now analzer name num: %d, text num: %d",
|
|
len(req.GetAnalyzerNames()), len(req.GetPlaceholder()))
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func validateMaxQueryResultWindow(offset int64, limit int64, largeTopKEnabled bool) error {
|
|
if offset < 0 {
|
|
return merr.WrapErrParameterInvalidMsg("%s [%d] is invalid, should be gte than 0", OffsetKey, offset)
|
|
}
|
|
if limit <= 0 {
|
|
return merr.WrapErrParameterInvalidMsg("%s [%d] is invalid, should be greater than 0", LimitKey, limit)
|
|
}
|
|
|
|
depth := offset + limit
|
|
maxQueryResultWindow := Params.QuotaConfig.MaxQueryResultWindow.GetAsInt64()
|
|
if largeTopKEnabled {
|
|
maxQueryResultWindow = Params.QuotaConfig.LargeMaxQueryResultWindow.GetAsInt64()
|
|
}
|
|
if depth <= 0 || depth > maxQueryResultWindow {
|
|
return merr.WrapErrParameterInvalidMsg("(offset+limit) should be in range [1, %d], but got %d", maxQueryResultWindow, depth)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateLimit(limit int64, largeTopKEnabled bool) error {
|
|
topKLimit := Params.QuotaConfig.TopKLimit.GetAsInt64()
|
|
if largeTopKEnabled {
|
|
topKLimit = Params.QuotaConfig.LargeTopKLimit.GetAsInt64()
|
|
}
|
|
if limit <= 0 || limit > topKLimit {
|
|
return merr.WrapErrParameterInvalidMsg("it should be in range [1, %d], but got %d", topKLimit, limit)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateNQLimit(limit int64) error {
|
|
nqLimit := Params.QuotaConfig.NQLimit.GetAsInt64()
|
|
if limit <= 0 || limit > nqLimit {
|
|
return merr.WrapErrParameterInvalidMsg("nq (number of search vector per search request) should be in range [1, %d], but got %d", nqLimit, limit)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateCollectionNameOrAlias(entity, entityType string) error {
|
|
if entity == "" {
|
|
return merr.WrapErrParameterInvalidMsg("collection %s should not be empty", entityType)
|
|
}
|
|
|
|
invalidMsg := fmt.Sprintf("Invalid collection %s: %s. ", entityType, entity)
|
|
if len(entity) > Params.ProxyCfg.MaxNameLength.GetAsInt() {
|
|
return merr.WrapErrParameterInvalidMsg("%s the length of a collection %s must be less than %s characters", invalidMsg, entityType,
|
|
Params.ProxyCfg.MaxNameLength.GetValue())
|
|
}
|
|
|
|
firstChar := entity[0]
|
|
if firstChar != '_' && !isAlpha(firstChar) {
|
|
return merr.WrapErrParameterInvalidMsg("%s the first character of a collection %s must be an underscore or letter", invalidMsg, entityType)
|
|
}
|
|
|
|
for i := 1; i < len(entity); i++ {
|
|
c := entity[i]
|
|
if c != '_' && !isAlpha(c) && !isNumber(c) {
|
|
return merr.WrapErrParameterInvalidMsg("%s collection %s can only contain numbers, letters and underscores", invalidMsg, entityType)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ValidatePrivilegeGroupName(groupName string) error {
|
|
if groupName == "" {
|
|
return merr.WrapErrPrivilegeGroupNameInvalid("privilege group name should not be empty")
|
|
}
|
|
|
|
if len(groupName) > Params.ProxyCfg.MaxNameLength.GetAsInt() {
|
|
return merr.WrapErrPrivilegeGroupNameInvalid(
|
|
"the length of a privilege group name %s must be less than %s characters", groupName, Params.ProxyCfg.MaxNameLength.GetValue())
|
|
}
|
|
|
|
firstChar := groupName[0]
|
|
if firstChar != '_' && !isAlpha(firstChar) {
|
|
return merr.WrapErrPrivilegeGroupNameInvalid(
|
|
"the first character of a privilege group name %s must be an underscore or letter", groupName)
|
|
}
|
|
|
|
for i := 1; i < len(groupName); i++ {
|
|
c := groupName[i]
|
|
if c != '_' && !isAlpha(c) && !isNumber(c) {
|
|
return merr.WrapErrParameterInvalidMsg(
|
|
"privilege group name %s can only contain numbers, letters and underscores", groupName)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ValidateResourceGroupName(entity string) error {
|
|
if entity == "" {
|
|
return merr.WrapErrParameterMissingMsg("resource group name couldn't be empty")
|
|
}
|
|
|
|
invalidMsg := fmt.Sprintf("Invalid resource group name %s.", entity)
|
|
if len(entity) > Params.ProxyCfg.MaxNameLength.GetAsInt() {
|
|
return merr.WrapErrParameterInvalidMsg("%s the length of a resource group name must be less than %s characters",
|
|
invalidMsg, Params.ProxyCfg.MaxNameLength.GetValue())
|
|
}
|
|
|
|
firstChar := entity[0]
|
|
if firstChar != '_' && !isAlpha(firstChar) {
|
|
return merr.WrapErrParameterInvalidMsg("%s the first character of a resource group name must be an underscore or letter", invalidMsg)
|
|
}
|
|
|
|
for i := 1; i < len(entity); i++ {
|
|
c := entity[i]
|
|
if c != '_' && !isAlpha(c) && !isNumber(c) {
|
|
return merr.WrapErrParameterInvalidMsg("%s resource group name can only contain numbers, letters and underscores", invalidMsg)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ValidateDatabaseName(dbName string) error {
|
|
if dbName == "" {
|
|
return merr.WrapErrDatabaseNameInvalid(dbName, "database name couldn't be empty")
|
|
}
|
|
|
|
if len(dbName) > Params.ProxyCfg.MaxNameLength.GetAsInt() {
|
|
return merr.WrapErrDatabaseNameInvalid(dbName,
|
|
fmt.Sprintf("the length of a database name must be less than %d characters", Params.ProxyCfg.MaxNameLength.GetAsInt()))
|
|
}
|
|
|
|
firstChar := dbName[0]
|
|
if firstChar != '_' && !isAlpha(firstChar) {
|
|
return merr.WrapErrDatabaseNameInvalid(dbName,
|
|
"the first character of a database name must be an underscore or letter")
|
|
}
|
|
|
|
for i := 1; i < len(dbName); i++ {
|
|
c := dbName[i]
|
|
if c != '_' && !isAlpha(c) && !isNumber(c) {
|
|
return merr.WrapErrDatabaseNameInvalid(dbName,
|
|
"database name can only contain numbers, letters and underscores")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ValidateCollectionAlias returns true if collAlias is a valid alias name for collection, otherwise returns false.
|
|
func ValidateCollectionAlias(collAlias string) error {
|
|
return validateCollectionNameOrAlias(collAlias, "alias")
|
|
}
|
|
|
|
func validateCollectionName(collName string) error {
|
|
return validateCollectionNameOrAlias(collName, "name")
|
|
}
|
|
|
|
func validateCollectionDescription(description string) error {
|
|
if len(description) > Params.ProxyCfg.MaxCollectionDescriptionLength.GetAsInt() {
|
|
return merr.WrapErrParameterInvalidMsg(
|
|
"the length of a collection description must not exceed %s bytes",
|
|
Params.ProxyCfg.MaxCollectionDescriptionLength.GetValue())
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validatePartitionTag(partitionTag string, strictCheck bool) error {
|
|
partitionTag = strings.TrimSpace(partitionTag)
|
|
|
|
invalidMsg := "Invalid partition name: " + partitionTag + ". "
|
|
if partitionTag == "" {
|
|
msg := invalidMsg + "Partition name should not be empty."
|
|
return merr.WrapErrParameterInvalidMsg("%s", msg)
|
|
}
|
|
if len(partitionTag) > Params.ProxyCfg.MaxNameLength.GetAsInt() {
|
|
msg := invalidMsg + "The length of a partition name must be less than " + Params.ProxyCfg.MaxNameLength.GetValue() + " characters."
|
|
return merr.WrapErrParameterInvalidMsg("%s", msg)
|
|
}
|
|
|
|
if strictCheck {
|
|
firstChar := partitionTag[0]
|
|
if firstChar != '_' && !isAlpha(firstChar) && !isNumber(firstChar) {
|
|
msg := invalidMsg + "The first character of a partition name must be an underscore or letter."
|
|
return merr.WrapErrParameterInvalidMsg("%s", msg)
|
|
}
|
|
|
|
tagSize := len(partitionTag)
|
|
for i := 1; i < tagSize; i++ {
|
|
c := partitionTag[i]
|
|
if c != '_' && !isAlpha(c) && !isNumber(c) && c != '-' {
|
|
msg := invalidMsg + "Partition name can only contain numbers, letters and underscores."
|
|
return merr.WrapErrParameterInvalidMsg("%s", msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func validateFieldName(fieldName string) error {
|
|
fieldName = strings.TrimSpace(fieldName)
|
|
|
|
if fieldName == "" {
|
|
return merr.WrapErrFieldNameInvalid(fieldName, "field name should not be empty")
|
|
}
|
|
|
|
invalidMsg := "Invalid field name: " + fieldName + ". "
|
|
if len(fieldName) > Params.ProxyCfg.MaxNameLength.GetAsInt() {
|
|
msg := invalidMsg + "The length of a field name must be less than " + Params.ProxyCfg.MaxNameLength.GetValue() + " characters."
|
|
return merr.WrapErrFieldNameInvalid(fieldName, msg)
|
|
}
|
|
|
|
firstChar := fieldName[0]
|
|
if firstChar != '_' && !isAlpha(firstChar) {
|
|
msg := invalidMsg + "The first character of a field name must be an underscore or letter."
|
|
return merr.WrapErrFieldNameInvalid(fieldName, msg)
|
|
}
|
|
|
|
fieldNameSize := len(fieldName)
|
|
for i := 1; i < fieldNameSize; i++ {
|
|
c := fieldName[i]
|
|
if c != '_' && !isAlpha(c) && !isNumber(c) {
|
|
msg := invalidMsg + "Field name can only contain numbers, letters, and underscores."
|
|
return merr.WrapErrFieldNameInvalid(fieldName, msg)
|
|
}
|
|
}
|
|
if common.IsFieldNameKeyword(fieldName) {
|
|
msg := invalidMsg + fmt.Sprintf("%s is keyword in milvus.", fieldName)
|
|
return merr.WrapErrFieldNameInvalid(fieldName, msg)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateDimension(field *schemapb.FieldSchema) error {
|
|
exist := false
|
|
var dim int64
|
|
for _, param := range field.TypeParams {
|
|
if param.Key == common.DimKey {
|
|
exist = true
|
|
tmp, err := strconv.ParseInt(param.Value, 10, 64)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
dim = tmp
|
|
break
|
|
}
|
|
}
|
|
// for sparse vector field, dim should not be specified
|
|
if typeutil.IsSparseFloatVectorType(field.DataType) {
|
|
if exist {
|
|
return merr.WrapErrParameterInvalidMsg("dim should not be specified for sparse vector field %s(%d)", field.GetName(), field.FieldID)
|
|
}
|
|
return nil
|
|
}
|
|
if !exist {
|
|
return merr.WrapErrParameterInvalidMsg("dimension is not defined in field type params of field %s, check type param `dim` for vector field", field.GetName())
|
|
}
|
|
|
|
if dim <= 1 {
|
|
return merr.WrapErrParameterInvalidMsg("invalid dimension: %d. should be in range 2 ~ %d", dim, Params.ProxyCfg.MaxDimension.GetAsInt())
|
|
}
|
|
|
|
// for dense vector field, dim will be limited by max_dimension
|
|
isBinaryDimension := typeutil.IsBinaryVectorType(field.DataType) ||
|
|
(field.GetDataType() == schemapb.DataType_ArrayOfVector && typeutil.IsBinaryVectorType(field.GetElementType()))
|
|
if isBinaryDimension {
|
|
if dim%8 != 0 {
|
|
return merr.WrapErrParameterInvalidMsg("invalid dimension: %d of field %s. binary vector dimension should be multiple of 8. ", dim, field.GetName())
|
|
}
|
|
if dim > Params.ProxyCfg.MaxDimension.GetAsInt64()*8 {
|
|
return merr.WrapErrParameterInvalidMsg("invalid dimension: %d of field %s. binary vector dimension should be in range 2 ~ %d", dim, field.GetName(), Params.ProxyCfg.MaxDimension.GetAsInt()*8)
|
|
}
|
|
} else {
|
|
if dim > Params.ProxyCfg.MaxDimension.GetAsInt64() {
|
|
return merr.WrapErrParameterInvalidMsg("invalid dimension: %d of field %s. float vector dimension should be in range 2 ~ %d", dim, field.GetName(), Params.ProxyCfg.MaxDimension.GetAsInt())
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateMaxLengthPerRow(collectionName string, field *schemapb.FieldSchema) error {
|
|
exist := false
|
|
for _, param := range field.TypeParams {
|
|
if param.Key != common.MaxLengthKey {
|
|
continue
|
|
}
|
|
|
|
maxLengthPerRow, err := strconv.ParseInt(param.Value, 10, 64)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
var defaultMaxLength int64
|
|
if field.DataType == schemapb.DataType_Text {
|
|
defaultMaxLength = Params.ProxyCfg.MaxTextLength.GetAsInt64()
|
|
} else {
|
|
defaultMaxLength = Params.ProxyCfg.MaxVarCharLength.GetAsInt64()
|
|
}
|
|
|
|
if maxLengthPerRow > defaultMaxLength || maxLengthPerRow <= 0 {
|
|
return merr.WrapErrParameterInvalidMsg("the maximum length specified for the field(%s) should be in (0, %d], but got %d instead", field.GetName(), defaultMaxLength, maxLengthPerRow)
|
|
}
|
|
exist = true
|
|
}
|
|
// if not exist type params max_length, return error
|
|
if !exist {
|
|
return merr.WrapErrParameterMissingMsg("type param(max_length) should be specified for the field(%s) of collection %s", field.GetName(), collectionName)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func getMaxCapacityPerRow(collectionName string, field *schemapb.FieldSchema) (int64, error) {
|
|
maxArrayCapacity := Params.ProxyCfg.MaxArrayCapacity.GetAsInt64()
|
|
exist := false
|
|
var maxCapacityPerRow int64
|
|
for _, param := range field.TypeParams {
|
|
if param.Key != common.MaxCapacityKey {
|
|
continue
|
|
}
|
|
|
|
var err error
|
|
maxCapacityPerRow, err = strconv.ParseInt(param.Value, 10, 64)
|
|
if err != nil {
|
|
return 0, merr.WrapErrParameterInvalidMsg("the value for %s of field %s must be an integer", common.MaxCapacityKey, field.GetName())
|
|
}
|
|
if maxCapacityPerRow > maxArrayCapacity || maxCapacityPerRow <= 0 {
|
|
return 0, merr.WrapErrParameterInvalidMsg("the maximum capacity specified for a Array should be in (0, %d]", maxArrayCapacity)
|
|
}
|
|
exist = true
|
|
}
|
|
// if not exist type params max_capacity, return error
|
|
if !exist {
|
|
return 0, merr.WrapErrParameterMissingMsg("type param(max_capacity) should be specified for array field %s of collection %s", field.GetName(), collectionName)
|
|
}
|
|
return maxCapacityPerRow, nil
|
|
}
|
|
|
|
func validateMaxCapacityPerRow(collectionName string, field *schemapb.FieldSchema) error {
|
|
_, err := getMaxCapacityPerRow(collectionName, field)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateVectorFieldMetricType(field *schemapb.FieldSchema) error {
|
|
if !typeutil.IsVectorType(field.DataType) {
|
|
return nil
|
|
}
|
|
for _, params := range field.IndexParams {
|
|
if params.Key == common.MetricTypeKey {
|
|
return nil
|
|
}
|
|
}
|
|
return merr.WrapErrParameterMissingMsg(`index param "metric_type" is not specified for index float vector %s`, field.GetName())
|
|
}
|
|
|
|
func validateDuplicatedFieldName(schema *schemapb.CollectionSchema) error {
|
|
names := make(map[string]bool)
|
|
validateFieldNames := func(name string) error {
|
|
_, ok := names[name]
|
|
if ok {
|
|
return merr.WrapErrParameterInvalidMsg("duplicated field name %s found", name)
|
|
}
|
|
names[name] = true
|
|
return nil
|
|
}
|
|
for _, field := range schema.Fields {
|
|
if err := validateFieldNames(field.Name); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
for _, structArrayField := range schema.StructArrayFields {
|
|
if err := validateFieldNames(structArrayField.Name); err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, field := range structArrayField.Fields {
|
|
if err := validateFieldNames(field.Name); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateElementType(dataType schemapb.DataType) error {
|
|
switch dataType {
|
|
case schemapb.DataType_Bool, schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32,
|
|
schemapb.DataType_Int64, schemapb.DataType_Float, schemapb.DataType_Double, schemapb.DataType_VarChar:
|
|
return nil
|
|
case schemapb.DataType_String:
|
|
return merr.WrapErrParameterInvalidMsg("string data type not supported yet, please use VarChar type instead")
|
|
case schemapb.DataType_None:
|
|
return merr.WrapErrParameterInvalidMsg("element data type None is not valid")
|
|
}
|
|
return merr.WrapErrParameterInvalidMsg("element type %s is not supported", dataType.String())
|
|
}
|
|
|
|
func validateFieldType(schema *schemapb.CollectionSchema) error {
|
|
for _, field := range schema.GetFields() {
|
|
switch field.GetDataType() {
|
|
case schemapb.DataType_String:
|
|
return merr.WrapErrParameterInvalidMsg("string data type not supported yet, please use VarChar type instead")
|
|
case schemapb.DataType_None:
|
|
return merr.WrapErrParameterInvalidMsg("data type None is not valid")
|
|
case schemapb.DataType_Array:
|
|
if err := validateElementType(field.GetElementType()); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
for _, structArrayField := range schema.StructArrayFields {
|
|
for _, field := range structArrayField.Fields {
|
|
if field.GetDataType() != schemapb.DataType_Array && field.GetDataType() != schemapb.DataType_ArrayOfVector {
|
|
return merr.WrapErrParameterInvalidMsg("fields in StructArrayField must be Array or ArrayOfVector, field name = %s, field type = %s",
|
|
field.GetName(), field.GetDataType().String())
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ValidateFieldAutoID call after validatePrimaryKey
|
|
func ValidateFieldAutoID(coll *schemapb.CollectionSchema) error {
|
|
idx := -1
|
|
for i, field := range coll.Fields {
|
|
if field.AutoID {
|
|
if idx != -1 {
|
|
return merr.WrapErrParameterInvalidMsg("only one field can speficy AutoID with true, field name = %s, %s", coll.Fields[idx].Name, field.Name)
|
|
}
|
|
idx = i
|
|
if !field.IsPrimaryKey {
|
|
return merr.WrapErrParameterInvalidMsg("only primary field can speficy AutoID with true, field name = %s", field.Name)
|
|
}
|
|
}
|
|
}
|
|
for _, structArrayField := range coll.StructArrayFields {
|
|
for _, field := range structArrayField.Fields {
|
|
if field.AutoID {
|
|
return merr.WrapErrParameterInvalidMsg("autoID is not supported for struct field, field name = %s", field.Name)
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ValidateField(field *schemapb.FieldSchema, schema *schemapb.CollectionSchema) error {
|
|
// validate field name
|
|
var err error
|
|
if err := validateFieldName(field.Name); err != nil {
|
|
return err
|
|
}
|
|
// validate dense vector field type parameters
|
|
isVectorType := typeutil.IsVectorType(field.DataType)
|
|
if isVectorType {
|
|
err = validateDimension(field)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
// valid max length per row parameters
|
|
// if max_length not specified, return error
|
|
if field.DataType == schemapb.DataType_VarChar ||
|
|
(field.GetDataType() == schemapb.DataType_Array && field.GetElementType() == schemapb.DataType_VarChar) {
|
|
err = validateMaxLengthPerRow(schema.Name, field)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
// valid max capacity for array per row parameters
|
|
// if max_capacity not specified, return error
|
|
if field.DataType == schemapb.DataType_Array {
|
|
if err := validateElementType(field.GetElementType()); err != nil {
|
|
return err
|
|
}
|
|
if err = validateMaxCapacityPerRow(schema.Name, field); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if field.DataType == schemapb.DataType_ArrayOfVector {
|
|
return merr.WrapErrParameterInvalidMsg("array of vector can only be in the struct array field, field name: %s", field.Name)
|
|
}
|
|
|
|
// TODO should remove the index params in the field schema
|
|
indexParams := funcutil.KeyValuePair2Map(field.GetIndexParams())
|
|
if err = ValidateAutoIndexMmapConfig(isVectorType, indexParams); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Validate warmup policy if specified in field TypeParams
|
|
if warmupPolicy, exist := common.GetWarmupPolicy(field.GetTypeParams()...); exist {
|
|
if err = common.ValidateWarmupPolicy(warmupPolicy); err != nil {
|
|
return merr.WrapErrParameterInvalidMsg("invalid warmup policy for field %s: %s", field.Name, err.Error())
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func ValidateFieldsInStruct(field *schemapb.FieldSchema, schema *schemapb.CollectionSchema) error {
|
|
// validate field name
|
|
var err error
|
|
if err := validateFieldName(field.Name); err != nil {
|
|
return err
|
|
}
|
|
|
|
if field.DataType != schemapb.DataType_Array && field.DataType != schemapb.DataType_ArrayOfVector {
|
|
return merr.WrapErrParameterInvalidMsg("fields in StructArrayField can only be array or array of struct, but field %s is %s", field.Name, field.DataType.String())
|
|
}
|
|
|
|
if field.ElementType == schemapb.DataType_ArrayOfStruct || field.ElementType == schemapb.DataType_ArrayOfVector ||
|
|
field.ElementType == schemapb.DataType_Array {
|
|
return merr.WrapErrParameterInvalidMsg("nested array is not supported %s", field.Name)
|
|
}
|
|
|
|
if field.DataType == schemapb.DataType_Array {
|
|
if err := validateElementType(field.GetElementType()); err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
// ArrayOfVector: support FloatVector, Float16Vector, BFloat16Vector, Int8Vector, BinaryVector
|
|
if !typeutil.IsFixDimVectorType(field.GetElementType()) {
|
|
return merr.WrapErrParameterInvalidMsg("Unsupported element type %s of ArrayOfVector field %s, only fixed dimension vector types are supported", field.GetElementType().String(), field.Name)
|
|
}
|
|
|
|
err = validateDimension(field)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// valid max length per row parameters
|
|
// if max_length not specified, return error
|
|
if field.ElementType == schemapb.DataType_VarChar {
|
|
err = validateMaxLengthPerRow(schema.Name, field)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if field.DataType == schemapb.DataType_Array || field.DataType == schemapb.DataType_ArrayOfVector {
|
|
err = validateMaxCapacityPerRow(schema.Name, field)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Validate warmup policy if specified in field TypeParams
|
|
if warmupPolicy, exist := common.GetWarmupPolicy(field.GetTypeParams()...); exist {
|
|
if err = common.ValidateWarmupPolicy(warmupPolicy); err != nil {
|
|
return merr.WrapErrParameterInvalidMsg("invalid warmup policy for field %s: %s", field.Name, err.Error())
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func validateStructArrayFieldMaxCapacity(structArrayField *schemapb.StructArrayFieldSchema, collectionName string) error {
|
|
var expectedMaxCapacity int64
|
|
hasExpectedMaxCapacity := false
|
|
for _, subField := range structArrayField.Fields {
|
|
maxCapacity, err := getMaxCapacityPerRow(collectionName, subField)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !hasExpectedMaxCapacity {
|
|
expectedMaxCapacity = maxCapacity
|
|
hasExpectedMaxCapacity = true
|
|
continue
|
|
}
|
|
if maxCapacity != expectedMaxCapacity {
|
|
return merr.WrapErrParameterInvalidMsg("all sub-fields in struct array field must have the same max_capacity: structName=%s, subFieldName=%s, max_capacity=%d, expected=%d",
|
|
structArrayField.Name, subField.Name, maxCapacity, expectedMaxCapacity)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ValidateStructArrayField validates the struct array field schema.
|
|
// When the struct is nullable, sub-field schemas are mutated in-place to set Nullable=true.
|
|
func ValidateStructArrayField(structArrayField *schemapb.StructArrayFieldSchema, schema *schemapb.CollectionSchema) error {
|
|
if len(structArrayField.Fields) == 0 {
|
|
return merr.WrapErrParameterInvalidMsg("struct array field %s has no sub-fields", structArrayField.Name)
|
|
}
|
|
|
|
// Validate warmup policy if specified in struct field TypeParams
|
|
if warmupPolicy, exist := common.GetWarmupPolicy(structArrayField.GetTypeParams()...); exist {
|
|
if err := common.ValidateWarmupPolicy(warmupPolicy); err != nil {
|
|
return merr.WrapErrParameterInvalidMsg("invalid warmup policy for struct field %s: %s", structArrayField.Name, err.Error())
|
|
}
|
|
}
|
|
|
|
for _, subField := range structArrayField.Fields {
|
|
if err := ValidateFieldsInStruct(subField, schema); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if err := validateStructArrayFieldMaxCapacity(structArrayField, schema.Name); err != nil {
|
|
return err
|
|
}
|
|
|
|
// If struct is nullable, propagate nullable to all sub-fields
|
|
if structArrayField.GetNullable() {
|
|
for _, subField := range structArrayField.Fields {
|
|
subField.Nullable = true
|
|
}
|
|
} else {
|
|
// If struct is not nullable, sub-fields must not be individually nullable
|
|
for _, subField := range structArrayField.Fields {
|
|
if subField.GetNullable() {
|
|
return merr.WrapErrParameterInvalidMsg("sub-field in non-nullable struct cannot be nullable individually, set nullable on the struct instead: structName=%s, subFieldName=%s",
|
|
structArrayField.Name, subField.Name)
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func validatePrimaryKey(coll *schemapb.CollectionSchema) error {
|
|
idx := -1
|
|
for i, field := range coll.Fields {
|
|
if field.IsPrimaryKey {
|
|
if idx != -1 {
|
|
return merr.WrapErrParameterInvalidMsg("there are more than one primary key, field name = %s, %s", coll.Fields[idx].Name, field.Name)
|
|
}
|
|
|
|
// The type of the primary key field can only be int64 and varchar
|
|
if field.DataType != schemapb.DataType_Int64 && field.DataType != schemapb.DataType_VarChar {
|
|
return merr.WrapErrParameterInvalidMsg("the data type of primary key should be Int64 or VarChar")
|
|
}
|
|
|
|
// varchar field do not support autoID
|
|
// If autoID is required, it is recommended to use int64 field as the primary key
|
|
//if field.DataType == schemapb.DataType_VarChar {
|
|
// if field.AutoID {
|
|
// return merr.WrapErrParameterInvalidMsg("autoID is not supported when the VarChar field is the primary key")
|
|
// }
|
|
//}
|
|
|
|
idx = i
|
|
}
|
|
}
|
|
if idx == -1 {
|
|
// External collections may not have a primary key
|
|
if !typeutil.IsExternalCollection(coll) {
|
|
return merr.WrapErrParameterMissingMsg("primary key is not specified")
|
|
}
|
|
}
|
|
|
|
for _, structArrayField := range coll.StructArrayFields {
|
|
for _, field := range structArrayField.Fields {
|
|
if field.IsPrimaryKey {
|
|
return merr.WrapErrParameterInvalidMsg("primary key is not supported for struct field, field name = %s", field.Name)
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// validateReservedFieldNames rejects user-supplied schema fields whose name
|
|
// collides with a system-reserved identifier (RowID, Timestamp,
|
|
// __virtual_pk__). Must be called BEFORE server-side injection of the
|
|
// virtual PK so the check only applies to user input. Applies to regular
|
|
// and struct-array fields alike. Fix for issue #49314.
|
|
func validateReservedFieldNames(schema *schemapb.CollectionSchema) error {
|
|
reserved := map[string]struct{}{
|
|
common.RowIDFieldName: {},
|
|
common.TimeStampFieldName: {},
|
|
common.VirtualPKFieldName: {},
|
|
}
|
|
check := func(name string) error {
|
|
if _, ok := reserved[name]; ok {
|
|
return merr.WrapErrFieldNameInvalid(name,
|
|
fmt.Sprintf("field name %q is reserved for internal use and cannot be used in user schemas", name))
|
|
}
|
|
return nil
|
|
}
|
|
for _, f := range schema.GetFields() {
|
|
if err := check(f.GetName()); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
for _, saf := range schema.GetStructArrayFields() {
|
|
if err := check(saf.GetName()); err != nil {
|
|
return err
|
|
}
|
|
for _, f := range saf.GetFields() {
|
|
if err := check(f.GetName()); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// injectVirtualPKForExternalCollection adds a virtual PK field for external collections
|
|
// if no primary key field exists. External collections use virtual PKs in the format:
|
|
// (segmentID << 32) | offset
|
|
func injectVirtualPKForExternalCollection(schema *schemapb.CollectionSchema) error {
|
|
// Check if a primary key already exists
|
|
for _, field := range schema.Fields {
|
|
if field.IsPrimaryKey {
|
|
// PK already exists, nothing to inject
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// Create virtual PK field with FieldID=0; RootCoord's assignFieldAndFunctionID
|
|
// will assign the actual field ID during collection creation.
|
|
virtualPKField := &schemapb.FieldSchema{
|
|
Name: common.VirtualPKFieldName,
|
|
Description: "auto-generated primary key for external collection",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
AutoID: true, // Virtual PKs are auto-generated
|
|
}
|
|
|
|
// Prepend virtual PK field to the schema fields
|
|
schema.Fields = append([]*schemapb.FieldSchema{virtualPKField}, schema.Fields...)
|
|
|
|
return nil
|
|
}
|
|
|
|
func validateDynamicField(coll *schemapb.CollectionSchema) error {
|
|
for _, field := range coll.Fields {
|
|
if field.IsDynamic {
|
|
return merr.WrapErrParameterInvalidMsg("cannot explicitly set a field as a dynamic field")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// RepeatedKeyValToMap transfer the kv pairs to map.
|
|
func RepeatedKeyValToMap(kvPairs []*commonpb.KeyValuePair) (map[string]string, error) {
|
|
resMap := make(map[string]string)
|
|
for _, kv := range kvPairs {
|
|
_, ok := resMap[kv.Key]
|
|
if ok {
|
|
return nil, merr.WrapErrParameterInvalidMsg("duplicated param key: %s", kv.Key)
|
|
}
|
|
resMap[kv.Key] = kv.Value
|
|
}
|
|
return resMap, nil
|
|
}
|
|
|
|
// isVector check if dataType belongs to vector type.
|
|
func isVector(dataType schemapb.DataType) (bool, error) {
|
|
switch dataType {
|
|
case schemapb.DataType_Bool, schemapb.DataType_Int8,
|
|
schemapb.DataType_Int16, schemapb.DataType_Int32,
|
|
schemapb.DataType_Int64,
|
|
schemapb.DataType_Float, schemapb.DataType_Double:
|
|
return false, nil
|
|
|
|
case schemapb.DataType_FloatVector, schemapb.DataType_BinaryVector, schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector, schemapb.DataType_SparseFloatVector:
|
|
return true, nil
|
|
}
|
|
|
|
return false, merr.WrapErrParameterInvalidMsg("invalid data type: %d", dataType)
|
|
}
|
|
|
|
func validateMetricType(dataType schemapb.DataType, metricTypeStrRaw string) error {
|
|
metricTypeStr := strings.ToUpper(metricTypeStrRaw)
|
|
switch metricTypeStr {
|
|
case metric.L2, metric.IP, metric.COSINE:
|
|
if typeutil.IsFloatVectorType(dataType) {
|
|
return nil
|
|
}
|
|
case metric.JACCARD, metric.HAMMING, metric.SUBSTRUCTURE, metric.SUPERSTRUCTURE, metric.MHJACCARD:
|
|
if dataType == schemapb.DataType_BinaryVector {
|
|
return nil
|
|
}
|
|
}
|
|
return merr.WrapErrParameterInvalidMsg("data_type %s mismatch with metric_type %s", dataType.String(), metricTypeStrRaw)
|
|
}
|
|
|
|
// validateMultipleVectorFields check if schema has multiple vector fields.
|
|
func validateMultipleVectorFields(schema *schemapb.CollectionSchema) error {
|
|
vecExist := false
|
|
var vecName string
|
|
|
|
for i := range schema.Fields {
|
|
name := schema.Fields[i].Name
|
|
dType := schema.Fields[i].DataType
|
|
isVec := typeutil.IsVectorType(dType)
|
|
if isVec && vecExist && !enableMultipleVectorFields {
|
|
return merr.WrapErrParameterInvalidMsg(
|
|
"multiple vector fields is not supported, fields name: %s, %s",
|
|
vecName,
|
|
name,
|
|
)
|
|
} else if isVec {
|
|
vecExist = true
|
|
vecName = name
|
|
}
|
|
}
|
|
|
|
// todo(Spadea): should be there any check between vectors in struct fields?
|
|
|
|
return nil
|
|
}
|
|
|
|
func validateLoadFieldsList(schema *schemapb.CollectionSchema) error {
|
|
var vectorCnt int
|
|
for _, field := range schema.Fields {
|
|
shouldLoad, err := common.ShouldFieldBeLoaded(field.GetTypeParams())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// shoud load field, skip other check
|
|
if shouldLoad {
|
|
if typeutil.IsVectorType(field.GetDataType()) {
|
|
vectorCnt++
|
|
}
|
|
continue
|
|
}
|
|
|
|
if field.IsPrimaryKey {
|
|
return merr.WrapErrParameterInvalidMsg("Primary key field %s cannot skip loading", field.GetName())
|
|
}
|
|
|
|
if field.IsPartitionKey {
|
|
return merr.WrapErrParameterInvalidMsg("Partition Key field %s cannot skip loading", field.GetName())
|
|
}
|
|
|
|
if field.IsClusteringKey {
|
|
return merr.WrapErrParameterInvalidMsg("Clustering Key field %s cannot skip loading", field.GetName())
|
|
}
|
|
}
|
|
|
|
for _, structArrayField := range schema.StructArrayFields {
|
|
for _, field := range structArrayField.Fields {
|
|
shouldLoad, err := common.ShouldFieldBeLoaded(field.GetTypeParams())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if shouldLoad {
|
|
if typeutil.IsVectorType(field.ElementType) {
|
|
vectorCnt++
|
|
}
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
|
|
if vectorCnt == 0 {
|
|
return merr.WrapErrParameterInvalidMsg("cannot config all vector field(s) skip loading")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// parsePrimaryFieldData2IDs get IDs to fill grpc result, for example insert request, delete request etc.
|
|
func parsePrimaryFieldData2IDs(fieldData *schemapb.FieldData) (*schemapb.IDs, error) {
|
|
primaryData := &schemapb.IDs{}
|
|
switch fieldData.Field.(type) {
|
|
case *schemapb.FieldData_Scalars:
|
|
scalarField := fieldData.GetScalars()
|
|
switch scalarField.Data.(type) {
|
|
case *schemapb.ScalarField_LongData:
|
|
primaryData.IdField = &schemapb.IDs_IntId{
|
|
IntId: scalarField.GetLongData(),
|
|
}
|
|
case *schemapb.ScalarField_StringData:
|
|
primaryData.IdField = &schemapb.IDs_StrId{
|
|
StrId: scalarField.GetStringData(),
|
|
}
|
|
default:
|
|
return nil, merr.WrapErrParameterInvalidMsg("currently only support DataType Int64 or VarChar as PrimaryField")
|
|
}
|
|
default:
|
|
return nil, merr.WrapErrParameterInvalidMsg("currently not support vector field as PrimaryField")
|
|
}
|
|
|
|
return primaryData, nil
|
|
}
|
|
|
|
// hasDuplicates checks if there are any duplicate values in the slice.
|
|
// Returns true immediately when the first duplicate is found (early exit).
|
|
func hasDuplicates[T comparable](ids []T) bool {
|
|
seen := make(map[T]struct{}, len(ids))
|
|
for _, id := range ids {
|
|
if _, exists := seen[id]; exists {
|
|
return true
|
|
}
|
|
seen[id] = struct{}{}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// CheckDuplicatePkExist checks if there are duplicate primary keys in the field data.
|
|
// Returns (true, nil) if duplicates exist, (false, nil) if no duplicates.
|
|
// Returns (false, error) if there's an error during checking.
|
|
func CheckDuplicatePkExist(primaryFieldSchema *schemapb.FieldSchema, fieldsData []*schemapb.FieldData) (bool, error) {
|
|
if len(fieldsData) == 0 {
|
|
return false, nil
|
|
}
|
|
|
|
// find primary field data
|
|
var primaryFieldData *schemapb.FieldData
|
|
for _, field := range fieldsData {
|
|
if field.GetFieldName() == primaryFieldSchema.GetName() {
|
|
primaryFieldData = field
|
|
break
|
|
}
|
|
}
|
|
|
|
if primaryFieldData == nil {
|
|
return false, merr.WrapErrParameterInvalidMsg("must assign pk when upsert, primary field: %v", primaryFieldSchema.GetName())
|
|
}
|
|
|
|
// check for duplicates based on primary key type
|
|
switch primaryFieldData.Field.(type) {
|
|
case *schemapb.FieldData_Scalars:
|
|
scalarField := primaryFieldData.GetScalars()
|
|
switch scalarField.Data.(type) {
|
|
case *schemapb.ScalarField_LongData:
|
|
intIDs := scalarField.GetLongData().GetData()
|
|
return hasDuplicates(intIDs), nil
|
|
case *schemapb.ScalarField_StringData:
|
|
strIDs := scalarField.GetStringData().GetData()
|
|
return hasDuplicates(strIDs), nil
|
|
default:
|
|
return false, merr.WrapErrParameterInvalidMsg("unsupported primary key type")
|
|
}
|
|
default:
|
|
return false, merr.WrapErrParameterInvalidMsg("primary field must be scalar type")
|
|
}
|
|
}
|
|
|
|
// autoGenPrimaryFieldData generate primary data when autoID == true
|
|
func autoGenPrimaryFieldData(fieldSchema *schemapb.FieldSchema, data interface{}) (*schemapb.FieldData, error) {
|
|
var fieldData schemapb.FieldData
|
|
fieldData.FieldName = fieldSchema.Name
|
|
fieldData.Type = fieldSchema.DataType
|
|
switch data := data.(type) {
|
|
case []int64:
|
|
switch fieldData.Type {
|
|
case schemapb.DataType_Int64:
|
|
fieldData.Field = &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_LongData{
|
|
LongData: &schemapb.LongArray{
|
|
Data: data,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
case schemapb.DataType_VarChar:
|
|
strIDs := make([]string, len(data))
|
|
for i, v := range data {
|
|
strIDs[i] = strconv.FormatInt(v, 10)
|
|
}
|
|
fieldData.Field = &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_StringData{
|
|
StringData: &schemapb.StringArray{
|
|
Data: strIDs,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
default:
|
|
return nil, merr.WrapErrParameterInvalidMsg("currently only support autoID for int64 and varchar PrimaryField")
|
|
}
|
|
default:
|
|
return nil, merr.WrapErrParameterInvalidMsg("currently only int64 is supported as the data source for the autoID of a PrimaryField")
|
|
}
|
|
|
|
return &fieldData, nil
|
|
}
|
|
|
|
func autoGenDynamicFieldData(schema *schemapb.CollectionSchema, data [][]byte) *schemapb.FieldData {
|
|
fd := &schemapb.FieldData{
|
|
FieldName: common.MetaFieldName,
|
|
Type: schemapb.DataType_JSON,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_JsonData{
|
|
JsonData: &schemapb.JSONArray{
|
|
Data: data,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
IsDynamic: true,
|
|
}
|
|
|
|
// Only set ValidData when the $meta field is nullable or has a default value.
|
|
// For 2.5 collections (non-nullable, no default), CheckValidData expects
|
|
// len(ValidData)==0, so we must NOT set it.
|
|
for _, f := range schema.Fields {
|
|
if f.GetIsDynamic() && (f.GetNullable() || f.GetDefaultValue() != nil) {
|
|
validData := make([]bool, len(data))
|
|
for i := range validData {
|
|
validData[i] = true
|
|
}
|
|
fd.ValidData = validData
|
|
break
|
|
}
|
|
}
|
|
|
|
return fd
|
|
}
|
|
|
|
// validateFieldDataColumns validates that all required fields are present and no unknown fields exist.
|
|
// It checks:
|
|
// 1. The number of columns matches the expected count (excluding BM25 output fields)
|
|
// 2. All field names exist in the schema
|
|
// Returns detailed error message listing expected and provided fields if validation fails.
|
|
func validateFieldDataColumns(columns []*schemapb.FieldData, schema *schemaInfo) error {
|
|
expectColumnNum := 0
|
|
|
|
// Count expected columns
|
|
for _, field := range schema.GetFields() {
|
|
if !typeutil.IsBM25FunctionOutputField(field, schema.CollectionSchema) && !typeutil.IsMinHashFunctionOutputField(field, schema.CollectionSchema) {
|
|
expectColumnNum++
|
|
}
|
|
}
|
|
for _, structField := range schema.GetStructArrayFields() {
|
|
expectColumnNum += len(structField.GetFields())
|
|
}
|
|
|
|
// Validate column count
|
|
if len(columns) != expectColumnNum {
|
|
return merr.WrapErrParameterInvalidMsg("len(columns) mismatch the expectColumnNum, expectColumnNum: %d, len(columns): %d",
|
|
expectColumnNum, len(columns))
|
|
}
|
|
|
|
// Validate field existence using schemaHelper
|
|
for _, fieldData := range columns {
|
|
_, err := schema.schemaHelper.GetFieldFromNameDefaultJSON(fieldData.FieldName)
|
|
if err != nil {
|
|
return merr.WrapErrParameterInvalidMsg("fieldName %v not exist in collection schema", fieldData.FieldName)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// fillFieldPropertiesOnly fills field properties (FieldId, Type, ElementType) from schema.
|
|
// It assumes that columns have been validated and does not perform validation.
|
|
// Use validateFieldDataColumns before calling this function if validation is needed.
|
|
func fillFieldPropertiesOnly(columns []*schemapb.FieldData, schema *schemaInfo) error {
|
|
for _, fieldData := range columns {
|
|
// Use schemaHelper to get field schema, automatically handles dynamic fields
|
|
fieldSchema, err := schema.schemaHelper.GetFieldFromNameDefaultJSON(fieldData.FieldName)
|
|
if err != nil {
|
|
return merr.WrapErrParameterInvalidMsg("fieldName %v not exist in collection schema", fieldData.FieldName)
|
|
}
|
|
|
|
fieldData.FieldId = fieldSchema.FieldID
|
|
fieldData.Type = fieldSchema.DataType
|
|
|
|
// Set the ElementType because it may not be set in the insert request.
|
|
switch fieldData.Type {
|
|
case schemapb.DataType_Array:
|
|
fd, ok := fieldData.Field.(*schemapb.FieldData_Scalars)
|
|
if !ok || fd.Scalars.GetArrayData() == nil {
|
|
return merr.WrapErrParameterInvalidMsg("field convert FieldData_Scalars fail in fieldData, fieldName: %s, collectionName: %s",
|
|
fieldData.FieldName, schema.Name)
|
|
}
|
|
fd.Scalars.GetArrayData().ElementType = fieldSchema.ElementType
|
|
case schemapb.DataType_ArrayOfVector:
|
|
fd, ok := fieldData.Field.(*schemapb.FieldData_Vectors)
|
|
if !ok || fd.Vectors.GetVectorArray() == nil {
|
|
return merr.WrapErrParameterInvalidMsg("field convert FieldData_Vectors fail in fieldData, fieldName: %s, collectionName: %s",
|
|
fieldData.FieldName, schema.Name)
|
|
}
|
|
fd.Vectors.GetVectorArray().ElementType = fieldSchema.ElementType
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func ValidateUsername(username string) error {
|
|
username = strings.TrimSpace(username)
|
|
|
|
if username == "" {
|
|
return merr.WrapErrParameterInvalidMsg("username must be not empty")
|
|
}
|
|
|
|
if len(username) > Params.ProxyCfg.MaxUsernameLength.GetAsInt() {
|
|
return merr.WrapErrParameterInvalidMsg("invalid username %s with length %d, the length of username must be less than %d", username, len(username), Params.ProxyCfg.MaxUsernameLength.GetAsInt())
|
|
}
|
|
|
|
firstChar := username[0]
|
|
if !isAlpha(firstChar) {
|
|
return merr.WrapErrParameterInvalidMsg("invalid user name %s, the first character must be a letter, but got %s", username, string(firstChar))
|
|
}
|
|
|
|
usernameSize := len(username)
|
|
for i := 1; i < usernameSize; i++ {
|
|
c := username[i]
|
|
if c != '_' && c != '-' && c != '.' && !isAlpha(c) && !isNumber(c) {
|
|
return merr.WrapErrParameterInvalidMsg("invalid user name %s, username must contain only numbers, letters, underscores, dots, and hyphens, but got %s", username, c)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ValidateUserDescription(description string) error {
|
|
if len(description) > Params.ProxyCfg.MaxUserDescriptionLength.GetAsInt() {
|
|
return merr.WrapErrParameterInvalidMsg("invalid user description with length %d, the length of user description must be less than or equal to %s", len(description), Params.ProxyCfg.MaxUserDescriptionLength.GetValue())
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ValidatePassword(password string) error {
|
|
if len(password) < Params.ProxyCfg.MinPasswordLength.GetAsInt() || len(password) > Params.ProxyCfg.MaxPasswordLength.GetAsInt() {
|
|
return merr.WrapErrParameterInvalidRange(Params.ProxyCfg.MinPasswordLength.GetAsInt(),
|
|
Params.ProxyCfg.MaxPasswordLength.GetAsInt(),
|
|
len(password), "invalid password length")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ReplaceID2Name(oldStr string, id int64, name string) string {
|
|
return strings.ReplaceAll(oldStr, strconv.FormatInt(id, 10), name)
|
|
}
|
|
|
|
func parseGuaranteeTsFromConsistency(ts, tMax typeutil.Timestamp, consistency commonpb.ConsistencyLevel) typeutil.Timestamp {
|
|
switch consistency {
|
|
case commonpb.ConsistencyLevel_Strong:
|
|
ts = tMax
|
|
case commonpb.ConsistencyLevel_Bounded:
|
|
ratio := Params.CommonCfg.GracefulTime.GetAsDuration(time.Millisecond)
|
|
ts = tsoutil.AddPhysicalDurationOnTs(tMax, -ratio)
|
|
case commonpb.ConsistencyLevel_Eventually:
|
|
ts = 1
|
|
}
|
|
return ts
|
|
}
|
|
|
|
func parseGuaranteeTs(ts, tMax typeutil.Timestamp) typeutil.Timestamp {
|
|
switch ts {
|
|
case strongTS:
|
|
ts = tMax
|
|
case boundedTS:
|
|
ratio := Params.CommonCfg.GracefulTime.GetAsDuration(time.Millisecond)
|
|
ts = tsoutil.AddPhysicalDurationOnTs(tMax, -ratio)
|
|
}
|
|
return ts
|
|
}
|
|
|
|
func getMaxMvccTsFromChannels(channelsTs map[string]uint64, beginTs typeutil.Timestamp) typeutil.Timestamp {
|
|
maxTs := typeutil.Timestamp(0)
|
|
for _, ts := range channelsTs {
|
|
if ts > maxTs {
|
|
maxTs = ts
|
|
}
|
|
}
|
|
|
|
if maxTs == 0 {
|
|
mlog.Warn(context.TODO(), "no channel ts found, use beginTs instead")
|
|
return beginTs
|
|
}
|
|
|
|
return maxTs
|
|
}
|
|
|
|
func validateName(entity string, nameType string) error {
|
|
return validateNameWithCustomChars(entity, nameType, Params.ProxyCfg.NameValidationAllowedChars.GetValue())
|
|
}
|
|
|
|
func validateNameWithCustomChars(entity string, nameType string, allowedChars string) error {
|
|
entity = strings.TrimSpace(entity)
|
|
|
|
if entity == "" {
|
|
return merr.WrapErrParameterInvalid("not empty", entity, nameType+" should be not empty")
|
|
}
|
|
|
|
if len(entity) > Params.ProxyCfg.MaxNameLength.GetAsInt() {
|
|
return merr.WrapErrParameterInvalidRange(0,
|
|
Params.ProxyCfg.MaxNameLength.GetAsInt(),
|
|
len(entity),
|
|
fmt.Sprintf("the length of %s must be not greater than limit", nameType))
|
|
}
|
|
|
|
firstChar := entity[0]
|
|
if firstChar != '_' && !isAlpha(firstChar) {
|
|
return merr.WrapErrParameterInvalid('_',
|
|
firstChar,
|
|
fmt.Sprintf("the first character of %s must be an underscore or letter", nameType))
|
|
}
|
|
|
|
for i := 1; i < len(entity); i++ {
|
|
c := entity[i]
|
|
if c != '_' && !isAlpha(c) && !isNumber(c) && !strings.ContainsRune(allowedChars, rune(c)) {
|
|
return merr.WrapErrParameterInvalidMsg("%s can only contain numbers, letters, underscores, and allowed characters (%s), found %c at %d", nameType, allowedChars, c, i)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ValidateRoleName(entity string) error {
|
|
return validateNameWithCustomChars(entity, "role name", Params.ProxyCfg.RoleNameValidationAllowedChars.GetValue())
|
|
}
|
|
|
|
func ValidateRoleDescription(description string) error {
|
|
return rbacutil.ValidateRoleDescription(description, Params.ProxyCfg.MaxRoleDescriptionLength.GetAsInt())
|
|
}
|
|
|
|
func IsDefaultRole(roleName string) bool {
|
|
for _, defaultRole := range util.DefaultRoles {
|
|
if defaultRole == roleName {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func ValidateObjectName(entity string) error {
|
|
if util.IsAnyWord(entity) {
|
|
return nil
|
|
}
|
|
return validateName(entity, "object name")
|
|
}
|
|
|
|
func ValidateCollectionName(entity string) error {
|
|
if util.IsAnyWord(entity) {
|
|
return nil
|
|
}
|
|
return validateName(entity, "collection name")
|
|
}
|
|
|
|
// ValidateSnapshotName validates snapshot name using standard naming rules.
|
|
func ValidateSnapshotName(snapshotName string) error {
|
|
return validateName(snapshotName, "snapshot name")
|
|
}
|
|
|
|
func ValidateObjectType(entity string) error {
|
|
return validateName(entity, "ObjectType")
|
|
}
|
|
|
|
func ValidatePrivilege(entity string) error {
|
|
if util.IsAnyWord(entity) {
|
|
return nil
|
|
}
|
|
return validateName(entity, "Privilege")
|
|
}
|
|
|
|
func GetCurUserFromContext(ctx context.Context) (string, error) {
|
|
return contextutil.GetCurUserFromContext(ctx)
|
|
}
|
|
|
|
func GetCurUserFromContextOrDefault(ctx context.Context) string {
|
|
username, _ := GetCurUserFromContext(ctx)
|
|
return username
|
|
}
|
|
|
|
func GetCurDBNameFromContextOrDefault(ctx context.Context) string {
|
|
md, ok := metadata.FromIncomingContext(ctx)
|
|
if !ok {
|
|
return util.DefaultDBName
|
|
}
|
|
dbNameData := md[strings.ToLower(util.HeaderDBName)]
|
|
if len(dbNameData) < 1 || dbNameData[0] == "" {
|
|
return util.DefaultDBName
|
|
}
|
|
return dbNameData[0]
|
|
}
|
|
|
|
// GetCurDBNameFromRequestOrContext returns the database a request actually
|
|
// operates on. It prefers the DbName carried in the request body (which is
|
|
// what downstream handlers execute against, after DatabaseInterceptor has
|
|
// normalized it) and only falls back to the connection-context db / cluster
|
|
// default when the request carries none.
|
|
//
|
|
// Privilege checks MUST use this rather than GetCurDBNameFromContextOrDefault:
|
|
// authorizing against the connection-context db while the operation runs
|
|
// against the request's DbName both falsely denies legitimate access and
|
|
// allows cross-database privilege escalation (see milvus-io/milvus#50678).
|
|
func GetCurDBNameFromRequestOrContext(ctx context.Context, req interface{}) string {
|
|
if getter, ok := req.(requestutil.DBNameGetter); ok {
|
|
if dbName := getter.GetDbName(); dbName != "" {
|
|
return dbName
|
|
}
|
|
}
|
|
return GetCurDBNameFromContextOrDefault(ctx)
|
|
}
|
|
|
|
func NewContextWithMetadata(ctx context.Context, username string, dbName string) context.Context {
|
|
dbKey := strings.ToLower(util.HeaderDBName)
|
|
if dbName != "" {
|
|
ctx = contextutil.AppendToIncomingContext(ctx, dbKey, dbName)
|
|
}
|
|
if username != "" {
|
|
originValue := fmt.Sprintf("%s%s%s", username, util.CredentialSeparator, username)
|
|
authKey := strings.ToLower(util.HeaderAuthorize)
|
|
authValue := crypto.Base64Encode(originValue)
|
|
ctx = contextutil.AppendToIncomingContext(ctx, authKey, authValue)
|
|
}
|
|
return ctx
|
|
}
|
|
|
|
func AppendUserInfoForRPC(ctx context.Context) context.Context {
|
|
curUser, _ := GetCurUserFromContext(ctx)
|
|
if curUser != "" {
|
|
originValue := fmt.Sprintf("%s%s%s", curUser, util.CredentialSeparator, curUser)
|
|
authKey := strings.ToLower(util.HeaderAuthorize)
|
|
authValue := crypto.Base64Encode(originValue)
|
|
ctx = metadata.AppendToOutgoingContext(ctx, authKey, authValue)
|
|
}
|
|
return ctx
|
|
}
|
|
|
|
func GetRole(username string) ([]string, error) {
|
|
privCache := privilege.GetPrivilegeCache()
|
|
if privCache == nil {
|
|
return []string{}, merr.WrapErrServiceUnavailable("internal: Milvus Proxy is not ready yet. please wait")
|
|
}
|
|
return privCache.GetUserRole(username), nil
|
|
}
|
|
|
|
func PasswordVerify(ctx context.Context, username, rawPwd string) bool {
|
|
return passwordVerify(ctx, username, rawPwd, privilege.GetPrivilegeCache())
|
|
}
|
|
|
|
func VerifyAPIKey(rawToken string) (string, error) {
|
|
hoo := hookutil.GetHook()
|
|
user, err := hoo.VerifyAPIKey(rawToken)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "fail to verify apikey", mlog.String("api_key", rawToken), mlog.Err(err))
|
|
return "", merr.WrapErrParameterInvalidMsg("invalid apikey: [%s]", rawToken)
|
|
}
|
|
return user, nil
|
|
}
|
|
|
|
// PasswordVerify verify password
|
|
func passwordVerify(ctx context.Context, username, rawPwd string, privilegeCache privilege.PrivilegeCache) bool {
|
|
// it represents the cache miss if Sha256Password is empty within credInfo, which shall be updated first connection.
|
|
// meanwhile, generating Sha256Password depends on raw password and encrypted password will not cache.
|
|
credInfo, err := privilege.GetPrivilegeCache().GetCredentialInfo(ctx, username)
|
|
if err != nil {
|
|
mlog.Error(context.TODO(), "found no credential", mlog.String("username", username), mlog.Err(err))
|
|
return false
|
|
}
|
|
|
|
// hit cache
|
|
sha256Pwd := crypto.SHA256(rawPwd, credInfo.Username)
|
|
if credInfo.Sha256Password != "" {
|
|
return sha256Pwd == credInfo.Sha256Password
|
|
}
|
|
|
|
// miss cache, verify against encrypted password from etcd
|
|
if err := bcrypt.CompareHashAndPassword([]byte(credInfo.EncryptedPassword), []byte(rawPwd)); err != nil {
|
|
mlog.Error(context.TODO(), "Verify password failed", mlog.Err(err))
|
|
return false
|
|
}
|
|
|
|
// update cache after miss cache
|
|
credInfo.Sha256Password = sha256Pwd
|
|
mlog.Debug(context.TODO(), "get credential miss cache, update cache with", mlog.Any("credential", credInfo))
|
|
privilegeCache.UpdateCredential(credInfo)
|
|
return true
|
|
}
|
|
|
|
func translatePkOutputFields(schema *schemapb.CollectionSchema) ([]string, []int64) {
|
|
pkNames := []string{}
|
|
fieldIDs := []int64{}
|
|
for _, field := range schema.Fields {
|
|
if field.IsPrimaryKey {
|
|
pkNames = append(pkNames, field.GetName())
|
|
fieldIDs = append(fieldIDs, field.GetFieldID())
|
|
}
|
|
}
|
|
return pkNames, fieldIDs
|
|
}
|
|
|
|
func recallCal[T string | int64](results []T, gts []T) float32 {
|
|
hit := 0
|
|
total := 0
|
|
for _, r := range results {
|
|
total++
|
|
for _, gt := range gts {
|
|
if r == gt {
|
|
hit++
|
|
break
|
|
}
|
|
}
|
|
}
|
|
return float32(hit) / float32(total)
|
|
}
|
|
|
|
func computeRecall(results *schemapb.SearchResultData, gts *schemapb.SearchResultData) error {
|
|
if results.GetNumQueries() != gts.GetNumQueries() {
|
|
return merr.WrapErrParameterInvalidMsg("num of queries is inconsistent between search results(%d) and ground truth(%d)", results.GetNumQueries(), gts.GetNumQueries())
|
|
}
|
|
|
|
// When search returns no results, IDs field is nil. Set recalls to 0 for all queries.
|
|
if results.GetIds() == nil || results.GetIds().GetIdField() == nil ||
|
|
gts.GetIds() == nil || gts.GetIds().GetIdField() == nil {
|
|
results.Recalls = make([]float32, results.GetNumQueries())
|
|
return nil
|
|
}
|
|
|
|
switch results.GetIds().GetIdField().(type) {
|
|
case *schemapb.IDs_IntId:
|
|
switch gts.GetIds().GetIdField().(type) {
|
|
case *schemapb.IDs_IntId:
|
|
currentResultIndex := int64(0)
|
|
currentGTIndex := int64(0)
|
|
recalls := make([]float32, 0, results.GetNumQueries())
|
|
for i := 0; i < int(results.GetNumQueries()); i++ {
|
|
currentResultTopk := results.GetTopks()[i]
|
|
currentGTTopk := gts.GetTopks()[i]
|
|
recalls = append(recalls, recallCal(results.GetIds().GetIntId().GetData()[currentResultIndex:currentResultIndex+currentResultTopk],
|
|
gts.GetIds().GetIntId().GetData()[currentGTIndex:currentGTIndex+currentGTTopk]))
|
|
currentResultIndex += currentResultTopk
|
|
currentGTIndex += currentGTTopk
|
|
}
|
|
results.Recalls = recalls
|
|
return nil
|
|
case *schemapb.IDs_StrId:
|
|
return merr.WrapErrParameterInvalidMsg("pk type is inconsistent between search results(int64) and ground truth(string)")
|
|
default:
|
|
return merr.WrapErrParameterInvalidMsg("unsupported pk type")
|
|
}
|
|
|
|
case *schemapb.IDs_StrId:
|
|
switch gts.GetIds().GetIdField().(type) {
|
|
case *schemapb.IDs_StrId:
|
|
currentResultIndex := int64(0)
|
|
currentGTIndex := int64(0)
|
|
recalls := make([]float32, 0, results.GetNumQueries())
|
|
for i := 0; i < int(results.GetNumQueries()); i++ {
|
|
currentResultTopk := results.GetTopks()[i]
|
|
currentGTTopk := gts.GetTopks()[i]
|
|
recalls = append(recalls, recallCal(results.GetIds().GetStrId().GetData()[currentResultIndex:currentResultIndex+currentResultTopk],
|
|
gts.GetIds().GetStrId().GetData()[currentGTIndex:currentGTIndex+currentGTTopk]))
|
|
currentResultIndex += currentResultTopk
|
|
currentGTIndex += currentGTTopk
|
|
}
|
|
results.Recalls = recalls
|
|
return nil
|
|
case *schemapb.IDs_IntId:
|
|
return merr.WrapErrParameterInvalidMsg("pk type is inconsistent between search results(string) and ground truth(int64)")
|
|
default:
|
|
return merr.WrapErrParameterInvalidMsg("unsupported pk type")
|
|
}
|
|
default:
|
|
return merr.WrapErrParameterInvalidMsg("unsupported pk type")
|
|
}
|
|
}
|
|
|
|
// Support wildcard in output fields:
|
|
//
|
|
// "*" - all fields
|
|
//
|
|
// For example, A and B are scalar fields, C and D are vector fields, duplicated fields will automatically be removed.
|
|
//
|
|
// output_fields=["*"] ==> [A,B,C,D]
|
|
// output_fields=["*",A] ==> [A,B,C,D]
|
|
// output_fields=["*",C] ==> [A,B,C,D]
|
|
//
|
|
// 4th return value is true if user requested pk field explicitly or using wildcard.
|
|
// if removePkField is true, pk field will not be include in the first(resultFieldNames)/second(userOutputFields)
|
|
// return value.
|
|
func translateOutputFields(outputFields []string, schema *schemaInfo, removePkField bool) ([]string, []string, []string, []agg.AggregateBase, bool, error) {
|
|
var primaryFieldName string
|
|
allFieldNameMap := make(map[string]*schemapb.FieldSchema)
|
|
resultFieldNameMap := make(map[string]bool)
|
|
resultFieldNames := make([]string, 0)
|
|
userOutputFieldsMap := make(map[string]bool)
|
|
userOutputFields := make([]string, 0)
|
|
userDynamicFieldsMap := make(map[string]bool)
|
|
userDynamicFields := make([]string, 0)
|
|
useAllDyncamicFields := false
|
|
aggregates := make([]agg.AggregateBase, 0)
|
|
for _, field := range schema.Fields {
|
|
if field.IsPrimaryKey {
|
|
primaryFieldName = field.Name
|
|
}
|
|
allFieldNameMap[field.Name] = field
|
|
}
|
|
|
|
// User may specify a struct array field or some specific fields in the struct array field
|
|
for _, subStruct := range schema.StructArrayFields {
|
|
for _, field := range subStruct.Fields {
|
|
allFieldNameMap[field.Name] = field
|
|
}
|
|
}
|
|
|
|
structArrayNameToFields := make(map[string][]*schemapb.FieldSchema)
|
|
for _, subStruct := range schema.StructArrayFields {
|
|
structArrayNameToFields[subStruct.Name] = subStruct.Fields
|
|
}
|
|
|
|
userRequestedPkFieldExplicitly := false
|
|
|
|
for _, outputFieldName := range outputFields {
|
|
outputFieldName = strings.TrimSpace(outputFieldName)
|
|
if outputFieldName == primaryFieldName {
|
|
userRequestedPkFieldExplicitly = true
|
|
}
|
|
if outputFieldName == "*" {
|
|
userRequestedPkFieldExplicitly = true
|
|
for fieldName, field := range allFieldNameMap {
|
|
if schema.CanRetrieveRawFieldData(field) {
|
|
resultFieldNameMap[fieldName] = true
|
|
userOutputFieldsMap[fieldName] = true
|
|
}
|
|
}
|
|
useAllDyncamicFields = true
|
|
} else {
|
|
if isAgg, aggregateName, aggFieldName := agg.MatchAggregationExpression(outputFieldName); isAgg {
|
|
if aggField, ok := allFieldNameMap[aggFieldName]; ok {
|
|
aggFuncs, aggErr := agg.NewAggregate(aggregateName, aggField.GetFieldID(), outputFieldName, aggField.GetDataType())
|
|
if aggErr != nil {
|
|
return nil, nil, nil, nil, false, aggErr
|
|
}
|
|
aggregates = append(aggregates, aggFuncs...)
|
|
} else if aggFieldName == "*" {
|
|
// only count(*) is allowed
|
|
if aggregateName != "count" {
|
|
return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("%s(*) is not supported, only count(*) is allowed", aggregateName)
|
|
}
|
|
if err := agg.ValidateAggFieldType(aggregateName, schemapb.DataType_None); err != nil {
|
|
return nil, nil, nil, nil, false, err
|
|
}
|
|
aggFuncs, aggErr := agg.NewAggregate(aggregateName, 0, outputFieldName, schemapb.DataType_None)
|
|
if aggErr != nil {
|
|
return nil, nil, nil, nil, false, aggErr
|
|
}
|
|
aggregates = append(aggregates, aggFuncs...)
|
|
} else {
|
|
return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("target field %s for aggregation:%s does not exist", aggFieldName, aggregateName)
|
|
}
|
|
userOutputFieldsMap[outputFieldName] = true
|
|
continue
|
|
}
|
|
|
|
if structArrayField, ok := structArrayNameToFields[outputFieldName]; ok {
|
|
for _, field := range structArrayField {
|
|
if schema.CanRetrieveRawFieldData(field) {
|
|
resultFieldNameMap[field.Name] = true
|
|
userOutputFieldsMap[field.Name] = true
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
if field, ok := allFieldNameMap[outputFieldName]; ok {
|
|
if !schema.CanRetrieveRawFieldData(field) {
|
|
return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("not allowed to retrieve raw data of field %s", outputFieldName)
|
|
}
|
|
resultFieldNameMap[outputFieldName] = true
|
|
userOutputFieldsMap[outputFieldName] = true
|
|
} else {
|
|
if schema.EnableDynamicField {
|
|
dynamicNestedPath := outputFieldName
|
|
err := planparserv2.ParseIdentifier(schema.schemaHelper, outputFieldName, func(expr *planpb.Expr) error {
|
|
columnInfo := expr.GetColumnExpr().GetInfo()
|
|
// there must be no error here
|
|
dynamicField, _ := schema.schemaHelper.GetDynamicField()
|
|
// only $meta["xxx"] is allowed for now
|
|
if dynamicField.GetFieldID() != columnInfo.GetFieldId() {
|
|
return merr.WrapErrParameterInvalidMsg("not support getting subkeys of json field yet")
|
|
}
|
|
nestedPaths := columnInfo.GetNestedPath()
|
|
// $meta["A"]["B"] not allowed for now
|
|
if len(nestedPaths) != 1 {
|
|
return merr.WrapErrParameterInvalidMsg("not support getting multiple level of dynamic field for now")
|
|
}
|
|
// $meta["dyn_field"], output field name could be:
|
|
// 1. "dyn_field", outputFieldName == nestedPath
|
|
// 2. `$meta["dyn_field"]` explicit form
|
|
if nestedPaths[0] != outputFieldName {
|
|
// use "dyn_field" as userDynamicFieldsMap when outputField = `$meta["dyn_field"]`
|
|
dynamicNestedPath = nestedPaths[0]
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
mlog.Info(context.TODO(), "parse output field name failed", mlog.String("field name", outputFieldName), mlog.Err(err))
|
|
return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("parse output field name failed: %s", outputFieldName)
|
|
}
|
|
resultFieldNameMap[common.MetaFieldName] = true
|
|
userOutputFieldsMap[outputFieldName] = true
|
|
userDynamicFieldsMap[dynamicNestedPath] = true
|
|
} else {
|
|
return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("field %s not exist", outputFieldName)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if removePkField {
|
|
delete(resultFieldNameMap, primaryFieldName)
|
|
delete(userOutputFieldsMap, primaryFieldName)
|
|
}
|
|
|
|
for fieldName := range resultFieldNameMap {
|
|
resultFieldNames = append(resultFieldNames, fieldName)
|
|
}
|
|
for fieldName := range userOutputFieldsMap {
|
|
userOutputFields = append(userOutputFields, fieldName)
|
|
}
|
|
if !useAllDyncamicFields {
|
|
for fieldName := range userDynamicFieldsMap {
|
|
userDynamicFields = append(userDynamicFields, fieldName)
|
|
}
|
|
}
|
|
|
|
return resultFieldNames, userOutputFields, userDynamicFields, aggregates, userRequestedPkFieldExplicitly, nil
|
|
}
|
|
|
|
func validCharInIndexName(c byte) bool {
|
|
return c == '_' || c == '[' || c == ']' || isAlpha(c) || isNumber(c)
|
|
}
|
|
|
|
func validateIndexName(indexName string) error {
|
|
// Shared with rootcoord's bound-index prepare (indexparamcheck).
|
|
return indexparamcheck.ValidateIndexName(indexName)
|
|
}
|
|
|
|
func isCollectionLoaded(ctx context.Context, mc types.MixCoordClient, collID int64) (bool, error) {
|
|
// get all loading collections
|
|
resp, err := mc.ShowLoadCollections(ctx, &querypb.ShowCollectionsRequest{
|
|
CollectionIDs: nil,
|
|
})
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if resp.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
|
|
return false, merr.Error(resp.GetStatus())
|
|
}
|
|
|
|
for _, loadedCollID := range resp.GetCollectionIDs() {
|
|
if collID == loadedCollID {
|
|
return true, nil
|
|
}
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
func isPartitionLoaded(ctx context.Context, mc types.MixCoordClient, collID int64, partID int64) (bool, error) {
|
|
// get all loading collections
|
|
resp, err := mc.ShowLoadPartitions(ctx, &querypb.ShowPartitionsRequest{
|
|
CollectionID: collID,
|
|
PartitionIDs: []int64{partID},
|
|
})
|
|
if err := merr.CheckRPCCall(resp, err); err != nil {
|
|
// qc returns error if partition not loaded
|
|
if errors.Is(err, merr.ErrPartitionNotLoaded) {
|
|
return false, nil
|
|
}
|
|
return false, err
|
|
}
|
|
|
|
return true, nil
|
|
}
|
|
|
|
func checkFieldsDataBySchema(ctx context.Context, allFields []*schemapb.FieldSchema, schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg, inInsert bool) error {
|
|
log := mlog.With(mlog.String("collection", schema.GetName()))
|
|
primaryKeyNum := 0
|
|
autoGenFieldNum := 0
|
|
|
|
dataNameSet := typeutil.NewSet[string]()
|
|
for _, data := range insertMsg.FieldsData {
|
|
fieldName := data.GetFieldName()
|
|
if dataNameSet.Contain(fieldName) {
|
|
return merr.WrapErrParameterInvalidMsg("duplicated field %s found", fieldName)
|
|
}
|
|
dataNameSet.Insert(fieldName)
|
|
}
|
|
|
|
allowInsertAutoID, _ := common.IsAllowInsertAutoID(schema.GetProperties()...)
|
|
hasPkData := false
|
|
needAutoGenPk := false
|
|
|
|
for _, fieldSchema := range allFields {
|
|
if fieldSchema.AutoID && !fieldSchema.IsPrimaryKey {
|
|
log.Warn(ctx, "not primary key field, but set autoID true", mlog.String("field", fieldSchema.GetName()))
|
|
return merr.WrapErrParameterInvalidMsg("only primary key could be with AutoID enabled")
|
|
}
|
|
|
|
if fieldSchema.IsPrimaryKey {
|
|
primaryKeyNum++
|
|
hasPkData = dataNameSet.Contain(fieldSchema.GetName())
|
|
needAutoGenPk = fieldSchema.AutoID && (!allowInsertAutoID || !hasPkData)
|
|
}
|
|
if fieldSchema.GetDefaultValue() != nil && fieldSchema.IsPrimaryKey {
|
|
return merr.WrapErrParameterInvalidMsg("primary key can't be with default value")
|
|
}
|
|
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && needAutoGenPk && inInsert) || typeutil.IsBM25FunctionOutputField(fieldSchema, schema) || typeutil.IsMinHashFunctionOutputField(fieldSchema, schema) {
|
|
// when inInsert, no need to pass when pk is autoid and SkipAutoIDCheck is false
|
|
autoGenFieldNum++
|
|
}
|
|
if _, ok := dataNameSet[fieldSchema.GetName()]; !ok {
|
|
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && needAutoGenPk && inInsert) || typeutil.IsBM25FunctionOutputField(fieldSchema, schema) || typeutil.IsMinHashFunctionOutputField(fieldSchema, schema) {
|
|
// autoGenField
|
|
continue
|
|
}
|
|
|
|
if fieldSchema.GetDefaultValue() == nil && !fieldSchema.GetNullable() {
|
|
log.Warn(ctx, "no corresponding fieldData pass in", mlog.String("fieldSchema", fieldSchema.GetName()))
|
|
return merr.WrapErrParameterInvalidMsg("fieldSchema(%s) has no corresponding fieldData pass in", fieldSchema.GetName())
|
|
}
|
|
// when use default_value or has set Nullable
|
|
// it's ok that no corresponding fieldData found
|
|
dataToAppend, err := typeutil.GenEmptyFieldData(fieldSchema)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
dataToAppend.ValidData = make([]bool, insertMsg.GetNumRows())
|
|
insertMsg.FieldsData = append(insertMsg.FieldsData, dataToAppend)
|
|
}
|
|
}
|
|
|
|
if primaryKeyNum > 1 {
|
|
log.Warn(ctx, "more than 1 primary keys not supported",
|
|
mlog.Int64("primaryKeyNum", int64(primaryKeyNum)))
|
|
return merr.WrapErrParameterInvalidMsg("more than 1 primary keys not supported, got %d", primaryKeyNum)
|
|
}
|
|
expectedNum := len(allFields)
|
|
actualNum := len(insertMsg.FieldsData) + autoGenFieldNum
|
|
|
|
if expectedNum != actualNum {
|
|
log.Warn(ctx, "the number of fields is not the same as needed", mlog.Int("expected", expectedNum), mlog.Int("actual", actualNum))
|
|
return merr.WrapErrParameterInvalid(expectedNum, actualNum, "more fieldData has pass in")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// subFieldHasData checks whether a struct sub-field contains actual data content,
|
|
// not just a protobuf-initialized empty wrapper (e.g. some SDKs may initialize the
|
|
// Vectors oneof by accessing .vectors.dim, making Field non-nil without real data).
|
|
func subFieldHasData(subField *schemapb.FieldData) bool {
|
|
switch fd := subField.Field.(type) {
|
|
case *schemapb.FieldData_Scalars:
|
|
return fd.Scalars.GetData() != nil
|
|
case *schemapb.FieldData_Vectors:
|
|
return fd.Vectors.GetData() != nil
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
// checkAndFlattenStructFieldData verifies the array length of the struct array field data in the insert message
|
|
// and then flattens the data so that data node and query node have not to handle the struct array field data.
|
|
func checkAndFlattenStructFieldData(schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg) error {
|
|
structSchemaMap := make(map[string]*schemapb.StructArrayFieldSchema, len(schema.GetStructArrayFields()))
|
|
for _, structField := range schema.GetStructArrayFields() {
|
|
structSchemaMap[structField.Name] = structField
|
|
}
|
|
|
|
fieldSchemaMap := make(map[string]*schemapb.FieldSchema, len(schema.GetFields()))
|
|
for _, fieldSchema := range schema.GetFields() {
|
|
fieldSchemaMap[fieldSchema.Name] = fieldSchema
|
|
}
|
|
|
|
structFieldCount := 0
|
|
flattenedFields := make([]*schemapb.FieldData, 0, len(insertMsg.GetFieldsData())+5)
|
|
|
|
for _, fieldData := range insertMsg.GetFieldsData() {
|
|
if _, ok := fieldSchemaMap[fieldData.FieldName]; ok {
|
|
flattenedFields = append(flattenedFields, fieldData)
|
|
continue
|
|
}
|
|
|
|
structName := fieldData.FieldName
|
|
structSchema, ok := structSchemaMap[structName]
|
|
if !ok {
|
|
return merr.WrapErrParameterInvalidMsg("fieldName %v not exist in collection schema, fieldType %v, fieldId %v", fieldData.FieldName, fieldData.Type, fieldData.FieldId)
|
|
}
|
|
|
|
structFieldCount++
|
|
structArrays, ok := fieldData.Field.(*schemapb.FieldData_StructArrays)
|
|
if !ok {
|
|
return merr.WrapErrParameterInvalidMsg("field convert FieldData_StructArrays fail in fieldData, fieldName: %s,"+
|
|
" collectionName:%s", structName, schema.Name)
|
|
}
|
|
|
|
if len(structArrays.StructArrays.Fields) != len(structSchema.GetFields()) {
|
|
return merr.WrapErrParameterInvalidMsg("length of fields of struct field mismatch length of the fields in schema, fieldName: %s,"+
|
|
" collectionName:%s, fieldData fields length:%d, schema fields length:%d",
|
|
structName, schema.Name, len(structArrays.StructArrays.Fields), len(structSchema.GetFields()))
|
|
}
|
|
|
|
// Check sub-field data consistency: within the same struct, all sub-fields must
|
|
// either have data or all be empty. Partial presence is invalid.
|
|
hasDataCount := 0
|
|
for _, subField := range structArrays.StructArrays.Fields {
|
|
if subFieldHasData(subField) {
|
|
hasDataCount++
|
|
}
|
|
}
|
|
totalSubFields := len(structArrays.StructArrays.Fields)
|
|
if hasDataCount == 0 {
|
|
// All sub-fields have empty payload — equivalent to the struct being
|
|
// omitted entirely. Reject illegal ValidData first: when no payload is
|
|
// provided, any ValidData[i]==true contradicts itself.
|
|
for _, subField := range structArrays.StructArrays.Fields {
|
|
for j, v := range subField.ValidData {
|
|
if v {
|
|
return merr.WrapErrParameterInvalidMsg("sub-field '%s' in struct '%s' claims row %d is valid but no payload is provided",
|
|
subField.FieldName, structName, j)
|
|
}
|
|
}
|
|
}
|
|
// Skip flatten and let checkFieldsDataBySchema backfill missing sub-fields
|
|
// uniformly, so scenario "struct omitted" and scenario "struct present but
|
|
// empty" share one code path downstream.
|
|
continue
|
|
}
|
|
if hasDataCount != totalSubFields {
|
|
return merr.WrapErrParameterInvalidMsg("inconsistent sub-field data in struct '%s': %d of %d sub-fields have data, all must be present or all absent",
|
|
structName, hasDataCount, totalSubFields)
|
|
}
|
|
|
|
// Validate that all sub-fields share the same ValidData mask.
|
|
// Nullable is a struct-level concept: a row is either entirely null or entirely present.
|
|
if structSchema.GetNullable() {
|
|
var refValidData []bool
|
|
var refFieldName string
|
|
refInitialized := false
|
|
for _, subField := range structArrays.StructArrays.Fields {
|
|
if !refInitialized {
|
|
refValidData = subField.ValidData
|
|
refFieldName = subField.FieldName
|
|
refInitialized = true
|
|
continue
|
|
}
|
|
if len(subField.ValidData) != len(refValidData) {
|
|
return merr.WrapErrParameterInvalidMsg("sub-field ValidData length mismatch in struct '%s': '%s' has %d, '%s' has %d",
|
|
structName, refFieldName, len(refValidData), subField.FieldName, len(subField.ValidData))
|
|
}
|
|
for j := range refValidData {
|
|
if subField.ValidData[j] != refValidData[j] {
|
|
return merr.WrapErrParameterInvalidMsg("sub-field ValidData mismatch in struct '%s' at row %d: '%s'=%v, '%s'=%v",
|
|
structName, j, refFieldName, refValidData[j], subField.FieldName, subField.ValidData[j])
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
subFieldSchemaByName := make(map[string]*schemapb.FieldSchema, len(structSchema.GetFields())*2)
|
|
for _, subFieldSchema := range structSchema.GetFields() {
|
|
subFieldSchemaByName[subFieldSchema.GetName()] = subFieldSchema
|
|
subFieldSchemaByName[storedStructSubFieldName(structName, subFieldSchema.GetName())] = subFieldSchema
|
|
if typeutil.IsStructSubField(subFieldSchema.GetName()) {
|
|
rawName, err := typeutil.ExtractStructFieldName(subFieldSchema.GetName())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
subFieldSchemaByName[rawName] = subFieldSchema
|
|
}
|
|
}
|
|
|
|
vectorElementWidth := func(subField *schemapb.FieldData, subFieldSchema *schemapb.FieldSchema) (int, error) {
|
|
dim, err := typeutil.GetDim(subFieldSchema)
|
|
if err != nil {
|
|
return 0, merr.WrapErrParameterInvalidErr(err, "sub-field '%s' in struct '%s'", subField.GetFieldName(), structName)
|
|
}
|
|
if dim <= 0 {
|
|
return 0, merr.WrapErrParameterInvalidMsg("sub-field '%s' in struct '%s': invalid dim %d", subField.GetFieldName(), structName, dim)
|
|
}
|
|
switch subFieldSchema.GetElementType() {
|
|
case schemapb.DataType_FloatVector, schemapb.DataType_Int8Vector:
|
|
return int(dim), nil
|
|
case schemapb.DataType_BinaryVector:
|
|
return int((dim + 7) / 8), nil
|
|
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
|
|
return int(dim * 2), nil
|
|
default:
|
|
return 0, merr.WrapErrParameterInvalidMsg("sub-field '%s' in struct '%s': unsupported array-of-vector element type %s",
|
|
subField.GetFieldName(), structName, subFieldSchema.GetElementType().String())
|
|
}
|
|
}
|
|
|
|
// Check the payload row count and, while those rows are in hand, verify the
|
|
// per-row struct element count. The outer row count only proves that every
|
|
// sub-field has the same number of physical payload rows. For each such row,
|
|
// every sub-field must also describe the same number of struct elements.
|
|
expectedArrayLen := -1
|
|
var firstValidData []bool
|
|
type rowElementCounter struct {
|
|
name string
|
|
count func(physicalRow int) (int, error)
|
|
}
|
|
rowElementCounters := make([]rowElementCounter, 0, totalSubFields)
|
|
for _, subField := range structArrays.StructArrays.Fields {
|
|
subFieldSchema := subFieldSchemaByName[subField.GetFieldName()]
|
|
if subFieldSchema == nil {
|
|
return merr.WrapErrParameterInvalidMsg("sub-field '%s' not found in struct schema '%s'", subField.GetFieldName(), structName)
|
|
}
|
|
|
|
var currentArrayLen int
|
|
|
|
switch subFieldData := subField.Field.(type) {
|
|
case *schemapb.FieldData_Scalars:
|
|
if scalarArray := subFieldData.Scalars.GetArrayData(); scalarArray != nil {
|
|
currentArrayLen = len(scalarArray.Data)
|
|
if totalSubFields > 1 {
|
|
rowElementCounters = append(rowElementCounters, rowElementCounter{
|
|
name: subField.GetFieldName(),
|
|
count: func(physicalRow int) (int, error) {
|
|
row := scalarArray.GetData()[physicalRow]
|
|
if row.GetData() == nil {
|
|
return 0, merr.WrapErrParameterInvalidMsg("nil array data")
|
|
}
|
|
switch subFieldSchema.GetElementType() {
|
|
case schemapb.DataType_Bool:
|
|
return len(row.GetBoolData().GetData()), nil
|
|
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
|
|
return len(row.GetIntData().GetData()), nil
|
|
case schemapb.DataType_Int64:
|
|
return len(row.GetLongData().GetData()), nil
|
|
case schemapb.DataType_Float:
|
|
return len(row.GetFloatData().GetData()), nil
|
|
case schemapb.DataType_Double:
|
|
return len(row.GetDoubleData().GetData()), nil
|
|
case schemapb.DataType_VarChar, schemapb.DataType_String:
|
|
return len(row.GetStringData().GetData()), nil
|
|
default:
|
|
return 0, merr.WrapErrParameterInvalidMsg("unsupported array element type %s", subFieldSchema.GetElementType().String())
|
|
}
|
|
},
|
|
})
|
|
}
|
|
} else {
|
|
return merr.WrapErrParameterInvalidMsg("scalar array data is nil in struct field '%s', sub-field '%s'",
|
|
structName, subField.FieldName)
|
|
}
|
|
case *schemapb.FieldData_Vectors:
|
|
if vectorArray := subFieldData.Vectors.GetVectorArray(); vectorArray != nil {
|
|
currentArrayLen = len(vectorArray.Data)
|
|
if totalSubFields > 1 {
|
|
var vectorWidth int
|
|
rowElementCounters = append(rowElementCounters, rowElementCounter{
|
|
name: subField.GetFieldName(),
|
|
count: func(physicalRow int) (int, error) {
|
|
if vectorWidth == 0 {
|
|
var err error
|
|
vectorWidth, err = vectorElementWidth(subField, subFieldSchema)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
row := vectorArray.GetData()[physicalRow]
|
|
if row.GetData() == nil {
|
|
return 0, merr.WrapErrParameterInvalidMsg("nil vector array data")
|
|
}
|
|
var payloadLen int
|
|
switch subFieldSchema.GetElementType() {
|
|
case schemapb.DataType_FloatVector:
|
|
payloadLen = len(row.GetFloatVector().GetData())
|
|
case schemapb.DataType_BinaryVector:
|
|
payloadLen = len(row.GetBinaryVector())
|
|
case schemapb.DataType_Float16Vector:
|
|
payloadLen = len(row.GetFloat16Vector())
|
|
case schemapb.DataType_BFloat16Vector:
|
|
payloadLen = len(row.GetBfloat16Vector())
|
|
case schemapb.DataType_Int8Vector:
|
|
payloadLen = len(row.GetInt8Vector())
|
|
}
|
|
if payloadLen%vectorWidth != 0 {
|
|
return 0, merr.WrapErrParameterInvalidMsg("payload length %d is not divisible by vector width %d", payloadLen, vectorWidth)
|
|
}
|
|
return payloadLen / vectorWidth, nil
|
|
},
|
|
})
|
|
}
|
|
} else {
|
|
return merr.WrapErrParameterInvalidMsg("vector array data is nil in struct field '%s', sub-field '%s'",
|
|
structName, subField.FieldName)
|
|
}
|
|
default:
|
|
return merr.WrapErrParameterInvalidMsg("unexpected field data type in struct array field, fieldName: %s", structName)
|
|
}
|
|
|
|
if expectedArrayLen == -1 {
|
|
expectedArrayLen = currentArrayLen
|
|
firstValidData = subField.GetValidData()
|
|
} else if currentArrayLen != expectedArrayLen {
|
|
return merr.WrapErrParameterInvalidMsg("inconsistent array length in struct field '%s': expected %d, got %d for sub-field '%s'",
|
|
structName, expectedArrayLen, currentArrayLen, subField.FieldName)
|
|
}
|
|
}
|
|
|
|
if totalSubFields > 1 && expectedArrayLen > 0 {
|
|
var physicalToLogical []int
|
|
if len(firstValidData) > 0 {
|
|
physicalToLogical = make([]int, 0, expectedArrayLen)
|
|
for logicalRow, valid := range firstValidData {
|
|
if valid {
|
|
physicalToLogical = append(physicalToLogical, logicalRow)
|
|
}
|
|
}
|
|
if len(physicalToLogical) != expectedArrayLen {
|
|
return merr.WrapErrParameterInvalidMsg("invalid ValidData for struct '%s': true count %d does not match payload row count %d",
|
|
structName, len(physicalToLogical), expectedArrayLen)
|
|
}
|
|
}
|
|
logicalRow := func(physicalRow int) int {
|
|
if len(physicalToLogical) == 0 {
|
|
return physicalRow
|
|
}
|
|
return physicalToLogical[physicalRow]
|
|
}
|
|
|
|
refCounter := rowElementCounters[0]
|
|
refElementCounts := make([]int, expectedArrayLen)
|
|
for physicalRow := 0; physicalRow < expectedArrayLen; physicalRow++ {
|
|
count, err := refCounter.count(physicalRow)
|
|
if err != nil {
|
|
return merr.WrapErrParameterInvalidErr(err, "struct '%s' row %d sub-field '%s'",
|
|
structName, logicalRow(physicalRow), refCounter.name)
|
|
}
|
|
refElementCounts[physicalRow] = count
|
|
}
|
|
for _, counter := range rowElementCounters[1:] {
|
|
for physicalRow := 0; physicalRow < expectedArrayLen; physicalRow++ {
|
|
count, err := counter.count(physicalRow)
|
|
if err != nil {
|
|
return merr.WrapErrParameterInvalidErr(err, "struct '%s' row %d sub-field '%s'",
|
|
structName, logicalRow(physicalRow), counter.name)
|
|
}
|
|
if count != refElementCounts[physicalRow] {
|
|
return merr.WrapErrParameterInvalidMsg("inconsistent struct element count in struct '%s' at row %d: '%s' has %d, '%s' has %d",
|
|
structName, logicalRow(physicalRow), refCounter.name, refElementCounts[physicalRow], counter.name, count)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for _, subField := range structArrays.StructArrays.Fields {
|
|
transformedFieldName := storedStructSubFieldName(structName, subField.FieldName)
|
|
subFieldCopy := &schemapb.FieldData{
|
|
FieldName: transformedFieldName,
|
|
FieldId: subField.FieldId,
|
|
Type: subField.Type,
|
|
Field: subField.Field,
|
|
IsDynamic: subField.IsDynamic,
|
|
ValidData: subField.ValidData,
|
|
}
|
|
|
|
flattenedFields = append(flattenedFields, subFieldCopy)
|
|
}
|
|
}
|
|
|
|
// Verify all required (non-nullable) struct array fields are provided
|
|
seenStructs := make(map[string]bool, structFieldCount)
|
|
for _, fieldData := range insertMsg.GetFieldsData() {
|
|
if _, ok := structSchemaMap[fieldData.FieldName]; ok {
|
|
seenStructs[fieldData.FieldName] = true
|
|
}
|
|
}
|
|
for _, sf := range schema.GetStructArrayFields() {
|
|
if !sf.GetNullable() && !seenStructs[sf.Name] {
|
|
return merr.WrapErrParameterInvalidMsg("required struct array field '%s' is missing in insert data", sf.Name)
|
|
}
|
|
}
|
|
|
|
insertMsg.FieldsData = flattenedFields
|
|
return nil
|
|
}
|
|
|
|
func checkPrimaryFieldData(ctx context.Context, allFields []*schemapb.FieldSchema, schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg) (*schemapb.IDs, error) {
|
|
log := mlog.With(mlog.String("collectionName", insertMsg.CollectionName))
|
|
rowNums := uint32(insertMsg.NRows())
|
|
// TODO(dragondriver): in fact, NumRows is not trustable, we should check all input fields
|
|
if insertMsg.NRows() <= 0 {
|
|
return nil, merr.WrapErrParameterInvalid("invalid num_rows", fmt.Sprint(rowNums), "num_rows should be greater than 0")
|
|
}
|
|
|
|
if err := checkFieldsDataBySchema(ctx, allFields, schema, insertMsg, true); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
primaryFieldSchema, err := typeutil.GetPrimaryFieldSchema(schema)
|
|
if err != nil {
|
|
log.Error(ctx, "get primary field schema failed", mlog.Any("schema", schema), mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
if primaryFieldSchema.GetNullable() {
|
|
return nil, merr.WrapErrParameterInvalidMsg("primary field not support null")
|
|
}
|
|
var primaryFieldData *schemapb.FieldData
|
|
// when checkPrimaryFieldData in insert
|
|
|
|
allowInsertAutoID, _ := common.IsAllowInsertAutoID(schema.GetProperties()...)
|
|
skipAutoIDCheck := primaryFieldSchema.AutoID &&
|
|
typeutil.IsPrimaryFieldDataExist(insertMsg.GetFieldsData(), primaryFieldSchema) && (Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() || allowInsertAutoID)
|
|
|
|
if !primaryFieldSchema.AutoID || skipAutoIDCheck {
|
|
primaryFieldData, err = typeutil.GetPrimaryFieldData(insertMsg.GetFieldsData(), primaryFieldSchema)
|
|
if err != nil {
|
|
log.Info(ctx, "get primary field data failed", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
} else {
|
|
// check primary key data not exist
|
|
if typeutil.IsPrimaryFieldDataExist(insertMsg.GetFieldsData(), primaryFieldSchema) {
|
|
return nil, merr.WrapErrParameterInvalidMsg("can not assign primary field data when auto id enabled and allow_insert_auto_id is false %v", primaryFieldSchema.Name)
|
|
}
|
|
// if autoID == true, currently support autoID for int64 and varchar PrimaryField
|
|
primaryFieldData, err = autoGenPrimaryFieldData(primaryFieldSchema, insertMsg.GetRowIDs())
|
|
if err != nil {
|
|
log.Info(ctx, "generate primary field data failed when autoID == true", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
// if autoID == true, set the primary field data
|
|
// insertMsg.fieldsData need append primaryFieldData
|
|
insertMsg.FieldsData = append(insertMsg.FieldsData, primaryFieldData)
|
|
}
|
|
|
|
// parse primaryFieldData to result.IDs, and as returned primary keys
|
|
ids, err := parsePrimaryFieldData2IDs(primaryFieldData)
|
|
if err != nil {
|
|
log.Warn(ctx, "parse primary field data to IDs failed", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
return ids, nil
|
|
}
|
|
|
|
// check whether insertMsg has all fields in schema
|
|
func LackOfFieldsDataBySchema(schema *schemapb.CollectionSchema, fieldsData []*schemapb.FieldData, skipPkFieldCheck bool, skipDynamicFieldCheck bool) error {
|
|
log := mlog.With(mlog.String("collection", schema.GetName()))
|
|
|
|
// find bm25 generated fields
|
|
bm25Fields := typeutil.NewSet[string](GetFunctionOutputFields(schema)...)
|
|
dataNameMap := make(map[string]*schemapb.FieldData)
|
|
for _, data := range fieldsData {
|
|
dataNameMap[data.GetFieldName()] = data
|
|
}
|
|
|
|
for _, fieldSchema := range schema.Fields {
|
|
if bm25Fields.Contain(fieldSchema.GetName()) {
|
|
continue
|
|
}
|
|
|
|
if fieldSchema.GetNullable() || fieldSchema.GetDefaultValue() != nil {
|
|
continue
|
|
}
|
|
|
|
if _, ok := dataNameMap[fieldSchema.GetName()]; !ok {
|
|
if (fieldSchema.IsPrimaryKey && fieldSchema.AutoID && !Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && skipPkFieldCheck) ||
|
|
typeutil.IsBM25FunctionOutputField(fieldSchema, schema) || typeutil.IsMinHashFunctionOutputField(fieldSchema, schema) ||
|
|
(skipDynamicFieldCheck && fieldSchema.GetIsDynamic()) {
|
|
// autoGenField
|
|
continue
|
|
}
|
|
|
|
log.Info(context.TODO(), "no corresponding fieldData pass in", mlog.String("fieldSchema", fieldSchema.GetName()))
|
|
return merr.WrapErrParameterInvalidMsg("fieldSchema(%s) has no corresponding fieldData pass in", fieldSchema.GetName())
|
|
}
|
|
}
|
|
for _, structSchema := range schema.GetStructArrayFields() {
|
|
if structSchema.GetNullable() {
|
|
continue
|
|
}
|
|
if _, ok := dataNameMap[structSchema.GetName()]; !ok {
|
|
log.Info(context.TODO(), "no corresponding struct fieldData pass in", mlog.String("structFieldSchema", structSchema.GetName()))
|
|
return merr.WrapErrParameterInvalidMsg("structFieldSchema(%s) has no corresponding fieldData pass in", structSchema.GetName())
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// for some varchar with analzyer
|
|
// we need check char format before insert it to message queue
|
|
// now only support utf-8
|
|
func checkInputUtf8Compatiable(allFields []*schemapb.FieldSchema, insertMsg *msgstream.InsertMsg) error {
|
|
checkeFields := lo.FilterMap(allFields, func(field *schemapb.FieldSchema, _ int) (int64, bool) {
|
|
if field.DataType == schemapb.DataType_VarChar {
|
|
return field.GetFieldID(), true
|
|
}
|
|
|
|
if field.DataType != schemapb.DataType_Text {
|
|
return 0, false
|
|
}
|
|
|
|
for _, kv := range field.GetTypeParams() {
|
|
if kv.Key == common.EnableAnalyzerKey {
|
|
return field.GetFieldID(), true
|
|
}
|
|
}
|
|
return 0, false
|
|
})
|
|
|
|
if len(checkeFields) == 0 {
|
|
return nil
|
|
}
|
|
|
|
for _, fieldData := range insertMsg.FieldsData {
|
|
if !lo.Contains(checkeFields, fieldData.GetFieldId()) {
|
|
continue
|
|
}
|
|
|
|
strData := fieldData.GetScalars().GetStringData()
|
|
for row, data := range strData.GetData() {
|
|
ok := utf8.ValidString(data)
|
|
if !ok {
|
|
mlog.Warn(context.TODO(), "string field data not utf-8 format", mlog.String("messageVersion", strData.ProtoReflect().Descriptor().Syntax().GoString()))
|
|
return merr.WrapErrAsInputError(merr.WrapErrParameterInvalidMsg("input with analyzer should be utf-8 format, but row: %d not utf-8 format. data: %s", row, data))
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func checkUpsertPrimaryFieldData(ctx context.Context, allFields []*schemapb.FieldSchema, schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg) (*schemapb.IDs, *schemapb.IDs, error) {
|
|
log := mlog.With(mlog.String("collectionName", insertMsg.CollectionName))
|
|
rowNums := uint32(insertMsg.NRows())
|
|
// TODO(dragondriver): in fact, NumRows is not trustable, we should check all input fields
|
|
if insertMsg.NRows() <= 0 {
|
|
return nil, nil, merr.WrapErrParameterInvalid("invalid num_rows", fmt.Sprint(rowNums), "num_rows should be greater than 0")
|
|
}
|
|
|
|
if err := checkFieldsDataBySchema(ctx, allFields, schema, insertMsg, false); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
primaryFieldSchema, err := typeutil.GetPrimaryFieldSchema(schema)
|
|
if err != nil {
|
|
log.Error(ctx, "get primary field schema failed", mlog.Any("schema", schema), mlog.Err(err))
|
|
return nil, nil, err
|
|
}
|
|
if primaryFieldSchema.GetNullable() {
|
|
return nil, nil, merr.WrapErrParameterInvalidMsg("primary field not support null")
|
|
}
|
|
// get primaryFieldData whether autoID is true or not
|
|
var primaryFieldData *schemapb.FieldData
|
|
var newPrimaryFieldData *schemapb.FieldData
|
|
|
|
primaryFieldID := primaryFieldSchema.FieldID
|
|
primaryFieldName := primaryFieldSchema.Name
|
|
for i, field := range insertMsg.GetFieldsData() {
|
|
if field.FieldId == primaryFieldID || field.FieldName == primaryFieldName {
|
|
primaryFieldData = field
|
|
if primaryFieldSchema.AutoID {
|
|
// use the passed pk as new pk when autoID == false
|
|
// automatic generate pk as new pk wehen autoID == true
|
|
newPrimaryFieldData, err = autoGenPrimaryFieldData(primaryFieldSchema, insertMsg.GetRowIDs())
|
|
if err != nil {
|
|
log.Info(ctx, "generate new primary field data failed when upsert", mlog.Err(err))
|
|
return nil, nil, err
|
|
}
|
|
insertMsg.FieldsData = append(insertMsg.GetFieldsData()[:i], insertMsg.GetFieldsData()[i+1:]...)
|
|
insertMsg.FieldsData = append(insertMsg.FieldsData, newPrimaryFieldData)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
// must assign primary field data when upsert
|
|
if primaryFieldData == nil {
|
|
return nil, nil, merr.WrapErrParameterInvalidMsg("must assign pk when upsert, primary field: %v", primaryFieldName)
|
|
}
|
|
|
|
// parse primaryFieldData to result.IDs, and as returned primary keys
|
|
ids, err := parsePrimaryFieldData2IDs(primaryFieldData)
|
|
if err != nil {
|
|
log.Warn(ctx, "parse primary field data to IDs failed", mlog.Err(err))
|
|
return nil, nil, err
|
|
}
|
|
if !primaryFieldSchema.GetAutoID() {
|
|
return ids, ids, nil
|
|
}
|
|
newIDs, err := parsePrimaryFieldData2IDs(newPrimaryFieldData)
|
|
if err != nil {
|
|
log.Warn(ctx, "parse primary field data to IDs failed", mlog.Err(err))
|
|
return nil, nil, err
|
|
}
|
|
return newIDs, ids, nil
|
|
}
|
|
|
|
func getPartitionKeyFieldData(fieldSchema *schemapb.FieldSchema, insertMsg *msgstream.InsertMsg) (*schemapb.FieldData, error) {
|
|
if len(insertMsg.GetPartitionName()) > 0 && !Params.ProxyCfg.SkipPartitionKeyCheck.GetAsBool() {
|
|
return nil, merr.WrapErrParameterInvalidMsg("not support manually specifying the partition names if partition key mode is used")
|
|
}
|
|
|
|
for _, fieldData := range insertMsg.GetFieldsData() {
|
|
if fieldData.GetFieldId() == fieldSchema.GetFieldID() {
|
|
return fieldData, nil
|
|
}
|
|
}
|
|
|
|
return nil, merr.WrapErrParameterInvalidMsg("partition key not specify when insert")
|
|
}
|
|
|
|
func getCollectionProgress(
|
|
ctx context.Context,
|
|
queryCoord types.QueryCoordClient,
|
|
msgBase *commonpb.MsgBase,
|
|
collectionID int64,
|
|
) (loadProgress int64, refreshProgress int64, err error) {
|
|
resp, err := queryCoord.ShowLoadCollections(ctx, &querypb.ShowCollectionsRequest{
|
|
Base: commonpbutil.UpdateMsgBase(
|
|
msgBase,
|
|
commonpbutil.WithMsgType(commonpb.MsgType_ShowCollections),
|
|
),
|
|
CollectionIDs: []int64{collectionID},
|
|
})
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "fail to show collections",
|
|
mlog.Int64("collectionID", collectionID),
|
|
mlog.Err(err),
|
|
)
|
|
return
|
|
}
|
|
|
|
err = merr.Error(resp.GetStatus())
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "fail to show collections",
|
|
mlog.Int64("collectionID", collectionID),
|
|
mlog.Err(err))
|
|
return
|
|
}
|
|
|
|
loadProgress = resp.GetInMemoryPercentages()[0]
|
|
if len(resp.GetRefreshProgress()) > 0 { // Compatibility for new Proxy with old QueryCoord
|
|
refreshProgress = resp.GetRefreshProgress()[0]
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func getPartitionProgress(
|
|
ctx context.Context,
|
|
queryCoord types.QueryCoordClient,
|
|
msgBase *commonpb.MsgBase,
|
|
partitionNames []string,
|
|
collectionName string,
|
|
collectionID int64,
|
|
dbName string,
|
|
) (loadProgress int64, refreshProgress int64, err error) {
|
|
IDs2Names := make(map[int64]string)
|
|
partitionIDs := make([]int64, 0)
|
|
for _, partitionName := range partitionNames {
|
|
var partitionID int64
|
|
partitionID, err = globalMetaCache.GetPartitionID(ctx, dbName, collectionName, partitionName)
|
|
if err != nil {
|
|
return
|
|
}
|
|
IDs2Names[partitionID] = partitionName
|
|
partitionIDs = append(partitionIDs, partitionID)
|
|
}
|
|
|
|
var resp *querypb.ShowPartitionsResponse
|
|
resp, err = queryCoord.ShowLoadPartitions(ctx, &querypb.ShowPartitionsRequest{
|
|
Base: commonpbutil.UpdateMsgBase(
|
|
msgBase,
|
|
commonpbutil.WithMsgType(commonpb.MsgType_ShowPartitions),
|
|
),
|
|
CollectionID: collectionID,
|
|
PartitionIDs: partitionIDs,
|
|
})
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "fail to show partitions", mlog.Int64("collection_id", collectionID),
|
|
mlog.String("collection_name", collectionName),
|
|
mlog.Strings("partition_names", partitionNames),
|
|
mlog.Err(err))
|
|
return
|
|
}
|
|
|
|
err = merr.Error(resp.GetStatus())
|
|
if err != nil {
|
|
err = merr.Error(resp.GetStatus())
|
|
mlog.Warn(context.TODO(), "fail to show partitions",
|
|
mlog.String("collectionName", collectionName),
|
|
mlog.Strings("partitionNames", partitionNames),
|
|
mlog.Err(err))
|
|
return
|
|
}
|
|
|
|
for _, p := range resp.InMemoryPercentages {
|
|
loadProgress += p
|
|
}
|
|
loadProgress /= int64(len(partitionIDs))
|
|
|
|
if len(resp.GetRefreshProgress()) > 0 { // Compatibility for new Proxy with old QueryCoord
|
|
refreshProgress = resp.GetRefreshProgress()[0]
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func isPartitionKeyMode(ctx context.Context, dbName string, colName string) (bool, error) {
|
|
colSchema, err := globalMetaCache.GetCollectionSchema(ctx, dbName, colName)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
|
|
for _, fieldSchema := range colSchema.GetFields() {
|
|
if fieldSchema.IsPartitionKey {
|
|
return true, nil
|
|
}
|
|
}
|
|
|
|
return false, nil
|
|
}
|
|
|
|
func hasPartitionKeyModeField(schema *schemapb.CollectionSchema) bool {
|
|
for _, fieldSchema := range schema.GetFields() {
|
|
if fieldSchema.IsPartitionKey {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// getDefaultPartitionsInPartitionKeyMode only used in partition key mode
|
|
func getDefaultPartitionsInPartitionKeyMode(ctx context.Context, dbName string, collectionName string) ([]string, error) {
|
|
partitions, err := globalMetaCache.GetPartitions(ctx, dbName, collectionName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Make sure the order of the partition names got every time is the same
|
|
partitionNames, _, err := typeutil.RearrangePartitionsForPartitionKey(partitions)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return partitionNames, nil
|
|
}
|
|
|
|
func assignChannelsByPK(pks *schemapb.IDs, channelNames []string, insertMsg *msgstream.InsertMsg) (map[string][]int, error) {
|
|
hashValues, err := typeutil.HashPK2Channels(pks, channelNames)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
insertMsg.HashValues = hashValues
|
|
|
|
numChannels := len(channelNames)
|
|
if numChannels == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
numRows := len(insertMsg.HashValues)
|
|
avgCapacity := (numRows / numChannels) + 1
|
|
|
|
channel2RowOffsets := make(map[string][]int, numChannels)
|
|
|
|
for offset, channelID := range insertMsg.HashValues {
|
|
idx := int(channelID)
|
|
if idx >= numChannels {
|
|
continue
|
|
}
|
|
|
|
channelName := channelNames[idx]
|
|
|
|
if _, ok := channel2RowOffsets[channelName]; !ok {
|
|
channel2RowOffsets[channelName] = make([]int, 0, avgCapacity)
|
|
}
|
|
channel2RowOffsets[channelName] = append(channel2RowOffsets[channelName], offset)
|
|
}
|
|
|
|
return channel2RowOffsets, nil
|
|
}
|
|
|
|
func assignChannelsByNamespace(namespace string, channelNames []string, insertMsg *msgstream.InsertMsg) (map[string][]int, error) {
|
|
if len(channelNames) == 0 {
|
|
return nil, merr.WrapErrServiceInternalMsg("no virtual channels available for namespace sharding")
|
|
}
|
|
channelID := typeutil.HashNamespace2Channels(namespace, channelNames)
|
|
return assignChannelsByChannel(channelID, channelNames, insertMsg), nil
|
|
}
|
|
|
|
func assignChannelsByChannel(channelID uint32, channelNames []string, insertMsg *msgstream.InsertMsg) map[string][]int {
|
|
insertMsg.HashValues = make([]uint32, insertMsg.NumRows)
|
|
for i := range insertMsg.HashValues {
|
|
insertMsg.HashValues[i] = channelID
|
|
}
|
|
|
|
channelName := channelNames[channelID]
|
|
channel2RowOffsets := map[string][]int{
|
|
channelName: make([]int, 0, insertMsg.NRows()),
|
|
}
|
|
for i := range insertMsg.HashValues {
|
|
channel2RowOffsets[channelName] = append(channel2RowOffsets[channelName], i)
|
|
}
|
|
return channel2RowOffsets
|
|
}
|
|
|
|
func assignPartitionKeys(ctx context.Context, dbName string, collName string, keys []*planpb.GenericValue) ([]string, error) {
|
|
partitionNames, err := globalMetaCache.GetPartitionsIndex(ctx, dbName, collName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
schema, err := globalMetaCache.GetCollectionSchema(ctx, dbName, collName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
partitionKeyFieldSchema, err := typeutil.GetPartitionKeyFieldSchema(schema.CollectionSchema)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
hashedPartitionNames, err := typeutil2.HashKey2Partitions(partitionKeyFieldSchema, keys, partitionNames)
|
|
return hashedPartitionNames, err
|
|
}
|
|
|
|
func assignNamespacePartitionKey(ctx context.Context, dbName string, collName string, namespace *string) ([]string, error) {
|
|
if namespace == nil {
|
|
return nil, nil
|
|
}
|
|
|
|
return assignPartitionKeys(ctx, dbName, collName, []*planpb.GenericValue{
|
|
{Val: &planpb.GenericValue_StringVal{StringVal: *namespace}},
|
|
})
|
|
}
|
|
|
|
func ErrWithLog(logger *mlog.Logger, msg string, err error) error {
|
|
wrapErr := errors.Wrap(err, msg)
|
|
if logger != nil {
|
|
logger.Warn(context.TODO(), msg, mlog.Err(err))
|
|
return wrapErr
|
|
}
|
|
mlog.Warn(context.TODO(), msg, mlog.Err(err))
|
|
return wrapErr
|
|
}
|
|
|
|
func verifyDynamicFieldData(schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg, skipStaticFieldNameCheck bool) error {
|
|
for _, field := range insertMsg.FieldsData {
|
|
if field.GetFieldName() == common.MetaFieldName {
|
|
if !schema.EnableDynamicField {
|
|
return merr.WrapErrParameterInvalidMsg("without dynamic schema enabled, the field name cannot be set to %s", common.MetaFieldName)
|
|
}
|
|
for _, rowData := range field.GetScalars().GetJsonData().GetData() {
|
|
jsonData := make(map[string]interface{})
|
|
if err := json.Unmarshal(rowData, &jsonData); err != nil {
|
|
mlog.Info(context.TODO(), "insert invalid dynamic data, milvus only support json map",
|
|
mlog.ByteString("data", rowData),
|
|
mlog.Err(err),
|
|
)
|
|
return merr.WrapErrParameterInvalidMsg("invalid dynamic field data, only json map is supported: %s", err.Error())
|
|
}
|
|
if _, ok := jsonData[common.MetaFieldName]; ok {
|
|
return merr.WrapErrParameterInvalidMsg("cannot set json key to: %s", common.MetaFieldName)
|
|
}
|
|
if !skipStaticFieldNameCheck {
|
|
for _, f := range schema.GetFields() {
|
|
if _, ok := jsonData[f.GetName()]; ok {
|
|
mlog.Info(context.TODO(), "dynamic field name include the static field name", mlog.String("fieldName", f.GetName()))
|
|
return merr.WrapErrParameterInvalidMsg("dynamic field name cannot include the static field name: %s", f.GetName())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// doCheckDynamicFieldData is the shared implementation for dynamic field validation.
|
|
// When skipStaticFieldNameCheck is true, $meta keys matching static field names are
|
|
// allowed — this is needed for partial updates after schema evolution, where $meta
|
|
// may legitimately contain keys that now correspond to static columns.
|
|
func doCheckDynamicFieldData(schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg, skipStaticFieldNameCheck bool) error {
|
|
for _, data := range insertMsg.FieldsData {
|
|
if data.IsDynamic {
|
|
data.FieldName = common.MetaFieldName
|
|
return verifyDynamicFieldData(schema, insertMsg, skipStaticFieldNameCheck)
|
|
}
|
|
}
|
|
defaultData := make([][]byte, insertMsg.NRows())
|
|
for i := range defaultData {
|
|
defaultData[i] = []byte("{}")
|
|
}
|
|
dynamicData := autoGenDynamicFieldData(schema, defaultData)
|
|
insertMsg.FieldsData = append(insertMsg.FieldsData, dynamicData)
|
|
return nil
|
|
}
|
|
|
|
func checkDynamicFieldData(schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg) error {
|
|
return doCheckDynamicFieldData(schema, insertMsg, false)
|
|
}
|
|
|
|
// checkDynamicFieldDataForPartialUpdate is a relaxed version of checkDynamicFieldData
|
|
// for partial updates. After schema evolution, $meta may legitimately contain keys
|
|
// matching static field names (e.g., a dynamic field "end_timestamp" that was later
|
|
// added as a static column). This function validates JSON format and rejects the
|
|
// reserved $meta key, but skips the static field name conflict check so that
|
|
// existing dynamic field data is preserved.
|
|
func checkDynamicFieldDataForPartialUpdate(schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg) error {
|
|
return doCheckDynamicFieldData(schema, insertMsg, true)
|
|
}
|
|
|
|
func namespaceShardingEnabled(schema *schemapb.CollectionSchema) bool {
|
|
if schema == nil || !schema.GetEnableNamespace() {
|
|
return false
|
|
}
|
|
enabled, err := common.IsNamespaceShardingEnabled(schema.GetProperties()...)
|
|
return err == nil && enabled
|
|
}
|
|
|
|
func namespacePartitionKeyMode(schema *schemapb.CollectionSchema) bool {
|
|
return schema != nil && schema.GetEnableNamespace() && common.IsNamespaceModePartitionKey(schema.GetProperties()...)
|
|
}
|
|
|
|
func namespacePartitionKeyModeEnabled(schema *schemapb.CollectionSchema) bool {
|
|
return namespaceShardingEnabled(schema) && namespacePartitionKeyMode(schema)
|
|
}
|
|
|
|
func namespacePartitionModeEnabled(schema *schemapb.CollectionSchema) bool {
|
|
return schema != nil && schema.GetEnableNamespace() && common.IsNamespaceModePartition(schema.GetProperties()...)
|
|
}
|
|
|
|
func namespaceShardingChannelID(schema *schemapb.CollectionSchema, namespace *string, channelNames []string) (uint32, bool, error) {
|
|
if namespace == nil || !namespacePartitionKeyModeEnabled(schema) {
|
|
return 0, false, nil
|
|
}
|
|
if len(channelNames) == 0 {
|
|
return 0, false, merr.WrapErrServiceInternalMsg("no virtual channels available for namespace sharding")
|
|
}
|
|
return typeutil.HashNamespace2Channels(*namespace, channelNames), true, nil
|
|
}
|
|
|
|
func namespaceShardingChannel(schema *schemapb.CollectionSchema, namespace *string, channelNames []string) (string, bool, error) {
|
|
channelID, ok, err := namespaceShardingChannelID(schema, namespace, channelNames)
|
|
if !ok || err != nil {
|
|
return "", ok, err
|
|
}
|
|
return channelNames[channelID], true, nil
|
|
}
|
|
|
|
func preferredNodeForChannel(preferredNodes map[string]int64, channel string) int64 {
|
|
if preferredNodes == nil {
|
|
return 0
|
|
}
|
|
preferredNodeID, ok := preferredNodes[channel]
|
|
if !ok {
|
|
return 0
|
|
}
|
|
return preferredNodeID
|
|
}
|
|
|
|
func preferredNodeFromConcurrentMap(preferredNodes *typeutil.ConcurrentMap[string, int64], channel string) int64 {
|
|
if preferredNodes == nil {
|
|
return 0
|
|
}
|
|
preferredNodeID, ok := preferredNodes.Get(channel)
|
|
if !ok {
|
|
return 0
|
|
}
|
|
return preferredNodeID
|
|
}
|
|
|
|
func resolveNamespacePartitionName(schema *schemapb.CollectionSchema, namespace *string, partitionName string) (string, bool, error) {
|
|
if err := common.CheckNamespace(schema, namespace); err != nil {
|
|
return "", false, err
|
|
}
|
|
if !namespacePartitionModeEnabled(schema) {
|
|
return partitionName, false, nil
|
|
}
|
|
|
|
namespacePartitionName := *namespace
|
|
if err := validatePartitionTag(namespacePartitionName, true); err != nil {
|
|
return "", true, err
|
|
}
|
|
if partitionName != "" && partitionName != namespacePartitionName {
|
|
return "", true, merr.WrapErrParameterInvalidMsg("partition name %q mismatches namespace %q", partitionName, namespacePartitionName)
|
|
}
|
|
return namespacePartitionName, true, nil
|
|
}
|
|
|
|
func resolveNamespacePartitionNames(schema *schemapb.CollectionSchema, namespace *string, partitionNames []string) ([]string, bool, error) {
|
|
if err := common.CheckNamespace(schema, namespace); err != nil {
|
|
return nil, false, err
|
|
}
|
|
if !namespacePartitionModeEnabled(schema) {
|
|
return partitionNames, false, nil
|
|
}
|
|
|
|
namespacePartitionName := *namespace
|
|
if err := validatePartitionTag(namespacePartitionName, true); err != nil {
|
|
return nil, true, err
|
|
}
|
|
if len(partitionNames) == 0 {
|
|
return []string{namespacePartitionName}, true, nil
|
|
}
|
|
if len(partitionNames) == 1 && partitionNames[0] == namespacePartitionName {
|
|
return partitionNames, true, nil
|
|
}
|
|
return nil, true, merr.WrapErrParameterInvalidMsg("partition names %v mismatch namespace %q", partitionNames, namespacePartitionName)
|
|
}
|
|
|
|
func namespaceForPlan(schema *schemapb.CollectionSchema, namespace *string) *string {
|
|
if namespacePartitionModeEnabled(schema) {
|
|
return nil
|
|
}
|
|
return namespace
|
|
}
|
|
|
|
func addNamespaceData(schema *schemapb.CollectionSchema, insertMsg *msgstream.InsertMsg) error {
|
|
partitionName, namespaceAsPartition, err := resolveNamespacePartitionName(schema, insertMsg.Namespace, insertMsg.GetPartitionName())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !schema.GetEnableNamespace() {
|
|
return nil
|
|
}
|
|
if namespaceAsPartition {
|
|
insertMsg.PartitionName = partitionName
|
|
return nil
|
|
}
|
|
|
|
// check namespace field exists
|
|
namespaceField := typeutil.GetFieldByName(schema, common.NamespaceFieldName)
|
|
if namespaceField == nil {
|
|
return merr.WrapErrParameterInvalidMsg("namespace field not found")
|
|
}
|
|
|
|
// If namespace field data is already present, validate it instead of rejecting outright.
|
|
for _, fieldData := range insertMsg.FieldsData {
|
|
if fieldData.FieldId == namespaceField.FieldID {
|
|
ns := ""
|
|
if insertMsg.Namespace != nil {
|
|
ns = *insertMsg.Namespace
|
|
}
|
|
scalars := fieldData.GetScalars()
|
|
if scalars == nil {
|
|
return merr.WrapErrParameterInvalidMsg("invalid namespace field data layout")
|
|
}
|
|
strData := scalars.GetStringData()
|
|
if strData == nil {
|
|
return merr.WrapErrParameterInvalidMsg("invalid namespace field data layout")
|
|
}
|
|
for _, v := range strData.GetData() {
|
|
if v != ns {
|
|
return merr.WrapErrParameterInvalidMsg("namespace field value %q mismatches namespace %q", v, ns)
|
|
}
|
|
}
|
|
// Values are consistent with the namespace; nothing more to do.
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// set namespace field data
|
|
namespaceData := make([]string, insertMsg.NRows())
|
|
namespace := *insertMsg.Namespace
|
|
for i := range namespaceData {
|
|
namespaceData[i] = namespace
|
|
}
|
|
insertMsg.FieldsData = append(insertMsg.FieldsData, &schemapb.FieldData{
|
|
FieldName: namespaceField.Name,
|
|
FieldId: namespaceField.FieldID,
|
|
Type: namespaceField.DataType,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_StringData{
|
|
StringData: &schemapb.StringArray{
|
|
Data: namespaceData,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
return nil
|
|
}
|
|
|
|
func GetCachedCollectionSchema(ctx context.Context, dbName string, colName string) (*schemaInfo, error) {
|
|
if globalMetaCache != nil {
|
|
return globalMetaCache.GetCollectionSchema(ctx, dbName, colName)
|
|
}
|
|
return nil, merr.WrapErrServiceNotReady(paramtable.GetRole(), paramtable.GetNodeID(), "initialization")
|
|
}
|
|
|
|
func CheckDatabase(ctx context.Context, dbName string) bool {
|
|
if globalMetaCache != nil {
|
|
return globalMetaCache.HasDatabase(ctx, dbName)
|
|
}
|
|
return false
|
|
}
|
|
|
|
func SetReportValue(status *commonpb.Status, value int) {
|
|
if value <= 0 {
|
|
return
|
|
}
|
|
if !merr.Ok(status) {
|
|
return
|
|
}
|
|
if status.ExtraInfo == nil {
|
|
status.ExtraInfo = make(map[string]string)
|
|
}
|
|
status.ExtraInfo["report_value"] = strconv.Itoa(value)
|
|
}
|
|
|
|
func SetStorageCost(status *commonpb.Status, storageCost segcore.StorageCost) {
|
|
if !Params.QueryNodeCfg.StorageUsageTrackingEnabled.GetAsBool() {
|
|
return
|
|
}
|
|
if storageCost.ScannedTotalBytes <= 0 {
|
|
return
|
|
}
|
|
if !merr.Ok(status) {
|
|
return
|
|
}
|
|
if status.ExtraInfo == nil {
|
|
status.ExtraInfo = make(map[string]string)
|
|
// set report_value to 0 for compatibility, when extra info is not nil, there are always the default report_value
|
|
// see https://github.com/milvus-io/pymilvus/pull/2999, pymilvus didn't check the report_value is set and use the value
|
|
status.ExtraInfo["report_value"] = strconv.Itoa(0)
|
|
}
|
|
|
|
status.ExtraInfo["scanned_remote_bytes"] = strconv.FormatInt(storageCost.ScannedRemoteBytes, 10)
|
|
status.ExtraInfo["scanned_total_bytes"] = strconv.FormatInt(storageCost.ScannedTotalBytes, 10)
|
|
cacheHitRatio := float64(storageCost.ScannedTotalBytes-storageCost.ScannedRemoteBytes) / float64(storageCost.ScannedTotalBytes)
|
|
status.ExtraInfo["cache_hit_ratio"] = strconv.FormatFloat(cacheHitRatio, 'f', -1, 64)
|
|
}
|
|
|
|
func GetCostValue(status *commonpb.Status) int {
|
|
if status == nil || status.ExtraInfo == nil {
|
|
return 0
|
|
}
|
|
value, err := strconv.Atoi(status.ExtraInfo["report_value"])
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
return value
|
|
}
|
|
|
|
// final return value means value is valid or not
|
|
func GetStorageCost(status *commonpb.Status) (int64, int64, float64, bool) {
|
|
if status == nil || status.ExtraInfo == nil {
|
|
return 0, 0, 0, false
|
|
}
|
|
var scannedRemoteBytes int64
|
|
var scannedTotalBytes int64
|
|
var cacheHitRatio float64
|
|
var err error
|
|
if value, ok := status.ExtraInfo["scanned_remote_bytes"]; ok {
|
|
scannedRemoteBytes, err = strconv.ParseInt(value, 10, 64)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "scanned_remote_bytes is not a valid int64", mlog.String("value", value), mlog.Err(err))
|
|
return 0, 0, 0, false
|
|
}
|
|
} else {
|
|
return 0, 0, 0, false
|
|
}
|
|
if value, ok := status.ExtraInfo["scanned_total_bytes"]; ok {
|
|
scannedTotalBytes, err = strconv.ParseInt(value, 10, 64)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "scanned_total_bytes is not a valid int64", mlog.String("value", value), mlog.Err(err))
|
|
return 0, 0, 0, false
|
|
}
|
|
} else {
|
|
return 0, 0, 0, false
|
|
}
|
|
if value, ok := status.ExtraInfo["cache_hit_ratio"]; ok {
|
|
cacheHitRatio, err = strconv.ParseFloat(value, 64)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "cache_hit_ratio is not a valid float64", mlog.String("value", value), mlog.Err(err))
|
|
return 0, 0, 0, false
|
|
}
|
|
} else {
|
|
return 0, 0, 0, false
|
|
}
|
|
return scannedRemoteBytes, scannedTotalBytes, cacheHitRatio, true
|
|
}
|
|
|
|
// GetRequestInfo returns collection name and rateType of request and return tokens needed.
|
|
func GetRequestInfo(ctx context.Context, req proto.Message) (int64, map[int64][]int64, internalpb.RateType, int, error) {
|
|
switch r := req.(type) {
|
|
case *milvuspb.InsertRequest:
|
|
dbID, collToPartIDs, err := getCollectionAndPartitionID(ctx, req.(reqPartName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DMLInsert, proto.Size(r), err
|
|
case *milvuspb.UpsertRequest:
|
|
dbID, collToPartIDs, err := getCollectionAndPartitionID(ctx, req.(reqPartName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DMLInsert, proto.Size(r), err
|
|
case *milvuspb.DeleteRequest:
|
|
dbID, collToPartIDs, err := getCollectionAndPartitionID(ctx, req.(reqPartName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DMLDelete, proto.Size(r), err
|
|
case *milvuspb.ImportRequest:
|
|
dbID, collToPartIDs, err := getCollectionAndPartitionID(ctx, req.(reqPartName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DMLBulkLoad, proto.Size(r), err
|
|
case *milvuspb.SearchRequest:
|
|
dbID, collToPartIDs, err := getCollectionAndPartitionIDs(ctx, req.(reqPartNames))
|
|
return dbID, collToPartIDs, internalpb.RateType_DQLSearch, int(r.GetNq()), err
|
|
case *milvuspb.HybridSearchRequest:
|
|
dbID, collToPartIDs, err := getCollectionAndPartitionIDs(ctx, req.(reqPartNames))
|
|
nq := 0
|
|
for _, subReq := range r.GetRequests() {
|
|
nq += int(subReq.GetNq())
|
|
}
|
|
return dbID, collToPartIDs, internalpb.RateType_DQLSearch, nq, err
|
|
case *milvuspb.QueryRequest:
|
|
dbID, collToPartIDs, err := getCollectionAndPartitionIDs(ctx, req.(reqPartNames))
|
|
return dbID, collToPartIDs, internalpb.RateType_DQLQuery, 1, err // think of the query request's nq as 1
|
|
case *milvuspb.CreateCollectionRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLCollection, 1, nil
|
|
case *milvuspb.RefreshExternalCollectionRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLCollection, 1, nil
|
|
case *milvuspb.RestoreExternalSnapshotRequest:
|
|
return getDatabaseID(r.GetDbName()), map[int64][]int64{}, internalpb.RateType_DDLCollection, 1, nil
|
|
case *milvuspb.ExportSnapshotRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLCollection, 1, nil
|
|
case *milvuspb.DropCollectionRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLCollection, 1, nil
|
|
case *milvuspb.LoadCollectionRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLCollection, 1, nil
|
|
case *milvuspb.ReleaseCollectionRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLCollection, 1, nil
|
|
case *milvuspb.CreatePartitionRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLPartition, 1, nil
|
|
case *milvuspb.DropPartitionRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLPartition, 1, nil
|
|
case *milvuspb.LoadPartitionsRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLPartition, 1, nil
|
|
case *milvuspb.ReleasePartitionsRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLPartition, 1, nil
|
|
case *milvuspb.CreateIndexRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLIndex, 1, nil
|
|
case *milvuspb.DropIndexRequest:
|
|
dbID, collToPartIDs := getCollectionID(req.(reqCollName))
|
|
return dbID, collToPartIDs, internalpb.RateType_DDLIndex, 1, nil
|
|
case *milvuspb.FlushRequest:
|
|
db, err := globalMetaCache.GetDatabaseInfo(ctx, r.GetDbName())
|
|
if err != nil {
|
|
return util.InvalidDBID, map[int64][]int64{}, 0, 0, err
|
|
}
|
|
|
|
collToPartIDs := make(map[int64][]int64, 0)
|
|
for _, collectionName := range r.GetCollectionNames() {
|
|
collectionID, err := globalMetaCache.GetCollectionID(ctx, r.GetDbName(), collectionName)
|
|
if err != nil {
|
|
return util.InvalidDBID, map[int64][]int64{}, 0, 0, err
|
|
}
|
|
collToPartIDs[collectionID] = []int64{}
|
|
}
|
|
return db.dbID, collToPartIDs, internalpb.RateType_DDLFlush, 1, nil
|
|
case *milvuspb.ManualCompactionRequest:
|
|
// Use the db the request actually targets (normalized by
|
|
// DatabaseInterceptor), consistent with the sibling cases, so quota is
|
|
// accounted against the correct database. See milvus-io/milvus#50678.
|
|
dbInfo, err := globalMetaCache.GetDatabaseInfo(ctx, r.GetDbName())
|
|
if err != nil {
|
|
return util.InvalidDBID, map[int64][]int64{}, 0, 0, err
|
|
}
|
|
return dbInfo.dbID, map[int64][]int64{
|
|
r.GetCollectionID(): {},
|
|
}, internalpb.RateType_DDLCompaction, 1, nil
|
|
case *milvuspb.CreateDatabaseRequest:
|
|
mlog.Info(context.TODO(), "rate limiter CreateDatabaseRequest")
|
|
return util.InvalidDBID, map[int64][]int64{}, internalpb.RateType_DDLDB, 1, nil
|
|
case *milvuspb.DropDatabaseRequest:
|
|
mlog.Info(context.TODO(), "rate limiter DropDatabaseRequest")
|
|
return util.InvalidDBID, map[int64][]int64{}, internalpb.RateType_DDLDB, 1, nil
|
|
case *milvuspb.AlterDatabaseRequest:
|
|
return util.InvalidDBID, map[int64][]int64{}, internalpb.RateType_DDLDB, 1, nil
|
|
default: // TODO: support more request
|
|
if req == nil {
|
|
return util.InvalidDBID, map[int64][]int64{}, 0, 0, merr.WrapErrParameterInvalidMsg("null request")
|
|
}
|
|
mlog.RatedWarn(context.TODO(), rate.Limit(60), "not supported request type for rate limiter", mlog.String("type", reflect.TypeOf(req).String()))
|
|
return util.InvalidDBID, map[int64][]int64{}, 0, 0, nil
|
|
}
|
|
}
|
|
|
|
// GetFailedResponse returns failed response.
|
|
func GetFailedResponse(req any, err error) any {
|
|
switch req.(type) {
|
|
case *milvuspb.InsertRequest, *milvuspb.DeleteRequest, *milvuspb.UpsertRequest:
|
|
return failedMutationResult(err)
|
|
case *milvuspb.ImportRequest:
|
|
return &milvuspb.ImportResponse{
|
|
Status: merr.Status(err),
|
|
}
|
|
case *milvuspb.SearchRequest:
|
|
return &milvuspb.SearchResults{
|
|
Status: merr.Status(err),
|
|
}
|
|
case *milvuspb.QueryRequest:
|
|
return &milvuspb.QueryResults{
|
|
Status: merr.Status(err),
|
|
}
|
|
case *milvuspb.CreateCollectionRequest, *milvuspb.DropCollectionRequest,
|
|
*milvuspb.LoadCollectionRequest, *milvuspb.ReleaseCollectionRequest,
|
|
*milvuspb.CreatePartitionRequest, *milvuspb.DropPartitionRequest,
|
|
*milvuspb.LoadPartitionsRequest, *milvuspb.ReleasePartitionsRequest,
|
|
*milvuspb.CreateIndexRequest, *milvuspb.DropIndexRequest,
|
|
*milvuspb.CreateDatabaseRequest, *milvuspb.DropDatabaseRequest,
|
|
*milvuspb.AlterDatabaseRequest:
|
|
return merr.Status(err)
|
|
case *milvuspb.RestoreExternalSnapshotRequest:
|
|
return &milvuspb.RestoreExternalSnapshotResponse{
|
|
Status: merr.Status(err),
|
|
}
|
|
case *milvuspb.ExportSnapshotRequest:
|
|
return &milvuspb.ExportSnapshotResponse{
|
|
Status: merr.Status(err),
|
|
}
|
|
case *milvuspb.FlushRequest:
|
|
return &milvuspb.FlushResponse{
|
|
Status: merr.Status(err),
|
|
}
|
|
case *milvuspb.ManualCompactionRequest:
|
|
return &milvuspb.ManualCompactionResponse{
|
|
Status: merr.Status(err),
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func GetFunctionOutputFields(collSchema *schemapb.CollectionSchema) []string {
|
|
fields := make([]string, 0)
|
|
for _, fSchema := range collSchema.Functions {
|
|
fields = append(fields, fSchema.OutputFieldNames...)
|
|
}
|
|
return fields
|
|
}
|
|
|
|
func GetBM25FunctionOutputFields(collSchema *schemapb.CollectionSchema) []string {
|
|
fields := make([]string, 0)
|
|
for _, fSchema := range collSchema.Functions {
|
|
if fSchema.Type == schemapb.FunctionType_BM25 {
|
|
fields = append(fields, fSchema.OutputFieldNames...)
|
|
}
|
|
}
|
|
return fields
|
|
}
|
|
|
|
func GetMinHashFunctionOutputFields(collSchema *schemapb.CollectionSchema) []string {
|
|
fields := make([]string, 0)
|
|
for _, fSchema := range collSchema.Functions {
|
|
if fSchema.Type == schemapb.FunctionType_MinHash {
|
|
fields = append(fields, fSchema.OutputFieldNames...)
|
|
}
|
|
}
|
|
return fields
|
|
}
|
|
|
|
// getCollectionTTL returns ttl if collection's ttl is specified
|
|
// or return global ttl if collection's ttl is not specified
|
|
// this is a helper util wrapping common.GetCollectionTTL without returning error
|
|
func getCollectionTTL(pairs []*commonpb.KeyValuePair) uint64 {
|
|
ttl, err := common.GetCollectionTTL(pairs)
|
|
if err != nil {
|
|
mlog.Error(context.TODO(), "failed to get collection ttl, use default ttl", mlog.Err(err))
|
|
}
|
|
if ttl < 0 {
|
|
return 0
|
|
}
|
|
return uint64(ttl)
|
|
}
|
|
|
|
// reconstructStructFieldData regroups flattened sub-fields (named "structName[fieldName]")
|
|
// back into StructArrayField entries, restoring original field names for the user-facing response.
|
|
// It modifies sub-field FieldName in place; callers must not reuse the input slice afterwards.
|
|
func reconstructStructFieldData(
|
|
fieldsData []*schemapb.FieldData,
|
|
outputFields []string,
|
|
schema *schemapb.CollectionSchema,
|
|
) ([]*schemapb.FieldData, []string) {
|
|
if len(outputFields) == 1 && outputFields[0] == "count(*)" {
|
|
return fieldsData, outputFields
|
|
}
|
|
|
|
if len(schema.StructArrayFields) == 0 {
|
|
return fieldsData, outputFields
|
|
}
|
|
|
|
regularFieldIDs := make(map[int64]interface{})
|
|
subFieldToStructMap := make(map[int64]int64)
|
|
groupedStructFields := make(map[int64][]*schemapb.FieldData)
|
|
structFieldNames := make(map[int64]string)
|
|
reconstructedOutputFields := make([]string, 0, len(fieldsData))
|
|
|
|
// record all regular field IDs
|
|
for _, field := range schema.Fields {
|
|
regularFieldIDs[field.GetFieldID()] = nil
|
|
}
|
|
|
|
// build the mapping from sub-field ID to struct field ID
|
|
for _, structField := range schema.StructArrayFields {
|
|
for _, subField := range structField.GetFields() {
|
|
subFieldToStructMap[subField.GetFieldID()] = structField.GetFieldID()
|
|
}
|
|
structFieldNames[structField.GetFieldID()] = structField.GetName()
|
|
}
|
|
|
|
newFieldsData := make([]*schemapb.FieldData, 0, len(fieldsData))
|
|
for _, field := range fieldsData {
|
|
fieldID := field.GetFieldId()
|
|
if _, ok := regularFieldIDs[fieldID]; ok {
|
|
newFieldsData = append(newFieldsData, field)
|
|
reconstructedOutputFields = append(reconstructedOutputFields, field.GetFieldName())
|
|
} else if structFieldID, ok := subFieldToStructMap[fieldID]; ok {
|
|
groupedStructFields[structFieldID] = append(groupedStructFields[structFieldID], field)
|
|
} else {
|
|
newFieldsData = append(newFieldsData, field)
|
|
reconstructedOutputFields = append(reconstructedOutputFields, field.GetFieldName())
|
|
}
|
|
}
|
|
|
|
for structFieldID, fields := range groupedStructFields {
|
|
// Restore original field names (from "structName[fieldName]" to "fieldName")
|
|
// for the user-facing response.
|
|
for _, field := range fields {
|
|
originalName, err := extractOriginalFieldName(field.FieldName)
|
|
if err != nil {
|
|
mlog.Error(context.TODO(), "failed to extract original field name from struct field",
|
|
mlog.String("fieldName", field.FieldName),
|
|
mlog.Err(err))
|
|
} else {
|
|
field.FieldName = originalName
|
|
}
|
|
}
|
|
|
|
newFieldsData = append(newFieldsData, &schemapb.FieldData{
|
|
FieldName: structFieldNames[structFieldID],
|
|
FieldId: structFieldID,
|
|
Type: schemapb.DataType_ArrayOfStruct,
|
|
Field: &schemapb.FieldData_StructArrays{StructArrays: &schemapb.StructArrayField{Fields: fields}},
|
|
})
|
|
reconstructedOutputFields = append(reconstructedOutputFields, structFieldNames[structFieldID])
|
|
}
|
|
|
|
return newFieldsData, reconstructedOutputFields
|
|
}
|
|
|
|
func reconstructStructFieldDataForQuery(results *milvuspb.QueryResults, schema *schemapb.CollectionSchema) {
|
|
fieldsData, outputFields := reconstructStructFieldData(
|
|
results.FieldsData,
|
|
results.OutputFields,
|
|
schema,
|
|
)
|
|
results.FieldsData = fieldsData
|
|
results.OutputFields = outputFields
|
|
}
|
|
|
|
func reconstructStructFieldDataForSearch(results *milvuspb.SearchResults, schema *schemapb.CollectionSchema) {
|
|
if results.Results == nil {
|
|
return
|
|
}
|
|
fieldsData, outputFields := reconstructStructFieldData(
|
|
results.Results.FieldsData,
|
|
results.Results.OutputFields,
|
|
schema,
|
|
)
|
|
results.Results.FieldsData = fieldsData
|
|
results.Results.OutputFields = outputFields
|
|
}
|
|
|
|
func getColTimezone(colInfo *collectionInfo) string {
|
|
timezone, _ := funcutil.TryGetAttrByKeyFromRepeatedKV(common.TimezoneKey, colInfo.properties)
|
|
if timezone == "" {
|
|
timezone = common.DefaultTimezone
|
|
}
|
|
return timezone
|
|
}
|
|
|
|
// timestamptzUTC2IsoStr converts Timestamptz (Unix Microsecond) data
|
|
// within FieldData results into ISO-8601 strings, applying the correct
|
|
// timezone offset and using the optimized format (microsecond precision, no trailing zeros).
|
|
func timestamptzUTC2IsoStr(results []*schemapb.FieldData, colTimezone string) error {
|
|
location, err := time.LoadLocation(colTimezone)
|
|
if err != nil {
|
|
mlog.Error(context.TODO(), "invalid timezone", mlog.String("timezone", colTimezone), mlog.Err(err))
|
|
return merr.WrapErrParameterInvalidMsg("got invalid default timezone: %s", colTimezone)
|
|
}
|
|
|
|
for _, fieldData := range results {
|
|
if fieldData.GetType() != schemapb.DataType_Timestamptz {
|
|
continue
|
|
}
|
|
|
|
scalarField := fieldData.GetScalars()
|
|
|
|
// Guard against nil scalars or missing timestamp data
|
|
if scalarField == nil || scalarField.GetTimestamptzData() == nil {
|
|
if longData := scalarField.GetLongData(); longData != nil && len(longData.GetData()) > 0 {
|
|
mlog.Warn(context.TODO(), "field data is not Timestamptz data", mlog.String("fieldName", fieldData.GetFieldName()))
|
|
return merr.WrapErrParameterInvalidMsg("field data for '%s' is not Timestamptz data", fieldData.GetFieldName())
|
|
}
|
|
// Handle the case of an empty field (e.g., all nulls), skip if no data to process.
|
|
continue
|
|
}
|
|
|
|
utcTimestamps := scalarField.GetTimestamptzData().GetData()
|
|
isoStrings := make([]string, len(utcTimestamps))
|
|
|
|
// CORE CHANGE: Use the optimized formatting function
|
|
for i, ts := range utcTimestamps {
|
|
// 1. Convert Unix Microsecond (UTC) to a time.Time object (still in UTC).
|
|
t := time.UnixMicro(ts).UTC()
|
|
|
|
// 2. Adjust the time object to the target location.
|
|
localTime := t.In(location)
|
|
|
|
// 3. Format using the optimized logic (max 6 digits, no trailing zeros)
|
|
isoStrings[i] = timestamptz.FormatTimeMicroWithoutTrailingZeros(localTime)
|
|
}
|
|
|
|
// Replace the TimestamptzData with the new StringData in place.
|
|
fieldData.GetScalars().Data = &schemapb.ScalarField_StringData{
|
|
StringData: &schemapb.StringArray{
|
|
Data: isoStrings,
|
|
},
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// extractFields is a helper function to extract specific integer fields from a time.Time object.
|
|
// Supported fields are: "year", "month", "day", "hour", "minute", "second", "microsecond", "nanosecond".
|
|
func extractFields(t time.Time, fieldList []string) ([]int64, error) {
|
|
extractedValues := make([]int64, 0, len(fieldList))
|
|
for _, field := range fieldList {
|
|
var val int64
|
|
switch strings.ToLower(field) {
|
|
case common.TszYear:
|
|
val = int64(t.Year())
|
|
case common.TszMonth:
|
|
val = int64(t.Month())
|
|
case common.TszDay:
|
|
val = int64(t.Day())
|
|
case common.TszHour:
|
|
val = int64(t.Hour())
|
|
case common.TszMinute:
|
|
val = int64(t.Minute())
|
|
case common.TszSecond:
|
|
val = int64(t.Second())
|
|
case common.TszMicrosecond:
|
|
val = int64(t.Nanosecond() / 1000)
|
|
default:
|
|
return nil, merr.WrapErrParameterInvalidMsg("unsupported field for extraction: %s, fields should be seprated by ',' or ' '", field)
|
|
}
|
|
extractedValues = append(extractedValues, val)
|
|
}
|
|
return extractedValues, nil
|
|
}
|
|
|
|
func extractFieldsFromResults(results []*schemapb.FieldData, timezone string, fieldList []string) error {
|
|
targetLocation, err := time.LoadLocation(timezone)
|
|
if err != nil {
|
|
mlog.Error(context.TODO(), "invalid timezone", mlog.String("timezone", timezone), mlog.Err(err))
|
|
return merr.WrapErrParameterInvalidMsg("got invalid timezone: %s", timezone)
|
|
}
|
|
|
|
for _, fieldData := range results {
|
|
if fieldData.GetType() != schemapb.DataType_Timestamptz {
|
|
continue
|
|
}
|
|
|
|
scalarField := fieldData.GetScalars()
|
|
if scalarField == nil || scalarField.GetTimestamptzData() == nil {
|
|
if longData := scalarField.GetLongData(); longData != nil && len(longData.GetData()) > 0 {
|
|
mlog.Warn(context.TODO(), "field data is not Timestamptz data, but found LongData instead", mlog.String("fieldName", fieldData.GetFieldName()))
|
|
return merr.WrapErrParameterInvalidMsg("field data for '%s' is not Timestamptz data", fieldData.GetFieldName())
|
|
}
|
|
continue
|
|
}
|
|
|
|
utcTimestamps := scalarField.GetTimestamptzData().GetData()
|
|
extractedResults := make([]*schemapb.ScalarField, 0, len(fieldList))
|
|
|
|
for _, ts := range utcTimestamps {
|
|
t := time.UnixMicro(ts).UTC()
|
|
localTime := t.In(targetLocation)
|
|
|
|
values, err := extractFields(localTime, fieldList)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
valuesScalarField := &schemapb.ScalarField_LongData{
|
|
LongData: &schemapb.LongArray{
|
|
Data: values,
|
|
},
|
|
}
|
|
extractedResults = append(extractedResults, &schemapb.ScalarField{
|
|
Data: valuesScalarField,
|
|
})
|
|
}
|
|
|
|
fieldData.GetScalars().Data = &schemapb.ScalarField_ArrayData{
|
|
ArrayData: &schemapb.ArrayArray{
|
|
Data: extractedResults,
|
|
ElementType: schemapb.DataType_Int64,
|
|
},
|
|
}
|
|
fieldData.Type = schemapb.DataType_Array
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func genFunctionFields(ctx context.Context, insertMsg *msgstream.InsertMsg, schema *schemaInfo, partialUpdate bool) error {
|
|
allowNonBM25Outputs := common.GetCollectionAllowInsertNonBM25FunctionOutputs(schema.Properties)
|
|
fieldIDs := lo.Map(insertMsg.FieldsData, func(fieldData *schemapb.FieldData, _ int) int64 {
|
|
id, _ := schema.MapFieldID(fieldData.FieldName)
|
|
return id
|
|
})
|
|
|
|
// Since PartialUpdate is supported, the field_data here may not be complete
|
|
needProcessFunctions, err := typeutil.GetNeedProcessFunctions(fieldIDs, schema.Functions, allowNonBM25Outputs, partialUpdate)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "Check upsert field error,", mlog.String("collectionName", schema.Name), mlog.Err(err))
|
|
return err
|
|
}
|
|
|
|
if embedding.HasNonBM25AndMinHashFunctions(schema.Functions, []int64{}) {
|
|
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-genFunctionFields-call-function-udf")
|
|
defer sp.End()
|
|
exec, err := embedding.NewFunctionExecutor(schema.CollectionSchema, needProcessFunctions, &models.ModelExtraInfo{ClusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), DBName: insertMsg.GetDbName()})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
sp.AddEvent("Create-function-udf")
|
|
if err := exec.ProcessInsert(ctx, insertMsg); err != nil {
|
|
return err
|
|
}
|
|
sp.AddEvent("Call-function-udf")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func getBM25FunctionOfAnnsField(fieldID int64, functions []*schemapb.FunctionSchema) (*schemapb.FunctionSchema, bool) {
|
|
return lo.Find(functions, func(function *schemapb.FunctionSchema) bool {
|
|
return function.GetType() == schemapb.FunctionType_BM25 && function.OutputFieldIds[0] == fieldID
|
|
})
|
|
}
|
|
|
|
// failMetricLabel classifies a request failure for the ProxyFunctionCall
|
|
// metric, matching the fail_input/fail_system split that
|
|
// requestutil.ParseMetricLabel emits at the gRPC interceptor; the legacy
|
|
// bare "fail" value must not reappear in this metric's value domain.
|
|
func failMetricLabel(err error) string {
|
|
// Client cancellation is neither party's failure; keep it out of the
|
|
// fail_system bucket (parity with ParseMetricLabel).
|
|
if errors.Is(err, context.Canceled) {
|
|
return metrics.CancelLabel
|
|
}
|
|
if merr.GetErrorType(err) == merr.InputError {
|
|
return metrics.FailInputLabel
|
|
}
|
|
return metrics.FailSystemLabel
|
|
}
|