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:
wei liu
2026-05-15 08:37:30 +08:00
committed by GitHub
co-authored by Zhen Ye
parent 8ae21f7150
commit ababab0d76
18 changed files with 2186 additions and 24 deletions
+4 -4
View File
@@ -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
+75 -2
View File
@@ -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")
}
+68 -3
View File
@@ -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",
+275
View File
@@ -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))
}
}
+103
View File
@@ -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)
+297
View File
@@ -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)
}
+467
View File
@@ -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)
}
+1
View File
@@ -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])
}