mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
feat: [2.6.16] support partial update ops for Array fields (#49811)
pr: [#49328](https://github.com/milvus-io/milvus/pull/49328) pr: [#49724](https://github.com/milvus-io/milvus/pull/49724) pr: [#49763](https://github.com/milvus-io/milvus/pull/49763) pr: [#49698](https://github.com/milvus-io/milvus/pull/49698) issue: [#49241](https://github.com/milvus-io/milvus/issues/49241) issue: [#49746](https://github.com/milvus-io/milvus/issues/49746) issue: [#49634](https://github.com/milvus-io/milvus/issues/49634) ## Summary Backport the 2.6 partial update op series to `hotfix-2.6.16`: - support `ARRAY_APPEND` and `ARRAY_REMOVE` partial update ops for Array fields - expose `fieldOps` through REST upsert - preserve existing Array rows when an op payload row is null ## Target branch Base branch: `hotfix-2.6.16` ## Cherry-picks - `5849977c408be2abd063d13c318e970bb4515f06` from [#49328](https://github.com/milvus-io/milvus/pull/49328) - `017ee8e5d97ec891c442489b1354b60676ca3b15` from [#49724](https://github.com/milvus-io/milvus/pull/49724) - `4e10473389843ff2cedc1911083af5a93fa83ec6` from [#49763](https://github.com/milvus-io/milvus/pull/49763) - `824c642c71154952bf70eaaf7b57cbcd7d08e67f` backports the applicable WAL test/recovery stabilization from `73dc8d4034fd352f5c69fd47266fccc062704feb` / [#49698](https://github.com/milvus-io/milvus/pull/49698) after omitting newer rate-limit API changes that do not exist on `hotfix-2.6.16` The partial update cherry-picks applied cleanly on top of `milvus/hotfix-2.6.16`. ## Additional revert - `2753c8defc7a860446d1f5f76b11a50c3cc71550` reverts `8ae21f715094abc92e00369e89a7208681eecfed` to restore the 2.6.16 build environment image version after CI reported Conan 2.x in the newer image. ## Validation - `git diff --check milvus/hotfix-2.6.16...HEAD` - attempted targeted Go test for `internal/streamingnode/server/wal/adaptor`, blocked locally because this worktree lacks `rocksdb.pc` and `milvus_core.pc` PR CI is the validation gate for this backport. --------- Signed-off-by: Wei Liu <wei.liu@zilliz.com> Signed-off-by: Zhen Ye <chyezh@outlook.com> Co-authored-by: Zhen Ye <chyezh@outlook.com>
This commit is contained in:
@@ -5,11 +5,11 @@ IMAGE_ARCH=amd64
|
||||
OS_NAME=ubuntu22.04
|
||||
|
||||
# for services.builder.image in docker-compose.yml
|
||||
DATE_VERSION=20260506-b5f57d9
|
||||
LATEST_DATE_VERSION=20260506-b5f57d9
|
||||
DATE_VERSION=20260318-39568a2
|
||||
LATEST_DATE_VERSION=20260318-39568a2
|
||||
# for services.gpubuilder.image in docker-compose.yml
|
||||
GPU_DATE_VERSION=20260506-b5f57d9
|
||||
LATEST_GPU_DATE_VERSION=20260506-b5f57d9
|
||||
GPU_DATE_VERSION=20260318-39568a2
|
||||
LATEST_GPU_DATE_VERSION=20260318-39568a2
|
||||
|
||||
# for other services in docker-compose.yml
|
||||
MINIO_ADDRESS=minio:9000
|
||||
|
||||
@@ -57,6 +57,11 @@ type columnBasedDataOption struct {
|
||||
// deferredErr captures construction-time errors from builder helpers (e.g. WithStructArrayColumn)
|
||||
// so they surface on InsertRequest/UpsertRequest rather than panicking in the chain.
|
||||
deferredErr error
|
||||
|
||||
// partialOps carries per-field FieldPartialUpdateOp directives. Keyed by
|
||||
// field name. Entries with REPLACE (or nil) are treated as no-ops and are
|
||||
// not serialized onto the wire.
|
||||
partialOps map[string]*schemapb.FieldPartialUpdateOp
|
||||
}
|
||||
|
||||
func (opt *columnBasedDataOption) WriteBackPKs(_ *entity.Schema, _ column.Column) error {
|
||||
@@ -379,6 +384,59 @@ func (opt *columnBasedDataOption) WithPartialUpdate(partialUpdate bool) *columnB
|
||||
return opt
|
||||
}
|
||||
|
||||
// WithArrayAppend declares that the Array field `fieldName` should be merged
|
||||
// with ARRAY_APPEND semantics during an Upsert. The server implicitly enables
|
||||
// partial_update when any non-REPLACE op is present, so callers do not need
|
||||
// to also invoke WithPartialUpdate(true).
|
||||
func (opt *columnBasedDataOption) WithArrayAppend(fieldName string) *columnBasedDataOption {
|
||||
return opt.WithFieldPartialOp(fieldName, schemapb.FieldPartialUpdateOp_ARRAY_APPEND)
|
||||
}
|
||||
|
||||
// WithArrayRemove declares that the Array field `fieldName` should be merged
|
||||
// with ARRAY_REMOVE semantics during an Upsert. See WithArrayAppend for the
|
||||
// implicit partial_update promotion.
|
||||
func (opt *columnBasedDataOption) WithArrayRemove(fieldName string) *columnBasedDataOption {
|
||||
return opt.WithFieldPartialOp(fieldName, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE)
|
||||
}
|
||||
|
||||
// WithFieldPartialOp attaches an explicit FieldPartialUpdateOp to the field
|
||||
// with name `fieldName`. Intended for advanced callers; typical users should
|
||||
// prefer the op-specific helpers (WithArrayAppend, WithArrayRemove).
|
||||
func (opt *columnBasedDataOption) WithFieldPartialOp(fieldName string, op schemapb.FieldPartialUpdateOp_OpType) *columnBasedDataOption {
|
||||
if op == schemapb.FieldPartialUpdateOp_REPLACE {
|
||||
// REPLACE is the default; clear any prior directive rather than
|
||||
// transmitting a no-op message.
|
||||
if opt.partialOps != nil {
|
||||
delete(opt.partialOps, fieldName)
|
||||
}
|
||||
return opt
|
||||
}
|
||||
if opt.partialOps == nil {
|
||||
opt.partialOps = make(map[string]*schemapb.FieldPartialUpdateOp)
|
||||
}
|
||||
opt.partialOps[fieldName] = &schemapb.FieldPartialUpdateOp{FieldName: fieldName, Op: op}
|
||||
return opt
|
||||
}
|
||||
|
||||
// buildFieldOps materializes the recorded FieldPartialUpdateOp directives
|
||||
// into a proto-ready slice. Only non-REPLACE ops are emitted — REPLACE is
|
||||
// the on-wire default and emitting it would waste bytes on every upsert.
|
||||
//
|
||||
// The returned slice is independent of the input fieldsData; a field
|
||||
// referenced by an op that was not in fieldsData is still emitted so the
|
||||
// server can surface a validation error rather than silently drop the
|
||||
// op. Client-side filtering would hide user typos.
|
||||
func (opt *columnBasedDataOption) buildFieldOps() []*schemapb.FieldPartialUpdateOp {
|
||||
if len(opt.partialOps) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*schemapb.FieldPartialUpdateOp, 0, len(opt.partialOps))
|
||||
for _, op := range opt.partialOps {
|
||||
out = append(out, op)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (opt *columnBasedDataOption) CollectionName() string {
|
||||
return opt.collName
|
||||
}
|
||||
@@ -408,13 +466,22 @@ func (opt *columnBasedDataOption) UpsertRequest(coll *entity.Collection) (*milvu
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Materialize any WithArrayAppend/WithArrayRemove/WithFieldPartialOp
|
||||
// directives into UpsertRequest.field_ops. Auto-promote partial_update
|
||||
// when any non-REPLACE op is present.
|
||||
fieldOps := opt.buildFieldOps()
|
||||
partialUpdate := opt.partialUpdate
|
||||
if len(fieldOps) > 0 {
|
||||
partialUpdate = true
|
||||
}
|
||||
return &milvuspb.UpsertRequest{
|
||||
CollectionName: opt.collName,
|
||||
PartitionName: opt.partitionName,
|
||||
FieldsData: fieldsData,
|
||||
NumRows: uint32(rowNum),
|
||||
SchemaTimestamp: coll.UpdateTimestamp,
|
||||
PartialUpdate: opt.partialUpdate,
|
||||
PartialUpdate: partialUpdate,
|
||||
FieldOps: fieldOps,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -470,12 +537,18 @@ func (opt *rowBasedDataOption) UpsertRequest(coll *entity.Collection) (*milvuspb
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
fieldOps := opt.buildFieldOps()
|
||||
partialUpdate := opt.partialUpdate
|
||||
if len(fieldOps) > 0 {
|
||||
partialUpdate = true
|
||||
}
|
||||
return &milvuspb.UpsertRequest{
|
||||
CollectionName: opt.collName,
|
||||
PartitionName: opt.partitionName,
|
||||
FieldsData: fieldsData,
|
||||
NumRows: uint32(rowNum),
|
||||
PartialUpdate: opt.partialUpdate,
|
||||
PartialUpdate: partialUpdate,
|
||||
FieldOps: fieldOps,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
// 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 milvusclient
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/client/v2/column"
|
||||
"github.com/milvus-io/milvus/client/v2/entity"
|
||||
)
|
||||
|
||||
func buildPartialOpTestCollection() *entity.Collection {
|
||||
schema := entity.NewSchema().
|
||||
WithField(entity.NewField().WithName("id").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
|
||||
WithField(entity.NewField().WithName("tags").WithDataType(entity.FieldTypeArray).WithElementType(entity.FieldTypeInt64).WithMaxCapacity(32))
|
||||
schema.CollectionName = "partial_op_test"
|
||||
return &entity.Collection{Name: "partial_op_test", Schema: schema}
|
||||
}
|
||||
|
||||
func buildPartialOpOption(modify func(opt *columnBasedDataOption)) *columnBasedDataOption {
|
||||
idCol := column.NewColumnInt64("id", []int64{1, 2})
|
||||
tagsCol := column.NewColumnInt64Array("tags", [][]int64{{1}, {2, 3}})
|
||||
opt := NewColumnBasedInsertOption("partial_op_test").WithColumns(idCol, tagsCol)
|
||||
if modify != nil {
|
||||
modify(opt)
|
||||
}
|
||||
return opt
|
||||
}
|
||||
|
||||
func findOp(ops []*schemapb.FieldPartialUpdateOp, name string) *schemapb.FieldPartialUpdateOp {
|
||||
for _, op := range ops {
|
||||
if op.GetFieldName() == name {
|
||||
return op
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestWithArrayAppendEmitsFieldOpAndAutoEnablesPartialUpdate(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithArrayAppend("tags")
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, req.GetPartialUpdate(), "ARRAY_APPEND should auto-enable partial_update")
|
||||
|
||||
tagsOp := findOp(req.GetFieldOps(), "tags")
|
||||
require.NotNil(t, tagsOp)
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_APPEND, tagsOp.GetOp())
|
||||
assert.Equal(t, "tags", tagsOp.GetFieldName())
|
||||
|
||||
// FieldData must remain clean — no op leakage into the data message.
|
||||
for _, fd := range req.GetFieldsData() {
|
||||
_ = fd
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithArrayRemoveEmitsFieldOp(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithArrayRemove("tags")
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, req.GetPartialUpdate())
|
||||
|
||||
tagsOp := findOp(req.GetFieldOps(), "tags")
|
||||
require.NotNil(t, tagsOp)
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, tagsOp.GetOp())
|
||||
}
|
||||
|
||||
func TestWithFieldPartialOpReplaceClearsPriorDirective(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithArrayAppend("tags")
|
||||
o.WithFieldPartialOp("tags", schemapb.FieldPartialUpdateOp_REPLACE)
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
assert.False(t, req.GetPartialUpdate(), "REPLACE should clear prior non-REPLACE op")
|
||||
assert.Empty(t, req.GetFieldOps())
|
||||
}
|
||||
|
||||
func TestWithFieldPartialOpReplaceWithoutPriorIsNoOp(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithFieldPartialOp("tags", schemapb.FieldPartialUpdateOp_REPLACE)
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
assert.False(t, req.GetPartialUpdate())
|
||||
assert.Empty(t, req.GetFieldOps())
|
||||
assert.Empty(t, opt.partialOps)
|
||||
}
|
||||
|
||||
func TestPartialOpDoesNotOverrideExplicitPartialUpdate(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithPartialUpdate(true)
|
||||
o.WithArrayAppend("tags")
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, req.GetPartialUpdate())
|
||||
}
|
||||
|
||||
func TestExplicitPartialUpdateFalseIsPromotedByOp(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithPartialUpdate(false)
|
||||
o.WithArrayAppend("tags")
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, req.GetPartialUpdate(), "non-REPLACE op should promote partial_update")
|
||||
}
|
||||
|
||||
func TestPartialOpForUnknownFieldStillEmitted(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithArrayAppend("does_not_exist")
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
// Unknown-field ops are forwarded as-is so the server can return a
|
||||
// descriptive validation error rather than the client silently
|
||||
// dropping the directive.
|
||||
assert.True(t, req.GetPartialUpdate())
|
||||
assert.Len(t, req.GetFieldOps(), 1)
|
||||
assert.Equal(t, "does_not_exist", req.GetFieldOps()[0].GetFieldName())
|
||||
}
|
||||
|
||||
func TestBuildFieldOpsReturnsNilWhenEmpty(t *testing.T) {
|
||||
opt := &columnBasedDataOption{}
|
||||
assert.Nil(t, opt.buildFieldOps())
|
||||
}
|
||||
|
||||
func TestRowBasedUpsertEmitsFieldOps(t *testing.T) {
|
||||
coll := buildPartialOpTestCollection()
|
||||
rows := []any{
|
||||
map[string]any{"id": int64(1), "tags": []int64{10}},
|
||||
map[string]any{"id": int64(2), "tags": []int64{20, 30}},
|
||||
}
|
||||
opt := NewRowBasedInsertOption(coll.Name, rows...)
|
||||
opt.WithArrayAppend("tags")
|
||||
|
||||
req, err := opt.UpsertRequest(coll)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, req.GetPartialUpdate())
|
||||
tagsOp := findOp(req.GetFieldOps(), "tags")
|
||||
require.NotNil(t, tagsOp)
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_APPEND, tagsOp.GetOp())
|
||||
}
|
||||
|
||||
func TestMultipleFieldOpsEmittedTogether(t *testing.T) {
|
||||
opt := buildPartialOpOption(func(o *columnBasedDataOption) {
|
||||
o.WithArrayAppend("tags")
|
||||
o.WithFieldPartialOp("other", schemapb.FieldPartialUpdateOp_ARRAY_REMOVE)
|
||||
})
|
||||
req, err := opt.UpsertRequest(buildPartialOpTestCollection())
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, req.GetFieldOps(), 2)
|
||||
seen := map[string]schemapb.FieldPartialUpdateOp_OpType{}
|
||||
for _, o := range req.GetFieldOps() {
|
||||
seen[o.GetFieldName()] = o.GetOp()
|
||||
}
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_APPEND, seen["tags"])
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, seen["other"])
|
||||
}
|
||||
@@ -832,6 +832,7 @@ func (h *HandlersV1) upsert(c *gin.Context) {
|
||||
httpReq.CollectionName = singleUpsertReq.CollectionName
|
||||
httpReq.Data = []map[string]interface{}{singleUpsertReq.Data}
|
||||
httpReq.PartialUpdate = singleUpsertReq.PartialUpdate
|
||||
httpReq.FieldOps = singleUpsertReq.FieldOps
|
||||
}
|
||||
if httpReq.CollectionName == "" || httpReq.Data == nil {
|
||||
log.Warn("high level restful api, upsert require parameter: [collectionName, data], but miss")
|
||||
@@ -847,6 +848,15 @@ func (h *HandlersV1) upsert(c *gin.Context) {
|
||||
NumRows: uint32(len(httpReq.Data)),
|
||||
PartialUpdate: httpReq.PartialUpdate,
|
||||
}
|
||||
fieldOps, err := buildFieldPartialUpdateOps(httpReq.FieldOps)
|
||||
if err != nil {
|
||||
HTTPAbortReturn(c, http.StatusOK, gin.H{
|
||||
HTTPReturnCode: merr.Code(err),
|
||||
HTTPReturnMessage: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
req.FieldOps = fieldOps
|
||||
c.Set(ContextRequest, req)
|
||||
username, _ := c.Get(ContextUsername)
|
||||
ctx := proxy.NewContextWithMetadata(c, username.(string), req.DbName)
|
||||
|
||||
@@ -1321,6 +1321,15 @@ func (h *HandlersV2) upsert(ctx context.Context, c *gin.Context, anyReq any, dbN
|
||||
PartialUpdate: httpReq.PartialUpdate,
|
||||
// PartitionName: "_default",
|
||||
}
|
||||
fieldOps, err := buildFieldPartialUpdateOps(httpReq.FieldOps)
|
||||
if err != nil {
|
||||
HTTPAbortReturn(c, http.StatusOK, gin.H{
|
||||
HTTPReturnCode: merr.Code(err),
|
||||
HTTPReturnMessage: err.Error(),
|
||||
})
|
||||
return nil, err
|
||||
}
|
||||
req.FieldOps = fieldOps
|
||||
c.Set(ContextRequest, req)
|
||||
|
||||
collSchema, err := h.GetCollectionSchema(ctx, c, dbName, httpReq.CollectionName)
|
||||
|
||||
@@ -68,17 +68,19 @@ type SingleInsertReq struct {
|
||||
}
|
||||
|
||||
type UpsertReq struct {
|
||||
DbName string `json:"dbName"`
|
||||
CollectionName string `json:"collectionName" validate:"required"`
|
||||
Data []map[string]interface{} `json:"data" validate:"required"`
|
||||
PartialUpdate bool `json:"partialUpdate"`
|
||||
DbName string `json:"dbName"`
|
||||
CollectionName string `json:"collectionName" validate:"required"`
|
||||
Data []map[string]interface{} `json:"data" validate:"required"`
|
||||
PartialUpdate bool `json:"partialUpdate"`
|
||||
FieldOps []FieldPartialUpdateOpReq `json:"fieldOps"`
|
||||
}
|
||||
|
||||
type SingleUpsertReq struct {
|
||||
DbName string `json:"dbName"`
|
||||
CollectionName string `json:"collectionName" validate:"required"`
|
||||
Data map[string]interface{} `json:"data" validate:"required"`
|
||||
PartialUpdate bool `json:"partialUpdate"`
|
||||
DbName string `json:"dbName"`
|
||||
CollectionName string `json:"collectionName" validate:"required"`
|
||||
Data map[string]interface{} `json:"data" validate:"required"`
|
||||
PartialUpdate bool `json:"partialUpdate"`
|
||||
FieldOps []FieldPartialUpdateOpReq `json:"fieldOps"`
|
||||
}
|
||||
|
||||
type SearchReq struct {
|
||||
|
||||
@@ -19,6 +19,7 @@ package httpserver
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -325,16 +326,55 @@ func (req *CollectionFilterReq) GetDbName() string { return req.DbName }
|
||||
func (req *CollectionFilterReq) GetCollectionName() string { return req.CollectionName }
|
||||
|
||||
type CollectionDataReq struct {
|
||||
DbName string `json:"dbName"`
|
||||
CollectionName string `json:"collectionName" binding:"required"`
|
||||
PartitionName string `json:"partitionName"`
|
||||
Data []map[string]interface{} `json:"data" binding:"required"`
|
||||
PartialUpdate bool `json:"partialUpdate"`
|
||||
DbName string `json:"dbName"`
|
||||
CollectionName string `json:"collectionName" binding:"required"`
|
||||
PartitionName string `json:"partitionName"`
|
||||
Data []map[string]interface{} `json:"data" binding:"required"`
|
||||
PartialUpdate bool `json:"partialUpdate"`
|
||||
FieldOps []FieldPartialUpdateOpReq `json:"fieldOps"`
|
||||
}
|
||||
|
||||
func (req *CollectionDataReq) GetDbName() string { return req.DbName }
|
||||
func (req *CollectionDataReq) GetCollectionName() string { return req.CollectionName }
|
||||
|
||||
type FieldPartialUpdateOpReq struct {
|
||||
FieldName string `json:"fieldName"`
|
||||
Op string `json:"op"`
|
||||
}
|
||||
|
||||
func buildFieldPartialUpdateOps(fieldOps []FieldPartialUpdateOpReq) ([]*schemapb.FieldPartialUpdateOp, error) {
|
||||
if len(fieldOps) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
ops := make([]*schemapb.FieldPartialUpdateOp, 0, len(fieldOps))
|
||||
for _, fieldOp := range fieldOps {
|
||||
op, err := parseFieldPartialUpdateOp(fieldOp.Op)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ops = append(ops, &schemapb.FieldPartialUpdateOp{
|
||||
FieldName: fieldOp.FieldName,
|
||||
Op: op,
|
||||
})
|
||||
}
|
||||
return ops, nil
|
||||
}
|
||||
|
||||
func parseFieldPartialUpdateOp(op string) (schemapb.FieldPartialUpdateOp_OpType, error) {
|
||||
switch strings.ToUpper(strings.TrimSpace(op)) {
|
||||
case "REPLACE":
|
||||
return schemapb.FieldPartialUpdateOp_REPLACE, nil
|
||||
case "ARRAY_APPEND":
|
||||
return schemapb.FieldPartialUpdateOp_ARRAY_APPEND, nil
|
||||
case "ARRAY_REMOVE":
|
||||
return schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, nil
|
||||
default:
|
||||
return schemapb.FieldPartialUpdateOp_REPLACE,
|
||||
merr.WrapErrParameterInvalidMsg("unsupported partial update op: " + op)
|
||||
}
|
||||
}
|
||||
|
||||
type SearchReqV2 struct {
|
||||
DbName string `json:"dbName"`
|
||||
CollectionName string `json:"collectionName" binding:"required"`
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/requestutil"
|
||||
)
|
||||
|
||||
@@ -50,3 +51,24 @@ func TestRequestV2_GetCollectionName(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildFieldPartialUpdateOps(t *testing.T) {
|
||||
ops, err := buildFieldPartialUpdateOps([]FieldPartialUpdateOpReq{
|
||||
{FieldName: "tags", Op: "ARRAY_APPEND"},
|
||||
{FieldName: "scores", Op: "ARRAY_REMOVE"},
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, ops, 2)
|
||||
assert.Equal(t, "tags", ops[0].GetFieldName())
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_APPEND, ops[0].GetOp())
|
||||
assert.Equal(t, "scores", ops[1].GetFieldName())
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, ops[1].GetOp())
|
||||
}
|
||||
|
||||
func TestBuildFieldPartialUpdateOps_RejectsUnknownOp(t *testing.T) {
|
||||
_, err := buildFieldPartialUpdateOps([]FieldPartialUpdateOpReq{
|
||||
{FieldName: "tags", Op: "ARRAY_EXTEND"},
|
||||
})
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported partial update op")
|
||||
}
|
||||
|
||||
@@ -384,6 +384,7 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error {
|
||||
it.insertFieldData = typeutil.PrepareResultFieldData(existFieldData, int64(upsertIDSize))
|
||||
|
||||
if len(updateIdxInUpsert) > 0 {
|
||||
fieldOpMap := buildFieldOpMap(it.req)
|
||||
// Note: For fields containing default values, default values need to be set according to valid data during insertion,
|
||||
// but query results fields do not set valid data when returning default value fields,
|
||||
// therefore valid data needs to be manually set to true
|
||||
@@ -411,6 +412,26 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error {
|
||||
existPKToIndex[pk] = j
|
||||
}
|
||||
|
||||
// Index upsert FieldData by name once, so per-row Array-op application
|
||||
// can locate the matching payload in O(1).
|
||||
upsertByName := lo.SliceToMap(it.upsertMsg.InsertMsg.GetFieldsData(), func(f *schemapb.FieldData) (string, *schemapb.FieldData) {
|
||||
return f.GetFieldName(), f
|
||||
})
|
||||
existByName := lo.SliceToMap(existFieldData, func(f *schemapb.FieldData) (string, *schemapb.FieldData) {
|
||||
return f.GetFieldName(), f
|
||||
})
|
||||
genericUpdateFieldData := it.upsertMsg.InsertMsg.GetFieldsData()
|
||||
if fieldOpMap != nil {
|
||||
genericUpdateFieldData = make([]*schemapb.FieldData, 0, len(genericUpdateFieldData))
|
||||
for _, fieldData := range it.upsertMsg.InsertMsg.GetFieldsData() {
|
||||
op, ok := fieldOpMap[fieldData.GetFieldName()]
|
||||
if ok && op != schemapb.FieldPartialUpdateOp_REPLACE {
|
||||
continue
|
||||
}
|
||||
genericUpdateFieldData = append(genericUpdateFieldData, fieldData)
|
||||
}
|
||||
}
|
||||
|
||||
baseIdx := 0
|
||||
for _, idx := range updateIdxInUpsert {
|
||||
typeutil.AppendIDs(it.deletePKs, upsertIDs, idx)
|
||||
@@ -420,12 +441,43 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error {
|
||||
return merr.WrapErrParameterInvalidMsg("primary key not found in exist data mapping")
|
||||
}
|
||||
typeutil.AppendFieldData(it.insertFieldData, existFieldData, int64(existIndex))
|
||||
err := typeutil.UpdateFieldData(it.insertFieldData, it.upsertMsg.InsertMsg.GetFieldsData(), int64(baseIdx), int64(idx))
|
||||
baseIdx += 1
|
||||
if err != nil {
|
||||
if err := typeutil.UpdateFieldData(it.insertFieldData, genericUpdateFieldData, int64(baseIdx), int64(idx)); err != nil {
|
||||
log.Info("update field data failed", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
// Per-row Array partial-op: generic REPLACE intentionally
|
||||
// skipped non-REPLACE op fields, so dstField still carries the
|
||||
// existing row appended above. Overwrite it with
|
||||
// ApplyArrayRowOp(existing row, upsert row).
|
||||
if fieldOpMap != nil {
|
||||
for _, dstField := range it.insertFieldData {
|
||||
op, ok := fieldOpMap[dstField.GetFieldName()]
|
||||
if !ok || op == schemapb.FieldPartialUpdateOp_REPLACE {
|
||||
continue
|
||||
}
|
||||
fieldSchema, schemaErr := it.schema.schemaHelper.GetFieldFromName(dstField.GetFieldName())
|
||||
if schemaErr != nil {
|
||||
return schemaErr
|
||||
}
|
||||
existField, ok := existByName[dstField.GetFieldName()]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
upsertField, ok := upsertByName[dstField.GetFieldName()]
|
||||
if !ok {
|
||||
return merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("partial-update op field %q missing from upsert payload", dstField.GetFieldName()))
|
||||
}
|
||||
if err := applyArrayPartialOpOnRow(
|
||||
dstField, existField, upsertField,
|
||||
int64(baseIdx), int64(existIndex), int64(idx),
|
||||
op, fieldSchema.GetElementType(), readMaxCapacity(fieldSchema),
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
baseIdx++
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1067,6 +1119,19 @@ func (it *upsertTask) PreExecute(ctx context.Context) error {
|
||||
}
|
||||
it.schema = schema
|
||||
|
||||
// Validate any FieldPartialUpdateOp directives attached to
|
||||
// UpsertRequest.field_ops. A non-REPLACE op implicitly promotes the
|
||||
// request to partial_update=true so users do not need to set both
|
||||
// fields explicitly.
|
||||
nonReplaceSeen, err := validateFieldPartialUpdateOps(it.req, schema.CollectionSchema)
|
||||
if err != nil {
|
||||
log.Warn("validate field partial update ops failed", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
if nonReplaceSeen && !it.req.GetPartialUpdate() {
|
||||
it.req.PartialUpdate = true
|
||||
}
|
||||
|
||||
it.partitionKeyMode, err = isPartitionKeyMode(ctx, it.req.GetDbName(), collectionName)
|
||||
if err != nil {
|
||||
log.Warn("check partition key mode failed",
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
// 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 (
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/common"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/merr"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
|
||||
)
|
||||
|
||||
// applyArrayPartialOpOnRow overwrites dstField[baseIdx] with
|
||||
// ApplyArrayRowOp(existField[existIdx], upsertField[upsertIdx]) using the
|
||||
// given op + elementType + maxCapacity. dstField is assumed to already carry
|
||||
// the existing row at baseIdx, placed there by AppendFieldData in the per-row
|
||||
// merge loop.
|
||||
func applyArrayPartialOpOnRow(
|
||||
dstField, existField, upsertField *schemapb.FieldData,
|
||||
baseIdx, existIdx, upsertIdx int64,
|
||||
op schemapb.FieldPartialUpdateOp_OpType,
|
||||
elementType schemapb.DataType,
|
||||
maxCapacity int,
|
||||
) error {
|
||||
if existField.GetScalars() == nil || existField.GetScalars().GetArrayData() == nil {
|
||||
return merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("existing data for field %q is not an Array", dstField.GetFieldName()))
|
||||
}
|
||||
if upsertField.GetScalars() == nil || upsertField.GetScalars().GetArrayData() == nil {
|
||||
return merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("upsert payload for field %q is not an Array", dstField.GetFieldName()))
|
||||
}
|
||||
dstArr := dstField.GetScalars().GetArrayData()
|
||||
if dstArr == nil {
|
||||
return merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("merged buffer for field %q is not an Array", dstField.GetFieldName()))
|
||||
}
|
||||
// Skip when the upsert row is explicitly null — nothing to append/remove.
|
||||
if len(upsertField.ValidData) > int(upsertIdx) && !upsertField.ValidData[upsertIdx] {
|
||||
return nil
|
||||
}
|
||||
baseRow := existField.GetScalars().GetArrayData().Data[existIdx]
|
||||
updateRow := upsertField.GetScalars().GetArrayData().Data[upsertIdx]
|
||||
merged, err := typeutil.ApplyArrayRowOp(baseRow, updateRow, op, elementType, maxCapacity)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
dstArr.Data[baseIdx] = merged
|
||||
if len(dstField.ValidData) > int(baseIdx) {
|
||||
dstField.ValidData[baseIdx] = true
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// hasNonReplacePartialOp reports whether req carries any non-REPLACE
|
||||
// FieldPartialUpdateOp. Used to decide implicit promotion to
|
||||
// partial_update=true and to short-circuit validation when no op is set.
|
||||
func hasNonReplacePartialOp(req *milvuspb.UpsertRequest) bool {
|
||||
for _, op := range req.GetFieldOps() {
|
||||
if op.GetOp() != schemapb.FieldPartialUpdateOp_REPLACE {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// validateFieldPartialUpdateOps validates FieldPartialUpdateOp entries
|
||||
// attached to an UpsertRequest. It enforces:
|
||||
//
|
||||
// - Each op targets a field present in the collection schema.
|
||||
// - Array-specific ops (ARRAY_APPEND, ARRAY_REMOVE) require the target
|
||||
// field to have DataType_Array.
|
||||
// - The target field must not be the primary key.
|
||||
// - The FieldData carried in fields_data must declare a matching
|
||||
// ElementType when the op targets an Array field.
|
||||
// - For ARRAY_APPEND, the per-row payload length must not, on its own,
|
||||
// already exceed max_capacity. The merged length is additionally
|
||||
// enforced at merge time by ApplyArrayRowOp.
|
||||
// - A given field_name appears at most once across field_ops.
|
||||
//
|
||||
// Returns (true, nil) when at least one non-REPLACE op was observed, so
|
||||
// the caller may implicitly promote partial_update=true.
|
||||
func validateFieldPartialUpdateOps(req *milvuspb.UpsertRequest, schema *schemapb.CollectionSchema) (bool, error) {
|
||||
fieldOps := req.GetFieldOps()
|
||||
if len(fieldOps) == 0 {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// Precompute PK names and a lookup table of FieldData by name so we
|
||||
// can validate payload alignment in O(1) per op.
|
||||
pkFields := make(map[string]struct{})
|
||||
for _, f := range schema.GetFields() {
|
||||
if f.GetIsPrimaryKey() {
|
||||
pkFields[f.GetName()] = struct{}{}
|
||||
}
|
||||
}
|
||||
fieldDataByName := make(map[string]*schemapb.FieldData, len(req.GetFieldsData()))
|
||||
for _, fd := range req.GetFieldsData() {
|
||||
fieldDataByName[fd.GetFieldName()] = fd
|
||||
}
|
||||
|
||||
nonReplaceSeen := false
|
||||
seenOpFields := make(map[string]struct{}, len(fieldOps))
|
||||
for _, opMsg := range fieldOps {
|
||||
name := opMsg.GetFieldName()
|
||||
if name == "" {
|
||||
return false, merr.WrapErrParameterInvalidMsg("FieldPartialUpdateOp.field_name is required")
|
||||
}
|
||||
if _, dup := seenOpFields[name]; dup {
|
||||
return false, merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("duplicate partial-update op for field %q", name))
|
||||
}
|
||||
seenOpFields[name] = struct{}{}
|
||||
|
||||
op := opMsg.GetOp()
|
||||
if op == schemapb.FieldPartialUpdateOp_REPLACE {
|
||||
// An explicit REPLACE is legal but indistinguishable from no
|
||||
// op at all. Accept silently — no further validation needed.
|
||||
continue
|
||||
}
|
||||
nonReplaceSeen = true
|
||||
|
||||
if _, isPK := pkFields[name]; isPK {
|
||||
return false, merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("field %q is the primary key and cannot carry a partial-update op", name))
|
||||
}
|
||||
|
||||
fieldSchema, err := findFieldSchemaByName(schema, name)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
switch op {
|
||||
case schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE:
|
||||
if fieldSchema.GetDataType() != schemapb.DataType_Array {
|
||||
return false, merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("op %s requires Array field, but field %q is %s",
|
||||
op.String(), name, fieldSchema.GetDataType().String()))
|
||||
}
|
||||
default:
|
||||
return false, merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("unsupported partial update op: %s", op.String()))
|
||||
}
|
||||
|
||||
fd, ok := fieldDataByName[name]
|
||||
if !ok {
|
||||
return false, merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("partial-update op targets field %q not present in fields_data", name))
|
||||
}
|
||||
// Reject malformed FieldData early -- a request that declares an
|
||||
// Array op but carries no ArrayData would otherwise panic on a nil
|
||||
// deref inside the merge path.
|
||||
if fd.GetScalars() == nil || fd.GetScalars().GetArrayData() == nil {
|
||||
return false, merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("partial-update op field %q payload is not an Array", name))
|
||||
}
|
||||
if got := fd.GetScalars().GetArrayData().GetElementType(); got != schemapb.DataType_None && got != fieldSchema.GetElementType() {
|
||||
return false, merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("field %q expects element type %s but request provides %s",
|
||||
name, fieldSchema.GetElementType().String(), got.String()))
|
||||
}
|
||||
|
||||
if op == schemapb.FieldPartialUpdateOp_ARRAY_APPEND {
|
||||
if err := checkArrayAppendPayloadWithinCapacity(fd, fieldSchema); err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nonReplaceSeen, nil
|
||||
}
|
||||
|
||||
// buildFieldOpMap returns a fieldName → op lookup. Callers use it during
|
||||
// merge to decide which op (if any) to apply to each Array field. A nil
|
||||
// map (no ops) is a valid return for the all-REPLACE default.
|
||||
func buildFieldOpMap(req *milvuspb.UpsertRequest) map[string]schemapb.FieldPartialUpdateOp_OpType {
|
||||
fieldOps := req.GetFieldOps()
|
||||
if len(fieldOps) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]schemapb.FieldPartialUpdateOp_OpType, len(fieldOps))
|
||||
for _, op := range fieldOps {
|
||||
out[op.GetFieldName()] = op.GetOp()
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// findFieldSchemaByName locates a field by name in the given collection
|
||||
// schema. Returns a descriptive parameter-invalid error when not found.
|
||||
func findFieldSchemaByName(schema *schemapb.CollectionSchema, name string) (*schemapb.FieldSchema, error) {
|
||||
for _, f := range schema.GetFields() {
|
||||
if f.GetName() == name {
|
||||
return f, nil
|
||||
}
|
||||
}
|
||||
return nil, merr.WrapErrParameterInvalidMsg(fmt.Sprintf("field %q not found in collection schema", name))
|
||||
}
|
||||
|
||||
// checkArrayAppendPayloadWithinCapacity checks that the per-row payload
|
||||
// length itself does not already exceed max_capacity. The final merged
|
||||
// length is enforced at merge time by ApplyArrayRowOp.
|
||||
//
|
||||
// When max_capacity is not declared or is non-positive, the check is a
|
||||
// no-op (matches the behavior of legacy upserts with no capacity gate).
|
||||
func checkArrayAppendPayloadWithinCapacity(fd *schemapb.FieldData, fieldSchema *schemapb.FieldSchema) error {
|
||||
maxCap := readMaxCapacity(fieldSchema)
|
||||
if maxCap < 0 {
|
||||
return nil
|
||||
}
|
||||
rows := fd.GetScalars().GetArrayData().GetData()
|
||||
for rowIdx, row := range rows {
|
||||
got := perRowArrayLen(row, fieldSchema.GetElementType())
|
||||
if got > maxCap {
|
||||
return merr.WrapErrParameterInvalidMsg(
|
||||
fmt.Sprintf("ARRAY_APPEND payload for field %q row %d has length %d exceeding max_capacity %d",
|
||||
fd.GetFieldName(), rowIdx, got, maxCap))
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// readMaxCapacity returns the declared max_capacity of an Array field, or
|
||||
// -1 when missing / malformed. -1 is a sentinel understood by both the
|
||||
// proxy pre-check and typeutil.ApplyArrayRowOp as "no capacity gate",
|
||||
// keeping the two sides consistent.
|
||||
func readMaxCapacity(fieldSchema *schemapb.FieldSchema) int {
|
||||
for _, kv := range fieldSchema.GetTypeParams() {
|
||||
if kv.GetKey() != common.MaxCapacityKey {
|
||||
continue
|
||||
}
|
||||
v, err := strconv.Atoi(kv.GetValue())
|
||||
if err != nil {
|
||||
return -1
|
||||
}
|
||||
return v
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// perRowArrayLen returns the element count of a single Array row given
|
||||
// its declared ElementType. Unsupported element types return 0.
|
||||
func perRowArrayLen(row *schemapb.ScalarField, elementType schemapb.DataType) int {
|
||||
switch elementType {
|
||||
case schemapb.DataType_Bool:
|
||||
return len(row.GetBoolData().GetData())
|
||||
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
|
||||
return len(row.GetIntData().GetData())
|
||||
case schemapb.DataType_Int64:
|
||||
return len(row.GetLongData().GetData())
|
||||
case schemapb.DataType_Float:
|
||||
return len(row.GetFloatData().GetData())
|
||||
case schemapb.DataType_Double:
|
||||
return len(row.GetDoubleData().GetData())
|
||||
case schemapb.DataType_VarChar, schemapb.DataType_String:
|
||||
return len(row.GetStringData().GetData())
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,355 @@
|
||||
// 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 (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/common"
|
||||
)
|
||||
|
||||
func arrayIntFieldSchema(name string, isPK bool, maxCap int) *schemapb.FieldSchema {
|
||||
typeParams := []*commonpb.KeyValuePair{}
|
||||
if maxCap > 0 {
|
||||
typeParams = append(typeParams, &commonpb.KeyValuePair{Key: common.MaxCapacityKey, Value: itoa(maxCap)})
|
||||
}
|
||||
return &schemapb.FieldSchema{
|
||||
Name: name,
|
||||
IsPrimaryKey: isPK,
|
||||
DataType: schemapb.DataType_Array,
|
||||
ElementType: schemapb.DataType_Int64,
|
||||
TypeParams: typeParams,
|
||||
}
|
||||
}
|
||||
|
||||
func itoa(n int) string {
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
neg := n < 0
|
||||
if neg {
|
||||
n = -n
|
||||
}
|
||||
buf := make([]byte, 0, 12)
|
||||
for n > 0 {
|
||||
buf = append([]byte{byte('0' + n%10)}, buf...)
|
||||
n /= 10
|
||||
}
|
||||
if neg {
|
||||
return "-" + string(buf)
|
||||
}
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
func arrayLongFieldData(name string, rows [][]int64) *schemapb.FieldData {
|
||||
rowFields := make([]*schemapb.ScalarField, len(rows))
|
||||
for i, row := range rows {
|
||||
rowFields[i] = &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: row}}}
|
||||
}
|
||||
return &schemapb.FieldData{
|
||||
Type: schemapb.DataType_Array,
|
||||
FieldName: name,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
|
||||
Data: rowFields,
|
||||
ElementType: schemapb.DataType_Int64,
|
||||
}},
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
func op(name string, t schemapb.FieldPartialUpdateOp_OpType) *schemapb.FieldPartialUpdateOp {
|
||||
return &schemapb.FieldPartialUpdateOp{FieldName: name, Op: t}
|
||||
}
|
||||
|
||||
func TestHasNonReplacePartialOp(t *testing.T) {
|
||||
assert.False(t, hasNonReplacePartialOp(&milvuspb.UpsertRequest{}))
|
||||
assert.False(t, hasNonReplacePartialOp(&milvuspb.UpsertRequest{FieldOps: []*schemapb.FieldPartialUpdateOp{
|
||||
op("tags", schemapb.FieldPartialUpdateOp_REPLACE),
|
||||
}}))
|
||||
assert.True(t, hasNonReplacePartialOp(&milvuspb.UpsertRequest{FieldOps: []*schemapb.FieldPartialUpdateOp{
|
||||
op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND),
|
||||
}}))
|
||||
assert.True(t, hasNonReplacePartialOp(&milvuspb.UpsertRequest{FieldOps: []*schemapb.FieldPartialUpdateOp{
|
||||
op("a", schemapb.FieldPartialUpdateOp_REPLACE),
|
||||
op("b", schemapb.FieldPartialUpdateOp_ARRAY_REMOVE),
|
||||
}}))
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_NoOp(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})}}
|
||||
seen, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, seen)
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_ReplaceIgnored(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_REPLACE)},
|
||||
}
|
||||
seen, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, seen)
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_AppendOnArrayField(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1, 2}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
seen, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, seen)
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RemoveOnArrayField(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_REMOVE)},
|
||||
}
|
||||
seen, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, seen)
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsEmptyFieldName(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "field_name is required")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsDuplicateField(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{
|
||||
op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND),
|
||||
op("tags", schemapb.FieldPartialUpdateOp_ARRAY_REMOVE),
|
||||
},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "duplicate")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsOpOnPK(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", true, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "primary key")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsUnknownField(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("other", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("other", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "not found")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsNonArrayField(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
|
||||
{Name: "tags", DataType: schemapb.DataType_VarChar},
|
||||
}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "Array field")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsElementTypeMismatch(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
|
||||
{Name: "tags", DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_VarChar},
|
||||
}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "element type")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsUnsupportedOpEnum(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_OpType(9999))},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported partial update op")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsOpWithoutFieldData(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
// FieldsData empty — op targets a field that was not sent
|
||||
FieldsData: []*schemapb.FieldData{},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "not present in fields_data")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsPayloadExceedingCapacity(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 2)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1, 2, 3, 4}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "max_capacity")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_SkipsCapacityWhenUnset(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 0)}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1, 2, 3, 4, 5}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
seen, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, seen)
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_RejectsNilArrayData(t *testing.T) {
|
||||
// FieldData declares type=Array but carries no ArrayData: the merge
|
||||
// path would deref nil; validate must reject up front.
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{arrayIntFieldSchema("tags", false, 8)}}
|
||||
fdNoArray := &schemapb.FieldData{
|
||||
FieldName: "tags",
|
||||
Type: schemapb.DataType_Array,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{}},
|
||||
}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{fdNoArray},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
_, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "not an Array")
|
||||
}
|
||||
|
||||
func TestValidateFieldPartialUpdateOps_SkipsCapacityWhenMalformed(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{{
|
||||
Name: "tags", DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_Int64,
|
||||
TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxCapacityKey, Value: "not-a-number"}},
|
||||
}}}
|
||||
req := &milvuspb.UpsertRequest{
|
||||
FieldsData: []*schemapb.FieldData{arrayLongFieldData("tags", [][]int64{{1, 2, 3, 4}})},
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
}
|
||||
seen, err := validateFieldPartialUpdateOps(req, schema)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, seen)
|
||||
}
|
||||
|
||||
func TestBuildFieldOpMap(t *testing.T) {
|
||||
assert.Nil(t, buildFieldOpMap(&milvuspb.UpsertRequest{}))
|
||||
got := buildFieldOpMap(&milvuspb.UpsertRequest{FieldOps: []*schemapb.FieldPartialUpdateOp{
|
||||
op("a", schemapb.FieldPartialUpdateOp_ARRAY_APPEND),
|
||||
op("b", schemapb.FieldPartialUpdateOp_ARRAY_REMOVE),
|
||||
}})
|
||||
require.Len(t, got, 2)
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_APPEND, got["a"])
|
||||
assert.Equal(t, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, got["b"])
|
||||
}
|
||||
|
||||
func TestPerRowArrayLen(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
row *schemapb.ScalarField
|
||||
et schemapb.DataType
|
||||
want int
|
||||
}{
|
||||
{"bool", &schemapb.ScalarField{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true, false}}}}, schemapb.DataType_Bool, 2},
|
||||
{"int32", &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}}}}, schemapb.DataType_Int32, 3},
|
||||
{"int8", &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1}}}}, schemapb.DataType_Int8, 1},
|
||||
{"int16", &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1, 2}}}}, schemapb.DataType_Int16, 2},
|
||||
{"int64", &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2, 3, 4}}}}, schemapb.DataType_Int64, 4},
|
||||
{"float", &schemapb.ScalarField{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1}}}}, schemapb.DataType_Float, 1},
|
||||
{"double", &schemapb.ScalarField{Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{1, 2}}}}, schemapb.DataType_Double, 2},
|
||||
{"varchar", &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"a", "b", "c"}}}}, schemapb.DataType_VarChar, 3},
|
||||
{"string", &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"a"}}}}, schemapb.DataType_String, 1},
|
||||
{"unsupported", &schemapb.ScalarField{}, schemapb.DataType_JSON, 0},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
assert.Equal(t, tc.want, perRowArrayLen(tc.row, tc.et), tc.name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadMaxCapacity(t *testing.T) {
|
||||
fs := &schemapb.FieldSchema{TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxCapacityKey, Value: "42"}}}
|
||||
assert.Equal(t, 42, readMaxCapacity(fs))
|
||||
|
||||
fs = &schemapb.FieldSchema{TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxCapacityKey, Value: "not-int"}}}
|
||||
assert.Equal(t, -1, readMaxCapacity(fs))
|
||||
|
||||
fs = &schemapb.FieldSchema{TypeParams: []*commonpb.KeyValuePair{{Key: "other", Value: "42"}}}
|
||||
assert.Equal(t, -1, readMaxCapacity(fs))
|
||||
|
||||
assert.Equal(t, -1, readMaxCapacity(&schemapb.FieldSchema{}))
|
||||
}
|
||||
|
||||
func TestFindFieldSchemaByName(t *testing.T) {
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{{Name: "tags"}, {Name: "scores"}}}
|
||||
got, err := findFieldSchemaByName(schema, "scores")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "scores", got.GetName())
|
||||
|
||||
_, err = findFieldSchemaByName(schema, "missing")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestItoa(t *testing.T) {
|
||||
cases := []struct {
|
||||
in int
|
||||
want string
|
||||
}{{0, "0"}, {1, "1"}, {42, "42"}, {-7, "-7"}, {1024, "1024"}}
|
||||
for _, tc := range cases {
|
||||
assert.Equal(t, tc.want, itoa(tc.in))
|
||||
}
|
||||
}
|
||||
@@ -1383,6 +1383,109 @@ func TestUpsertTask_queryPreExecute_PureUpdate(t *testing.T) {
|
||||
assert.Equal(t, []int32{600, 700}, valueField.GetScalars().GetIntData().GetData())
|
||||
}
|
||||
|
||||
func TestUpsertTask_queryPreExecute_ArrayPartialOpSkipsGenericReplace(t *testing.T) {
|
||||
schema := newSchemaInfo(&schemapb.CollectionSchema{
|
||||
Name: "test_array_partial_op_skip_replace",
|
||||
Fields: []*schemapb.FieldSchema{
|
||||
{FieldID: 100, Name: "id", IsPrimaryKey: true, DataType: schemapb.DataType_Int64},
|
||||
{
|
||||
FieldID: 101,
|
||||
Name: "tags",
|
||||
DataType: schemapb.DataType_Array,
|
||||
ElementType: schemapb.DataType_Int64,
|
||||
Nullable: true,
|
||||
TypeParams: []*commonpb.KeyValuePair{
|
||||
{Key: common.MaxCapacityKey, Value: "16"},
|
||||
},
|
||||
},
|
||||
{FieldID: 102, Name: "note", DataType: schemapb.DataType_VarChar, Nullable: true},
|
||||
},
|
||||
})
|
||||
|
||||
upsertTags := arrayLongFieldData("tags", [][]int64{{3}})
|
||||
upsertTags.FieldId = 101
|
||||
upsertTags.ValidData = []bool{true, false}
|
||||
upsertData := []*schemapb.FieldData{
|
||||
{
|
||||
FieldName: "id", FieldId: 100, Type: schemapb.DataType_Int64,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2}}}}},
|
||||
},
|
||||
upsertTags,
|
||||
{
|
||||
FieldName: "note", FieldId: 102, Type: schemapb.DataType_VarChar,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"new1"}}}}},
|
||||
ValidData: []bool{true, false},
|
||||
},
|
||||
}
|
||||
numRows := uint64(2)
|
||||
|
||||
existingTags := arrayLongFieldData("tags", [][]int64{{10, 20}, {100, 200}})
|
||||
existingTags.FieldId = 101
|
||||
existingTags.ValidData = []bool{true, true}
|
||||
mockQueryResult := &milvuspb.QueryResults{
|
||||
Status: merr.Success(),
|
||||
FieldsData: []*schemapb.FieldData{
|
||||
{
|
||||
FieldName: "id", FieldId: 100, Type: schemapb.DataType_Int64,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1, 2}}}}},
|
||||
},
|
||||
existingTags,
|
||||
{
|
||||
FieldName: "note", FieldId: 102, Type: schemapb.DataType_VarChar,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"old1", "old2"}}}}},
|
||||
ValidData: []bool{true, true},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
task := &upsertTask{
|
||||
ctx: context.Background(),
|
||||
schema: schema,
|
||||
req: &milvuspb.UpsertRequest{
|
||||
FieldsData: upsertData,
|
||||
NumRows: uint32(numRows),
|
||||
PartialUpdate: true,
|
||||
FieldOps: []*schemapb.FieldPartialUpdateOp{op("tags", schemapb.FieldPartialUpdateOp_ARRAY_APPEND)},
|
||||
},
|
||||
upsertMsg: &msgstream.UpsertMsg{
|
||||
InsertMsg: &msgstream.InsertMsg{
|
||||
InsertRequest: &msgpb.InsertRequest{
|
||||
FieldsData: upsertData,
|
||||
NumRows: numRows,
|
||||
Version: msgpb.InsertDataVersion_ColumnBased,
|
||||
},
|
||||
},
|
||||
},
|
||||
node: &Proxy{},
|
||||
}
|
||||
|
||||
mockRetrieve := mockey.Mock(retrieveByPKs).Return(mockQueryResult, segcore.StorageCost{}, nil).Build()
|
||||
defer mockRetrieve.UnPatch()
|
||||
|
||||
err := task.queryPreExecute(context.Background())
|
||||
assert.NoError(t, err)
|
||||
|
||||
var tagsField *schemapb.FieldData
|
||||
var noteField *schemapb.FieldData
|
||||
for _, f := range task.insertFieldData {
|
||||
switch f.GetFieldName() {
|
||||
case "tags":
|
||||
tagsField = f
|
||||
case "note":
|
||||
noteField = f
|
||||
}
|
||||
}
|
||||
assert.NotNil(t, tagsField)
|
||||
assert.Equal(t, []bool{true, true}, tagsField.ValidData)
|
||||
tagRows := tagsField.GetScalars().GetArrayData().GetData()
|
||||
assert.Equal(t, []int64{10, 20, 3}, tagRows[0].GetLongData().GetData())
|
||||
assert.Equal(t, []int64{100, 200}, tagRows[1].GetLongData().GetData())
|
||||
|
||||
assert.NotNil(t, noteField)
|
||||
assert.Equal(t, []bool{true, false}, noteField.ValidData)
|
||||
assert.Equal(t, []string{"new1"}, noteField.GetScalars().GetStringData().GetData())
|
||||
}
|
||||
|
||||
// Test ToCompressedFormatNullable for Geometry and Timestamptz types
|
||||
func TestToCompressedFormatNullable_GeometryAndTimestamptz(t *testing.T) {
|
||||
t.Run("timestamptz with null values", func(t *testing.T) {
|
||||
|
||||
@@ -54,6 +54,7 @@ func TestFencedError(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestWAL(t *testing.T) {
|
||||
walimplstest.Reset()
|
||||
initResourceForTest(t)
|
||||
b := registry.MustGetBuilder(message.WALNameTest,
|
||||
redo.NewInterceptorBuilder(),
|
||||
@@ -178,9 +179,9 @@ func (f *testOneWALFramework) Run() {
|
||||
MustBuildMutable()
|
||||
|
||||
result, err := rwWAL.Append(ctx, createMsg)
|
||||
walimplstest.DisableFenced(pChannel.Name)
|
||||
require.Nil(f.t, result)
|
||||
require.True(f.t, status.AsStreamingError(err).IsFenced())
|
||||
walimplstest.DisableFenced(pChannel.Name)
|
||||
rwWAL.Close()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -24,6 +24,13 @@ var (
|
||||
enableFenceError = atomic.NewBool(true)
|
||||
)
|
||||
|
||||
// Reset clears global state of the in-memory WAL test implementation.
|
||||
func Reset() {
|
||||
logs = typeutil.NewConcurrentMap[string, *messageLog]()
|
||||
fenced = typeutil.NewConcurrentSet[string]()
|
||||
enableFenceError.Store(true)
|
||||
}
|
||||
|
||||
// EnableFenced enables fenced mode for the given channel.
|
||||
func EnableFenced(channel string) {
|
||||
fenced.Insert(channel)
|
||||
@@ -50,7 +57,7 @@ func (w *walImpls) Append(ctx context.Context, msg message.MutableMessage) (mess
|
||||
if fenced.Contain(w.Channel().Name) {
|
||||
return nil, errors.Mark(errors.New("err"), walimpls.ErrFenced)
|
||||
}
|
||||
if enableFenceError.Load() && rand.Int31n(30) == 0 {
|
||||
if enableFenceError.Load() && msg.MessageType() != message.MessageTypeTimeTick && rand.Int31n(30) == 0 {
|
||||
return nil, errors.New("random error")
|
||||
}
|
||||
return w.datas.Append(ctx, msg)
|
||||
|
||||
@@ -1388,6 +1388,73 @@ func UpdateFieldData(base, update []*schemapb.FieldData, baseIdx, updateIdx int6
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateArrayFieldByColumnWithOp merges an Array field's update rows into
|
||||
// the base rows while applying a FieldPartialUpdateOp. Non-Array field
|
||||
// types or a REPLACE op fall back to UpdateFieldDataByColumn's behavior.
|
||||
//
|
||||
// maxCapacity caps ARRAY_APPEND's post-merge length; pass -1 to skip the
|
||||
// check (proxy should pass the schema-declared max_capacity).
|
||||
func UpdateArrayFieldByColumnWithOp(
|
||||
base, update *schemapb.FieldData,
|
||||
baseIndices, updateIndices []int64,
|
||||
op schemapb.FieldPartialUpdateOp_OpType,
|
||||
maxCapacity int,
|
||||
) error {
|
||||
if op == schemapb.FieldPartialUpdateOp_REPLACE {
|
||||
// 2.6 has no column-level UpdateFieldDataByColumn; emulate REPLACE
|
||||
// by writing update rows into base row-by-row via UpdateFieldData.
|
||||
if base == nil || update == nil {
|
||||
return nil
|
||||
}
|
||||
if len(baseIndices) != len(updateIndices) {
|
||||
return fmt.Errorf("baseIndices and updateIndices length mismatch: %d vs %d", len(baseIndices), len(updateIndices))
|
||||
}
|
||||
for i := range baseIndices {
|
||||
if err := UpdateFieldData([]*schemapb.FieldData{base}, []*schemapb.FieldData{update}, baseIndices[i], updateIndices[i]); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if base == nil || update == nil || len(baseIndices) == 0 {
|
||||
return nil
|
||||
}
|
||||
if len(baseIndices) != len(updateIndices) {
|
||||
return fmt.Errorf("baseIndices and updateIndices length mismatch: %d vs %d", len(baseIndices), len(updateIndices))
|
||||
}
|
||||
if base.GetType() != schemapb.DataType_Array {
|
||||
return fmt.Errorf("op %s requires Array field, got %s", op.String(), base.GetType().String())
|
||||
}
|
||||
baseScalar := base.GetScalars()
|
||||
updateScalar := update.GetScalars()
|
||||
if baseScalar.GetArrayData() == nil || updateScalar.GetArrayData() == nil {
|
||||
return fmt.Errorf("op %s requires non-nil ArrayData on both base and update", op.String())
|
||||
}
|
||||
baseData := baseScalar.GetArrayData().Data
|
||||
updateData := updateScalar.GetArrayData().Data
|
||||
elementType := baseScalar.GetArrayData().GetElementType()
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
// If the upsert payload row is explicitly null, there is nothing to
|
||||
// append/remove; leave the existing base row untouched.
|
||||
if len(update.ValidData) > 0 && !update.ValidData[updateIdx] {
|
||||
continue
|
||||
}
|
||||
merged, err := ApplyArrayRowOp(baseData[baseIdx], updateData[updateIdx], op, elementType, maxCapacity)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
baseData[baseIdx] = merged
|
||||
// After a successful merge the base row carries concrete data and
|
||||
// must be marked valid, otherwise downstream readers keep treating
|
||||
// it as null and drop the merged payload silently.
|
||||
if len(base.ValidData) > 0 {
|
||||
base.ValidData[baseIdx] = true
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MergeFieldData appends fields data to dst
|
||||
func MergeFieldData(dst []*schemapb.FieldData, src []*schemapb.FieldData) error {
|
||||
fieldID2Data := make(map[int64]*schemapb.FieldData)
|
||||
@@ -2547,3 +2614,233 @@ func IsBm25FunctionInputField(coll *schemapb.CollectionSchema, field *schemapb.F
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ApplyArrayRowOp applies a FieldPartialUpdateOp to a single Array-field row.
|
||||
//
|
||||
// base and update are per-row ScalarField values (as stored in
|
||||
// ArrayArray.Data[i]). elementType is the Array field's declared element type,
|
||||
// used to dispatch the concrete typed-array handling. maxCapacity caps the
|
||||
// resulting array length for ARRAY_APPEND; pass -1 to skip the check (the
|
||||
// proxy is expected to enforce the schema-declared max_capacity).
|
||||
//
|
||||
// Returned ScalarField is a newly constructed value; base/update are not
|
||||
// mutated. Nil base or update is treated as an empty row for that side.
|
||||
func ApplyArrayRowOp(
|
||||
base, update *schemapb.ScalarField,
|
||||
op schemapb.FieldPartialUpdateOp_OpType,
|
||||
elementType schemapb.DataType,
|
||||
maxCapacity int,
|
||||
) (*schemapb.ScalarField, error) {
|
||||
switch op {
|
||||
case schemapb.FieldPartialUpdateOp_REPLACE:
|
||||
return update, nil
|
||||
case schemapb.FieldPartialUpdateOp_ARRAY_APPEND:
|
||||
return appendArrayRow(base, update, elementType, maxCapacity)
|
||||
case schemapb.FieldPartialUpdateOp_ARRAY_REMOVE:
|
||||
return removeArrayRow(base, update, elementType)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported FieldPartialUpdateOp: %s", op.String())
|
||||
}
|
||||
}
|
||||
|
||||
func appendArrayRow(
|
||||
base, update *schemapb.ScalarField,
|
||||
elementType schemapb.DataType,
|
||||
maxCapacity int,
|
||||
) (*schemapb.ScalarField, error) {
|
||||
switch elementType {
|
||||
case schemapb.DataType_Bool:
|
||||
b := base.GetBoolData().GetData()
|
||||
u := update.GetBoolData().GetData()
|
||||
merged := make([]bool, 0, len(b)+len(u))
|
||||
merged = append(merged, b...)
|
||||
merged = append(merged, u...)
|
||||
if maxCapacity >= 0 && len(merged) > maxCapacity {
|
||||
return nil, newArrayCapacityError(len(merged), maxCapacity)
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: merged}}}, nil
|
||||
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
|
||||
b := base.GetIntData().GetData()
|
||||
u := update.GetIntData().GetData()
|
||||
merged := make([]int32, 0, len(b)+len(u))
|
||||
merged = append(merged, b...)
|
||||
merged = append(merged, u...)
|
||||
if maxCapacity >= 0 && len(merged) > maxCapacity {
|
||||
return nil, newArrayCapacityError(len(merged), maxCapacity)
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: merged}}}, nil
|
||||
case schemapb.DataType_Int64:
|
||||
b := base.GetLongData().GetData()
|
||||
u := update.GetLongData().GetData()
|
||||
merged := make([]int64, 0, len(b)+len(u))
|
||||
merged = append(merged, b...)
|
||||
merged = append(merged, u...)
|
||||
if maxCapacity >= 0 && len(merged) > maxCapacity {
|
||||
return nil, newArrayCapacityError(len(merged), maxCapacity)
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: merged}}}, nil
|
||||
case schemapb.DataType_Float:
|
||||
b := base.GetFloatData().GetData()
|
||||
u := update.GetFloatData().GetData()
|
||||
merged := make([]float32, 0, len(b)+len(u))
|
||||
merged = append(merged, b...)
|
||||
merged = append(merged, u...)
|
||||
if maxCapacity >= 0 && len(merged) > maxCapacity {
|
||||
return nil, newArrayCapacityError(len(merged), maxCapacity)
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: merged}}}, nil
|
||||
case schemapb.DataType_Double:
|
||||
b := base.GetDoubleData().GetData()
|
||||
u := update.GetDoubleData().GetData()
|
||||
merged := make([]float64, 0, len(b)+len(u))
|
||||
merged = append(merged, b...)
|
||||
merged = append(merged, u...)
|
||||
if maxCapacity >= 0 && len(merged) > maxCapacity {
|
||||
return nil, newArrayCapacityError(len(merged), maxCapacity)
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: merged}}}, nil
|
||||
case schemapb.DataType_VarChar, schemapb.DataType_String:
|
||||
b := base.GetStringData().GetData()
|
||||
u := update.GetStringData().GetData()
|
||||
merged := make([]string, 0, len(b)+len(u))
|
||||
merged = append(merged, b...)
|
||||
merged = append(merged, u...)
|
||||
if maxCapacity >= 0 && len(merged) > maxCapacity {
|
||||
return nil, newArrayCapacityError(len(merged), maxCapacity)
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: merged}}}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("ARRAY_APPEND does not support element type: %s", elementType.String())
|
||||
}
|
||||
}
|
||||
|
||||
func removeArrayRow(
|
||||
base, update *schemapb.ScalarField,
|
||||
elementType schemapb.DataType,
|
||||
) (*schemapb.ScalarField, error) {
|
||||
switch elementType {
|
||||
case schemapb.DataType_Bool:
|
||||
b := base.GetBoolData().GetData()
|
||||
u := update.GetBoolData().GetData()
|
||||
if len(b) == 0 || len(u) == 0 {
|
||||
return base, nil
|
||||
}
|
||||
remove := make(map[bool]struct{}, len(u))
|
||||
for _, v := range u {
|
||||
remove[v] = struct{}{}
|
||||
}
|
||||
out := make([]bool, 0, len(b))
|
||||
for _, v := range b {
|
||||
if _, skip := remove[v]; !skip {
|
||||
out = append(out, v)
|
||||
}
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: out}}}, nil
|
||||
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
|
||||
b := base.GetIntData().GetData()
|
||||
u := update.GetIntData().GetData()
|
||||
if len(b) == 0 || len(u) == 0 {
|
||||
return base, nil
|
||||
}
|
||||
remove := make(map[int32]struct{}, len(u))
|
||||
for _, v := range u {
|
||||
remove[v] = struct{}{}
|
||||
}
|
||||
out := make([]int32, 0, len(b))
|
||||
for _, v := range b {
|
||||
if _, skip := remove[v]; !skip {
|
||||
out = append(out, v)
|
||||
}
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: out}}}, nil
|
||||
case schemapb.DataType_Int64:
|
||||
b := base.GetLongData().GetData()
|
||||
u := update.GetLongData().GetData()
|
||||
if len(b) == 0 || len(u) == 0 {
|
||||
return base, nil
|
||||
}
|
||||
remove := make(map[int64]struct{}, len(u))
|
||||
for _, v := range u {
|
||||
remove[v] = struct{}{}
|
||||
}
|
||||
out := make([]int64, 0, len(b))
|
||||
for _, v := range b {
|
||||
if _, skip := remove[v]; !skip {
|
||||
out = append(out, v)
|
||||
}
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: out}}}, nil
|
||||
case schemapb.DataType_Float:
|
||||
// Linear scan: float32 equality is position-wise exact; hashing is
|
||||
// valid but yields no speedup in typical cardinalities.
|
||||
b := base.GetFloatData().GetData()
|
||||
u := update.GetFloatData().GetData()
|
||||
if len(b) == 0 || len(u) == 0 {
|
||||
return base, nil
|
||||
}
|
||||
out := make([]float32, 0, len(b))
|
||||
for _, v := range b {
|
||||
if !containsFloat32(u, v) {
|
||||
out = append(out, v)
|
||||
}
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: out}}}, nil
|
||||
case schemapb.DataType_Double:
|
||||
b := base.GetDoubleData().GetData()
|
||||
u := update.GetDoubleData().GetData()
|
||||
if len(b) == 0 || len(u) == 0 {
|
||||
return base, nil
|
||||
}
|
||||
out := make([]float64, 0, len(b))
|
||||
for _, v := range b {
|
||||
if !containsFloat64(u, v) {
|
||||
out = append(out, v)
|
||||
}
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: out}}}, nil
|
||||
case schemapb.DataType_VarChar, schemapb.DataType_String:
|
||||
b := base.GetStringData().GetData()
|
||||
u := update.GetStringData().GetData()
|
||||
if len(b) == 0 || len(u) == 0 {
|
||||
return base, nil
|
||||
}
|
||||
remove := make(map[string]struct{}, len(u))
|
||||
for _, v := range u {
|
||||
remove[v] = struct{}{}
|
||||
}
|
||||
out := make([]string, 0, len(b))
|
||||
for _, v := range b {
|
||||
if _, skip := remove[v]; !skip {
|
||||
out = append(out, v)
|
||||
}
|
||||
}
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: out}}}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("ARRAY_REMOVE does not support element type: %s", elementType.String())
|
||||
}
|
||||
}
|
||||
|
||||
func containsFloat32(haystack []float32, needle float32) bool {
|
||||
// NaN is intentionally treated as never equal to any value, matching
|
||||
// IEEE-754 semantics — an ARRAY_REMOVE request targeting NaN cannot
|
||||
// delete NaN elements from the base.
|
||||
for _, v := range haystack {
|
||||
if v == needle {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func containsFloat64(haystack []float64, needle float64) bool {
|
||||
for _, v := range haystack {
|
||||
if v == needle {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func newArrayCapacityError(got, max int) error {
|
||||
return fmt.Errorf("array length %d exceeds max_capacity %d", got, max)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,467 @@
|
||||
// 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 typeutil
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
)
|
||||
|
||||
// boolRow/intRow/... helpers build a ScalarField carrying a single row's
|
||||
// typed-array payload. They keep tests concise and match how ArrayArray
|
||||
// stores per-row data.
|
||||
|
||||
func boolRow(vs ...bool) *schemapb.ScalarField {
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: vs}}}
|
||||
}
|
||||
|
||||
func intRow(vs ...int32) *schemapb.ScalarField {
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: vs}}}
|
||||
}
|
||||
|
||||
func longRow(vs ...int64) *schemapb.ScalarField {
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: vs}}}
|
||||
}
|
||||
|
||||
func floatRow(vs ...float32) *schemapb.ScalarField {
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: vs}}}
|
||||
}
|
||||
|
||||
func doubleRow(vs ...float64) *schemapb.ScalarField {
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: vs}}}
|
||||
}
|
||||
|
||||
func stringRow(vs ...string) *schemapb.ScalarField {
|
||||
return &schemapb.ScalarField{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: vs}}}
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_Replace(t *testing.T) {
|
||||
base := longRow(1, 2, 3)
|
||||
update := longRow(9, 8)
|
||||
got, err := ApplyArrayRowOp(base, update, schemapb.FieldPartialUpdateOp_REPLACE, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{9, 8}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_UnsupportedOp(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1), longRow(2), schemapb.FieldPartialUpdateOp_OpType(999), schemapb.DataType_Int64, -1)
|
||||
assert.Error(t, err)
|
||||
assert.Nil(t, got)
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendBool(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(boolRow(true, false), boolRow(true), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Bool, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []bool{true, false, true}, got.GetBoolData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendInt32(t *testing.T) {
|
||||
for _, et := range []schemapb.DataType{schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32} {
|
||||
got, err := ApplyArrayRowOp(intRow(1, 2), intRow(3, 4), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, et, -1)
|
||||
require.NoError(t, err, et.String())
|
||||
assert.Equal(t, []int32{1, 2, 3, 4}, got.GetIntData().GetData(), et.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendInt64(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2), longRow(3, 4), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 2, 3, 4}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendFloat(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(floatRow(1.5, 2.5), floatRow(3.5), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Float, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []float32{1.5, 2.5, 3.5}, got.GetFloatData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendDouble(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(doubleRow(1.25), doubleRow(2.5, 3.5), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Double, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []float64{1.25, 2.5, 3.5}, got.GetDoubleData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendVarChar(t *testing.T) {
|
||||
for _, et := range []schemapb.DataType{schemapb.DataType_VarChar, schemapb.DataType_String} {
|
||||
got, err := ApplyArrayRowOp(stringRow("a"), stringRow("b", "c"), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, et, -1)
|
||||
require.NoError(t, err, et.String())
|
||||
assert.Equal(t, []string{"a", "b", "c"}, got.GetStringData().GetData(), et.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendEmptyBase(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(), longRow(1, 2), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 2}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendEmptyUpdate(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2), longRow(), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 2}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendBothEmpty(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(), longRow(), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendOverflowsCapacity(t *testing.T) {
|
||||
_, err := ApplyArrayRowOp(longRow(1, 2), longRow(3), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Int64, 2)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "max_capacity")
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendCapacityExactlyMet(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2), longRow(3), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Int64, 3)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 2, 3}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendNegativeCapacitySkipsCheck(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2, 3, 4, 5), longRow(6), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, got.GetLongData().GetData(), 6)
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendCapacityBoolInt32FloatDoubleString(t *testing.T) {
|
||||
cases := []struct {
|
||||
et schemapb.DataType
|
||||
base *schemapb.ScalarField
|
||||
upd *schemapb.ScalarField
|
||||
}{
|
||||
{schemapb.DataType_Bool, boolRow(true), boolRow(false)},
|
||||
{schemapb.DataType_Int32, intRow(1), intRow(2)},
|
||||
{schemapb.DataType_Float, floatRow(1), floatRow(2)},
|
||||
{schemapb.DataType_Double, doubleRow(1), doubleRow(2)},
|
||||
{schemapb.DataType_VarChar, stringRow("a"), stringRow("b")},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
_, err := ApplyArrayRowOp(tc.base, tc.upd, schemapb.FieldPartialUpdateOp_ARRAY_APPEND, tc.et, 1)
|
||||
assert.Error(t, err, tc.et.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_AppendUnsupportedElementType(t *testing.T) {
|
||||
_, err := ApplyArrayRowOp(longRow(), longRow(), schemapb.FieldPartialUpdateOp_ARRAY_APPEND, schemapb.DataType_JSON, -1)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveBool(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(boolRow(true, false, true, false), boolRow(true), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Bool, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []bool{false, false}, got.GetBoolData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveInt32(t *testing.T) {
|
||||
for _, et := range []schemapb.DataType{schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32} {
|
||||
got, err := ApplyArrayRowOp(intRow(1, 2, 3, 2, 1), intRow(2), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, et, -1)
|
||||
require.NoError(t, err, et.String())
|
||||
assert.Equal(t, []int32{1, 3, 1}, got.GetIntData().GetData(), et.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveInt64(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2, 3, 4), longRow(2, 4), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 3}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveFloat(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(floatRow(1.5, 2.5, 3.5, 2.5), floatRow(2.5), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Float, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []float32{1.5, 3.5}, got.GetFloatData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveDouble(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(doubleRow(1.1, 2.2, 3.3), doubleRow(2.2), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Double, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []float64{1.1, 3.3}, got.GetDoubleData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveVarChar(t *testing.T) {
|
||||
for _, et := range []schemapb.DataType{schemapb.DataType_VarChar, schemapb.DataType_String} {
|
||||
got, err := ApplyArrayRowOp(stringRow("a", "b", "a", "c"), stringRow("a"), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, et, -1)
|
||||
require.NoError(t, err, et.String())
|
||||
assert.Equal(t, []string{"b", "c"}, got.GetStringData().GetData(), et.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveMultipleValues(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2, 3, 4, 5), longRow(2, 4), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 3, 5}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveNoMatch(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2, 3), longRow(99), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 2, 3}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveEmptyBase(t *testing.T) {
|
||||
base := longRow()
|
||||
got, err := ApplyArrayRowOp(base, longRow(1, 2), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Same(t, base, got)
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveEmptyUpdate(t *testing.T) {
|
||||
base := longRow(1, 2, 3)
|
||||
got, err := ApplyArrayRowOp(base, longRow(), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Same(t, base, got)
|
||||
}
|
||||
|
||||
// TestApplyArrayRowOp_RemoveEmptyAcrossTypes exercises the early-exit
|
||||
// "empty base or update" branches for every supported element type.
|
||||
func TestApplyArrayRowOp_RemoveEmptyAcrossTypes(t *testing.T) {
|
||||
type rowCase struct {
|
||||
et schemapb.DataType
|
||||
emptyBase *schemapb.ScalarField
|
||||
emptyUpd *schemapb.ScalarField
|
||||
valuedBase *schemapb.ScalarField
|
||||
valuedUpd *schemapb.ScalarField
|
||||
}
|
||||
cases := []rowCase{
|
||||
{schemapb.DataType_Bool, boolRow(), boolRow(), boolRow(true), boolRow(true)},
|
||||
{schemapb.DataType_Int8, intRow(), intRow(), intRow(1), intRow(1)},
|
||||
{schemapb.DataType_Int16, intRow(), intRow(), intRow(1), intRow(1)},
|
||||
{schemapb.DataType_Int32, intRow(), intRow(), intRow(1), intRow(1)},
|
||||
{schemapb.DataType_Float, floatRow(), floatRow(), floatRow(1), floatRow(1)},
|
||||
{schemapb.DataType_Double, doubleRow(), doubleRow(), doubleRow(1), doubleRow(1)},
|
||||
{schemapb.DataType_VarChar, stringRow(), stringRow(), stringRow("a"), stringRow("a")},
|
||||
{schemapb.DataType_String, stringRow(), stringRow(), stringRow("a"), stringRow("a")},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
// empty base
|
||||
got, err := ApplyArrayRowOp(tc.emptyBase, tc.valuedUpd, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, tc.et, -1)
|
||||
require.NoError(t, err, tc.et.String())
|
||||
assert.Same(t, tc.emptyBase, got, tc.et.String())
|
||||
// empty update
|
||||
got, err = ApplyArrayRowOp(tc.valuedBase, tc.emptyUpd, schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, tc.et, -1)
|
||||
require.NoError(t, err, tc.et.String())
|
||||
assert.Same(t, tc.valuedBase, got, tc.et.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveDuplicatesInUpdate(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(longRow(1, 2, 3, 2), longRow(2, 2, 2), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Int64, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{1, 3}, got.GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveFloatNaN(t *testing.T) {
|
||||
nan32 := float32(math.NaN())
|
||||
got, err := ApplyArrayRowOp(floatRow(1.0, nan32, 2.0), floatRow(nan32), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Float, -1)
|
||||
require.NoError(t, err)
|
||||
// NaN != NaN → base NaN retained
|
||||
out := got.GetFloatData().GetData()
|
||||
require.Len(t, out, 3)
|
||||
assert.Equal(t, float32(1.0), out[0])
|
||||
assert.True(t, math.IsNaN(float64(out[1])))
|
||||
assert.Equal(t, float32(2.0), out[2])
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveDoubleNaN(t *testing.T) {
|
||||
got, err := ApplyArrayRowOp(doubleRow(math.NaN(), 1.0), doubleRow(math.NaN()), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_Double, -1)
|
||||
require.NoError(t, err)
|
||||
out := got.GetDoubleData().GetData()
|
||||
require.Len(t, out, 2)
|
||||
assert.True(t, math.IsNaN(out[0]))
|
||||
assert.Equal(t, 1.0, out[1])
|
||||
}
|
||||
|
||||
func TestApplyArrayRowOp_RemoveUnsupportedElementType(t *testing.T) {
|
||||
_, err := ApplyArrayRowOp(longRow(), longRow(), schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, schemapb.DataType_JSON, -1)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestContainsFloat32(t *testing.T) {
|
||||
assert.True(t, containsFloat32([]float32{1, 2, 3}, 2))
|
||||
assert.False(t, containsFloat32([]float32{1, 2, 3}, 9))
|
||||
assert.False(t, containsFloat32(nil, 1))
|
||||
// NaN behavior
|
||||
nan := float32(math.NaN())
|
||||
assert.False(t, containsFloat32([]float32{nan}, nan))
|
||||
}
|
||||
|
||||
func TestContainsFloat64(t *testing.T) {
|
||||
assert.True(t, containsFloat64([]float64{1, 2, 3}, 2))
|
||||
assert.False(t, containsFloat64([]float64{1, 2, 3}, 9))
|
||||
assert.False(t, containsFloat64(nil, 1))
|
||||
assert.False(t, containsFloat64([]float64{math.NaN()}, math.NaN()))
|
||||
}
|
||||
|
||||
func arrayField(rows []*schemapb.ScalarField, et schemapb.DataType) *schemapb.FieldData {
|
||||
return &schemapb.FieldData{
|
||||
Type: schemapb.DataType_Array,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
|
||||
Data: rows,
|
||||
ElementType: et,
|
||||
}},
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateArrayFieldByColumnWithOp_Append(t *testing.T) {
|
||||
base := arrayField([]*schemapb.ScalarField{longRow(1), longRow(10, 20)}, schemapb.DataType_Int64)
|
||||
update := arrayField([]*schemapb.ScalarField{longRow(2, 3), longRow(30)}, schemapb.DataType_Int64)
|
||||
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0, 1}, []int64{0, 1},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1)
|
||||
require.NoError(t, err)
|
||||
|
||||
rows := base.GetScalars().GetArrayData().GetData()
|
||||
assert.Equal(t, []int64{1, 2, 3}, rows[0].GetLongData().GetData())
|
||||
assert.Equal(t, []int64{10, 20, 30}, rows[1].GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestUpdateArrayFieldByColumnWithOp_Remove(t *testing.T) {
|
||||
base := arrayField([]*schemapb.ScalarField{longRow(1, 2, 3, 2), longRow(10)}, schemapb.DataType_Int64)
|
||||
update := arrayField([]*schemapb.ScalarField{longRow(2), longRow(99)}, schemapb.DataType_Int64)
|
||||
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0, 1}, []int64{0, 1},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_REMOVE, -1)
|
||||
require.NoError(t, err)
|
||||
|
||||
rows := base.GetScalars().GetArrayData().GetData()
|
||||
assert.Equal(t, []int64{1, 3}, rows[0].GetLongData().GetData())
|
||||
assert.Equal(t, []int64{10}, rows[1].GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestUpdateArrayFieldByColumnWithOp_ReplaceFallback(t *testing.T) {
|
||||
base := arrayField([]*schemapb.ScalarField{longRow(1, 2)}, schemapb.DataType_Int64)
|
||||
update := arrayField([]*schemapb.ScalarField{longRow(9)}, schemapb.DataType_Int64)
|
||||
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0}, []int64{0},
|
||||
schemapb.FieldPartialUpdateOp_REPLACE, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{9}, base.GetScalars().GetArrayData().GetData()[0].GetLongData().GetData())
|
||||
}
|
||||
|
||||
func TestUpdateArrayFieldByColumnWithOp_RejectsNonArray(t *testing.T) {
|
||||
base := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_Int64,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1}}},
|
||||
}},
|
||||
}
|
||||
update := arrayField([]*schemapb.ScalarField{longRow(2)}, schemapb.DataType_Int64)
|
||||
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0}, []int64{0},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "Array field")
|
||||
}
|
||||
|
||||
func TestUpdateArrayFieldByColumnWithOp_RejectsIndexLenMismatch(t *testing.T) {
|
||||
base := arrayField([]*schemapb.ScalarField{longRow()}, schemapb.DataType_Int64)
|
||||
update := arrayField([]*schemapb.ScalarField{longRow()}, schemapb.DataType_Int64)
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0}, []int64{0, 1},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "length mismatch")
|
||||
}
|
||||
|
||||
func TestUpdateArrayFieldByColumnWithOp_NilInputsReturnNil(t *testing.T) {
|
||||
assert.NoError(t, UpdateArrayFieldByColumnWithOp(nil, nil, nil, nil,
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1))
|
||||
base := arrayField([]*schemapb.ScalarField{longRow()}, schemapb.DataType_Int64)
|
||||
assert.NoError(t, UpdateArrayFieldByColumnWithOp(base, nil, nil, nil,
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1))
|
||||
}
|
||||
|
||||
func TestUpdateArrayFieldByColumnWithOp_PropagatesMergeError(t *testing.T) {
|
||||
// Capacity exceeded on APPEND triggers ApplyArrayRowOp's error path.
|
||||
base := arrayField([]*schemapb.ScalarField{longRow(1, 2)}, schemapb.DataType_Int64)
|
||||
update := arrayField([]*schemapb.ScalarField{longRow(3)}, schemapb.DataType_Int64)
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0}, []int64{0},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, 2)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "max_capacity")
|
||||
}
|
||||
|
||||
func TestNewArrayCapacityError(t *testing.T) {
|
||||
err := newArrayCapacityError(10, 5)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "10")
|
||||
assert.Contains(t, err.Error(), "5")
|
||||
}
|
||||
|
||||
// Regression: a successful APPEND / REMOVE on a previously-null base row
|
||||
// must flip ValidData back to true, otherwise downstream readers keep the
|
||||
// row as null and silently drop the merged payload.
|
||||
func TestUpdateArrayFieldByColumnWithOp_FlipsValidDataOnMerge(t *testing.T) {
|
||||
base := arrayField([]*schemapb.ScalarField{longRow(), longRow(1, 2)}, schemapb.DataType_Int64)
|
||||
base.ValidData = []bool{false, true}
|
||||
update := arrayField([]*schemapb.ScalarField{longRow(5, 6), longRow(3)}, schemapb.DataType_Int64)
|
||||
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0, 1}, []int64{0, 1},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []bool{true, true}, base.ValidData)
|
||||
rows := base.GetScalars().GetArrayData().GetData()
|
||||
assert.Equal(t, []int64{5, 6}, rows[0].GetLongData().GetData())
|
||||
assert.Equal(t, []int64{1, 2, 3}, rows[1].GetLongData().GetData())
|
||||
}
|
||||
|
||||
// Regression: a null upsert payload row (update.ValidData[i]=false) must
|
||||
// leave the base row and its ValidData bit untouched.
|
||||
func TestUpdateArrayFieldByColumnWithOp_SkipsNullUpdatePayload(t *testing.T) {
|
||||
base := arrayField([]*schemapb.ScalarField{longRow(1, 2)}, schemapb.DataType_Int64)
|
||||
base.ValidData = []bool{true}
|
||||
update := arrayField([]*schemapb.ScalarField{longRow(9)}, schemapb.DataType_Int64)
|
||||
update.ValidData = []bool{false}
|
||||
|
||||
err := UpdateArrayFieldByColumnWithOp(base, update, []int64{0}, []int64{0},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1)
|
||||
require.NoError(t, err)
|
||||
rows := base.GetScalars().GetArrayData().GetData()
|
||||
assert.Equal(t, []int64{1, 2}, rows[0].GetLongData().GetData())
|
||||
assert.Equal(t, []bool{true}, base.ValidData)
|
||||
}
|
||||
|
||||
// Regression: a malformed FieldData (type=Array but ArrayData=nil) must
|
||||
// be rejected deterministically instead of panicking inside the merge.
|
||||
func TestUpdateArrayFieldByColumnWithOp_RejectsNilArrayData(t *testing.T) {
|
||||
base := arrayField([]*schemapb.ScalarField{longRow(1)}, schemapb.DataType_Int64)
|
||||
|
||||
updateNoArray := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_Array,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{}},
|
||||
}
|
||||
err := UpdateArrayFieldByColumnWithOp(base, updateNoArray, []int64{0}, []int64{0},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1)
|
||||
require.Error(t, err)
|
||||
|
||||
baseNoArray := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_Array,
|
||||
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{}},
|
||||
}
|
||||
good := arrayField([]*schemapb.ScalarField{longRow(1)}, schemapb.DataType_Int64)
|
||||
err = UpdateArrayFieldByColumnWithOp(baseNoArray, good, []int64{0}, []int64{0},
|
||||
schemapb.FieldPartialUpdateOp_ARRAY_APPEND, -1)
|
||||
require.Error(t, err)
|
||||
}
|
||||
@@ -3,6 +3,7 @@ module github.com/milvus-io/milvus/tests/go_client
|
||||
go 1.25.8
|
||||
|
||||
require (
|
||||
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.16
|
||||
github.com/milvus-io/milvus/client/v2 v2.0.0-20241125024034-0b9edb62a92d
|
||||
github.com/milvus-io/milvus/pkg/v2 v2.6.7-0.20251202033909-b71a123d25ad
|
||||
github.com/peterstace/simplefeatures v0.54.0
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
// 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 testcases
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/client/v2/column"
|
||||
"github.com/milvus-io/milvus/client/v2/entity"
|
||||
"github.com/milvus-io/milvus/client/v2/index"
|
||||
client "github.com/milvus-io/milvus/client/v2/milvusclient"
|
||||
"github.com/milvus-io/milvus/tests/go_client/base"
|
||||
"github.com/milvus-io/milvus/tests/go_client/common"
|
||||
hp "github.com/milvus-io/milvus/tests/go_client/testcases/helper"
|
||||
)
|
||||
|
||||
// End-to-end tests for ARRAY_APPEND / ARRAY_REMOVE partial-update
|
||||
// operators on Array fields. Unlike the in-process integration test,
|
||||
// these run against a live Milvus deployment through the Go SDK.
|
||||
|
||||
const (
|
||||
arrayPartialOpDim = 4
|
||||
arrayPartialOpCapacity = 16
|
||||
arrayPartialOpTagsFld = "tags"
|
||||
)
|
||||
|
||||
// setupArrayPartialOpCollection creates a minimal collection
|
||||
// (pk int64, vec float-vec dim=4, tags Array<Int64> max_capacity=16),
|
||||
// inserts seed rows (one per element in seeds), then indexes + loads
|
||||
// the collection so it is queryable.
|
||||
func setupArrayPartialOpCollection(
|
||||
ctx context.Context, t *testing.T, mc *base.MilvusClient,
|
||||
collName string, seeds [][]int64,
|
||||
) *entity.Schema {
|
||||
schema := entity.NewSchema().WithName(collName).
|
||||
WithField(entity.NewField().
|
||||
WithName(common.DefaultInt64FieldName).
|
||||
WithDataType(entity.FieldTypeInt64).
|
||||
WithIsPrimaryKey(true)).
|
||||
WithField(entity.NewField().
|
||||
WithName(common.DefaultFloatVecFieldName).
|
||||
WithDataType(entity.FieldTypeFloatVector).
|
||||
WithDim(arrayPartialOpDim)).
|
||||
WithField(entity.NewField().
|
||||
WithName(arrayPartialOpTagsFld).
|
||||
WithDataType(entity.FieldTypeArray).
|
||||
WithElementType(entity.FieldTypeInt64).
|
||||
WithMaxCapacity(arrayPartialOpCapacity))
|
||||
|
||||
err := mc.CreateCollection(ctx, client.NewCreateCollectionOption(collName, schema))
|
||||
common.CheckErr(t, err, true)
|
||||
|
||||
pks := make([]int64, len(seeds))
|
||||
vecs := make([][]float32, len(seeds))
|
||||
for i := range seeds {
|
||||
pks[i] = int64(i)
|
||||
row := make([]float32, arrayPartialOpDim)
|
||||
for j := 0; j < arrayPartialOpDim; j++ {
|
||||
row[j] = float32((i*arrayPartialOpDim+j)%7) / 10.0
|
||||
}
|
||||
vecs[i] = row
|
||||
}
|
||||
_, err = mc.Insert(ctx, client.NewColumnBasedInsertOption(collName).
|
||||
WithColumns(
|
||||
column.NewColumnInt64(common.DefaultInt64FieldName, pks),
|
||||
column.NewColumnFloatVector(common.DefaultFloatVecFieldName, arrayPartialOpDim, vecs),
|
||||
column.NewColumnInt64Array(arrayPartialOpTagsFld, seeds),
|
||||
))
|
||||
common.CheckErr(t, err, true)
|
||||
|
||||
_, err = mc.Flush(ctx, client.NewFlushOption(collName))
|
||||
common.CheckErr(t, err, true)
|
||||
|
||||
indexTask, err := mc.CreateIndex(ctx, client.NewCreateIndexOption(
|
||||
collName, common.DefaultFloatVecFieldName, index.NewAutoIndex(entity.COSINE)))
|
||||
common.CheckErr(t, err, true)
|
||||
require.NoError(t, indexTask.Await(ctx))
|
||||
|
||||
loadTask, err := mc.LoadCollection(ctx, client.NewLoadCollectionOption(collName))
|
||||
common.CheckErr(t, err, true)
|
||||
require.NoError(t, loadTask.Await(ctx))
|
||||
|
||||
return schema
|
||||
}
|
||||
|
||||
// queryTagsByPK returns pk → tags map for all rows with pk >= 0.
|
||||
func queryTagsByPK(
|
||||
ctx context.Context, t *testing.T, mc *base.MilvusClient, collName string,
|
||||
) map[int64][]int64 {
|
||||
resSet, err := mc.Query(ctx, client.NewQueryOption(collName).
|
||||
WithFilter(fmt.Sprintf("%s >= 0", common.DefaultInt64FieldName)).
|
||||
WithOutputFields(common.DefaultInt64FieldName, arrayPartialOpTagsFld).
|
||||
WithConsistencyLevel(entity.ClStrong))
|
||||
common.CheckErr(t, err, true)
|
||||
|
||||
var pks []int64
|
||||
var tags [][]int64
|
||||
for _, col := range resSet.Fields {
|
||||
switch col.Name() {
|
||||
case common.DefaultInt64FieldName:
|
||||
pks = col.(*column.ColumnInt64).Data()
|
||||
case arrayPartialOpTagsFld:
|
||||
tags = col.(*column.ColumnInt64Array).Data()
|
||||
}
|
||||
}
|
||||
require.Len(t, tags, len(pks))
|
||||
out := make(map[int64][]int64, len(pks))
|
||||
for i, pk := range pks {
|
||||
out[pk] = tags[i]
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// equalAsMultiset returns true when got and want contain the same
|
||||
// elements with the same multiplicity, ignoring order.
|
||||
func equalAsMultiset(got, want []int64) bool {
|
||||
if len(got) != len(want) {
|
||||
return false
|
||||
}
|
||||
a := append([]int64(nil), got...)
|
||||
b := append([]int64(nil), want...)
|
||||
sort.Slice(a, func(i, j int) bool { return a[i] < a[j] })
|
||||
sort.Slice(b, func(i, j int) bool { return b[i] < b[j] })
|
||||
for i := range a {
|
||||
if a[i] != b[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func TestArrayPartialOpAppend(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
|
||||
mc := hp.CreateDefaultMilvusClient(ctx, t)
|
||||
|
||||
collName := common.GenRandomString("array_partial_op_append_", 6)
|
||||
_ = setupArrayPartialOpCollection(ctx, t, mc,
|
||||
collName, [][]int64{{1, 2}, {10, 20}})
|
||||
|
||||
// Append: row 0 payload = [3,4]; row 1 payload = [30].
|
||||
// Expected: row 0 → [1,2,3,4], row 1 → [10,20,30].
|
||||
_, err := mc.Upsert(ctx, client.NewColumnBasedInsertOption(collName).
|
||||
WithColumns(
|
||||
column.NewColumnInt64(common.DefaultInt64FieldName, []int64{0, 1}),
|
||||
column.NewColumnInt64Array(arrayPartialOpTagsFld,
|
||||
[][]int64{{3, 4}, {30}}),
|
||||
).
|
||||
WithArrayAppend(arrayPartialOpTagsFld))
|
||||
common.CheckErr(t, err, true)
|
||||
|
||||
got := queryTagsByPK(ctx, t, mc, collName)
|
||||
require.True(t, equalAsMultiset(got[0], []int64{1, 2, 3, 4}), "row 0 got=%v", got[0])
|
||||
require.True(t, equalAsMultiset(got[1], []int64{10, 20, 30}), "row 1 got=%v", got[1])
|
||||
}
|
||||
|
||||
func TestArrayPartialOpRemove(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
|
||||
mc := hp.CreateDefaultMilvusClient(ctx, t)
|
||||
|
||||
collName := common.GenRandomString("array_partial_op_remove_", 6)
|
||||
_ = setupArrayPartialOpCollection(ctx, t, mc,
|
||||
collName, [][]int64{{1, 2, 3, 2, 1}, {10, 20, 30}})
|
||||
|
||||
// Remove: row 0 payload = [2] → [1,3,1]; row 1 payload = [40] → no-op.
|
||||
_, err := mc.Upsert(ctx, client.NewColumnBasedInsertOption(collName).
|
||||
WithColumns(
|
||||
column.NewColumnInt64(common.DefaultInt64FieldName, []int64{0, 1}),
|
||||
column.NewColumnInt64Array(arrayPartialOpTagsFld,
|
||||
[][]int64{{2}, {40}}),
|
||||
).
|
||||
WithArrayRemove(arrayPartialOpTagsFld))
|
||||
common.CheckErr(t, err, true)
|
||||
|
||||
got := queryTagsByPK(ctx, t, mc, collName)
|
||||
require.True(t, equalAsMultiset(got[0], []int64{1, 3, 1}), "row 0 got=%v", got[0])
|
||||
require.True(t, equalAsMultiset(got[1], []int64{10, 20, 30}), "row 1 got=%v", got[1])
|
||||
}
|
||||
|
||||
func TestArrayPartialOpAppendExceedsCapacity(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
|
||||
mc := hp.CreateDefaultMilvusClient(ctx, t)
|
||||
|
||||
collName := common.GenRandomString("array_partial_op_overflow_", 6)
|
||||
// Seed with length-15 base; appending 5 more exceeds capacity 16.
|
||||
base := make([]int64, 15)
|
||||
for i := range base {
|
||||
base[i] = int64(i)
|
||||
}
|
||||
payload := make([]int64, 5)
|
||||
for i := range payload {
|
||||
payload[i] = int64(100 + i)
|
||||
}
|
||||
_ = setupArrayPartialOpCollection(ctx, t, mc, collName, [][]int64{base})
|
||||
|
||||
_, err := mc.Upsert(ctx, client.NewColumnBasedInsertOption(collName).
|
||||
WithColumns(
|
||||
column.NewColumnInt64(common.DefaultInt64FieldName, []int64{0}),
|
||||
column.NewColumnInt64Array(arrayPartialOpTagsFld, [][]int64{payload}),
|
||||
).
|
||||
WithArrayAppend(arrayPartialOpTagsFld))
|
||||
common.CheckErr(t, err, false, "max_capacity")
|
||||
}
|
||||
|
||||
// TestArrayPartialOpAutoPromotesPartialUpdate asserts that a non-REPLACE
|
||||
// op alone is enough to trigger partial-update semantics: the caller
|
||||
// does not need to set partial_update=true explicitly.
|
||||
func TestArrayPartialOpAutoPromotesPartialUpdate(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
|
||||
mc := hp.CreateDefaultMilvusClient(ctx, t)
|
||||
|
||||
collName := common.GenRandomString("array_partial_op_auto_", 6)
|
||||
_ = setupArrayPartialOpCollection(ctx, t, mc, collName, [][]int64{{1, 2}})
|
||||
|
||||
// Payload intentionally omits the vector field. Without partial-update
|
||||
// semantics the server would reject the request for missing fields;
|
||||
// here the op must auto-promote partial_update=true so only the tags
|
||||
// field is merged.
|
||||
opt := client.NewColumnBasedInsertOption(collName).
|
||||
WithColumns(
|
||||
column.NewColumnInt64(common.DefaultInt64FieldName, []int64{0}),
|
||||
column.NewColumnInt64Array(arrayPartialOpTagsFld, [][]int64{{3}}),
|
||||
).
|
||||
WithFieldPartialOp(arrayPartialOpTagsFld, schemapb.FieldPartialUpdateOp_ARRAY_APPEND)
|
||||
_, err := mc.Upsert(ctx, opt)
|
||||
common.CheckErr(t, err, true)
|
||||
|
||||
got := queryTagsByPK(ctx, t, mc, collName)
|
||||
require.True(t, equalAsMultiset(got[0], []int64{1, 2, 3}), "row 0 got=%v", got[0])
|
||||
}
|
||||
Reference in New Issue
Block a user