mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
fix: change queryPreExecute merge logic from row-based to column-based (#47122)
issue: #46953 - GenNullableFieldData only supported scalar types, causing error when partial upsert after adding nullable vector fields issue: #47160 - pickFieldData used row index directly to access Contents array without converting to data index, causing index out of range panic relate: #45993 Row-based processing had two problems: - Nullable vector handling was tightly coupled with other field types, requiring complex special-case logic - Update path was inefficient for nullable vectors: data had to be read first then updated in-place. For example, changing a valid vector to null requires shifting subsequent data since null values are not stored, which is very inefficient. Column-based processing solves both: - Each field type is processed independently, nullable vector logic is cleanly separated from other types - Nullable vectors are appended directly without read-then-update, avoiding expensive data shifting operations - Iterate over fields instead of rows when merging upsert and existing data key changes: - Remove nullable vector special handling: nullableVectorMergeContext, buildNullableVectorIdxMap, rebuildNullableVectorFieldData, etc. - Use generic column utilities: AppendFieldDataByColumn, UpdateFieldDataByColumn - Add vector type support to GenNullableFieldData - Add unit test for nullable vector upsert scenarios --------- Signed-off-by: marcelo-cjl <marcelo.chen@zilliz.com>
This commit is contained in:
marcelo-cjl
@@ -561,13 +561,16 @@ func pickFieldData(ids *schemapb.IDs, pkOffset map[any]int, fields []*schemapb.F
|
||||
// v3 v2 v5 v4 v1 (result vectors)
|
||||
// ===========================================
|
||||
fieldsData := make([]*schemapb.FieldData, len(fields))
|
||||
idxComputer := typeutil.NewFieldDataIdxComputer(fields)
|
||||
for i := 0; i < typeutil.GetSizeOfIDs(ids); i++ {
|
||||
id := typeutil.GetPK(ids, int64(i))
|
||||
if _, ok := pkOffset[id]; !ok {
|
||||
return nil, merr.WrapErrInconsistentRequery(fmt.Sprintf("incomplete query result, missing id %s, len(searchIDs) = %d, len(queryIDs) = %d, collection=%d",
|
||||
id, typeutil.GetSizeOfIDs(ids), len(pkOffset), collectionID))
|
||||
}
|
||||
typeutil.AppendFieldData(fieldsData, fields, int64(pkOffset[id]))
|
||||
rowIdx := int64(pkOffset[id])
|
||||
fieldIdxs := idxComputer.Compute(rowIdx)
|
||||
typeutil.AppendFieldData(fieldsData, fields, rowIdx, fieldIdxs...)
|
||||
}
|
||||
|
||||
return fieldsData, nil
|
||||
|
||||
@@ -4149,3 +4149,148 @@ func (s *SearchPipelineSuite) TestNewRequeryOperator_WithHighlighterNoDynamicFie
|
||||
// Verify only translated output fields are included (no dynamic fields added)
|
||||
s.Equal([]string{"title"}, reqOp.outputFieldNames)
|
||||
}
|
||||
|
||||
// TestPickFieldDataWithNullableSparseVector tests pickFieldData with nullable sparse vector fields.
|
||||
// This test ensures that when reordering query results (which may return in different order than requested),
|
||||
// nullable sparse vectors are correctly handled using FieldDataIdxComputer to map row indices to data indices.
|
||||
func (s *SearchPipelineSuite) TestPickFieldDataWithNullableSparseVector() {
|
||||
// Scenario: Search returns IDs [3, 1, 2], but Query returns them in order [1, 2, 3]
|
||||
// with a nullable sparse vector where row 1 (middle) is null.
|
||||
// pickFieldData should correctly reorder the data to match search result order.
|
||||
|
||||
// Search result IDs (the order we want)
|
||||
searchIDs := &schemapb.IDs{
|
||||
IdField: &schemapb.IDs_IntId{
|
||||
IntId: &schemapb.LongArray{
|
||||
Data: []int64{3, 1, 2}, // Want this order in final result
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Query returns data in different order: [1, 2, 3]
|
||||
// pkOffset maps each PK to its position in the query result
|
||||
pkOffset := map[any]int{
|
||||
int64(1): 0, // PK 1 is at index 0 in query result
|
||||
int64(2): 1, // PK 2 is at index 1 in query result
|
||||
int64(3): 2, // PK 3 is at index 2 in query result
|
||||
}
|
||||
|
||||
// Create sparse vector data for 3 rows, but row index 1 (PK=2) is null
|
||||
// ValidData: [true, false, true] means row 0 and 2 have data, row 1 is null
|
||||
// Contents only has 2 entries (for the non-null rows)
|
||||
sparseContent0, _ := testutils.GenerateSparseFloatVectorsData(1)
|
||||
sparseContent2, _ := testutils.GenerateSparseFloatVectorsData(1)
|
||||
|
||||
queryFields := []*schemapb.FieldData{
|
||||
{
|
||||
Type: schemapb.DataType_Int64,
|
||||
FieldName: "pk",
|
||||
FieldId: 100,
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{
|
||||
LongData: &schemapb.LongArray{
|
||||
Data: []int64{1, 2, 3}, // Query result order
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: schemapb.DataType_SparseFloatVector,
|
||||
FieldName: "sparse_vec",
|
||||
FieldId: 101,
|
||||
ValidData: []bool{true, false, true}, // Row 1 is null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 700,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 700,
|
||||
Contents: [][]byte{sparseContent0[0], sparseContent2[0]}, // Only 2 entries for non-null rows
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Call pickFieldData - this should NOT panic
|
||||
result, err := pickFieldData(searchIDs, pkOffset, queryFields, 12345)
|
||||
s.NoError(err)
|
||||
s.NotNil(result)
|
||||
s.Len(result, 2)
|
||||
|
||||
// Verify PK field is reordered correctly: [3, 1, 2]
|
||||
pkData := result[0].GetScalars().GetLongData().GetData()
|
||||
s.Equal([]int64{3, 1, 2}, pkData)
|
||||
|
||||
// Verify sparse vector ValidData is reordered correctly
|
||||
// Original ValidData: [true, false, true] for rows [1, 2, 3]
|
||||
// After reorder to [3, 1, 2]: ValidData should be [true, true, false]
|
||||
sparseValidData := result[1].GetValidData()
|
||||
s.Equal([]bool{true, true, false}, sparseValidData)
|
||||
|
||||
// Verify sparse vector Contents has correct number of entries (2 non-null values)
|
||||
sparseContents := result[1].GetVectors().GetSparseFloatVector().GetContents()
|
||||
s.Len(sparseContents, 2)
|
||||
}
|
||||
|
||||
// TestPickFieldDataWithAllNullSparseVector tests pickFieldData when all sparse vector values are null.
|
||||
func (s *SearchPipelineSuite) TestPickFieldDataWithAllNullSparseVector() {
|
||||
searchIDs := &schemapb.IDs{
|
||||
IdField: &schemapb.IDs_IntId{
|
||||
IntId: &schemapb.LongArray{
|
||||
Data: []int64{2, 1},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
pkOffset := map[any]int{
|
||||
int64(1): 0,
|
||||
int64(2): 1,
|
||||
}
|
||||
|
||||
queryFields := []*schemapb.FieldData{
|
||||
{
|
||||
Type: schemapb.DataType_Int64,
|
||||
FieldName: "pk",
|
||||
FieldId: 100,
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{
|
||||
LongData: &schemapb.LongArray{
|
||||
Data: []int64{1, 2},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: schemapb.DataType_SparseFloatVector,
|
||||
FieldName: "sparse_vec",
|
||||
FieldId: 101,
|
||||
ValidData: []bool{false, false}, // All null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 700,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 700,
|
||||
Contents: [][]byte{}, // Empty - all null
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
result, err := pickFieldData(searchIDs, pkOffset, queryFields, 12345)
|
||||
s.NoError(err)
|
||||
s.NotNil(result)
|
||||
|
||||
// All should still be null after reorder
|
||||
sparseValidData := result[1].GetValidData()
|
||||
s.Equal([]bool{false, false}, sparseValidData)
|
||||
s.Len(result[1].GetVectors().GetSparseFloatVector().GetContents(), 0)
|
||||
}
|
||||
|
||||
+215
-279
@@ -368,88 +368,92 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error {
|
||||
upsertFieldMap := lo.SliceToMap(it.upsertMsg.InsertMsg.GetFieldsData(), func(field *schemapb.FieldData) (int64, *schemapb.FieldData) {
|
||||
return field.FieldId, field
|
||||
})
|
||||
nullableVectorContexts := make(map[int64]*nullableVectorMergeContext)
|
||||
for _, fieldData := range existFieldData {
|
||||
fieldSchema, err := it.schema.schemaHelper.GetFieldFromName(fieldData.GetFieldName())
|
||||
if err != nil {
|
||||
log.Info("get field schema failed", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
if fieldSchema.GetNullable() && typeutil.IsVectorType(fieldSchema.GetDataType()) {
|
||||
ctx := &nullableVectorMergeContext{
|
||||
existIdxMap: buildNullableVectorIdxMap(fieldData.GetValidData()),
|
||||
existField: fieldData,
|
||||
mergedValid: make([]bool, 0, len(updateIdxInUpsert)),
|
||||
}
|
||||
|
||||
if upsertField, ok := upsertFieldMap[fieldData.FieldId]; ok {
|
||||
ctx.upsertIdxMap = buildNullableVectorIdxMap(upsertField.GetValidData())
|
||||
ctx.upsertField = upsertField
|
||||
ctx.hasUpsertData = true
|
||||
}
|
||||
|
||||
nullableVectorContexts[fieldData.FieldId] = ctx
|
||||
} else if fieldSchema.GetDefaultValue() != nil {
|
||||
// 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
|
||||
if len(fieldData.GetValidData()) == 0 {
|
||||
fieldData.ValidData = make([]bool, upsertIDSize)
|
||||
for i := range fieldData.ValidData {
|
||||
fieldData.ValidData[i] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build mapping from existing primary keys to their positions in query result
|
||||
// This ensures we can correctly locate data even if query results are not in the same order as request
|
||||
existIDsLen := typeutil.GetSizeOfIDs(existIDs)
|
||||
existPKToIndex := make(map[interface{}]int, existIDsLen)
|
||||
for j := 0; j < existIDsLen; j++ {
|
||||
pk := typeutil.GetPK(existIDs, int64(j))
|
||||
existPKToIndex[pk] = j
|
||||
for i := 0; i < existIDsLen; i++ {
|
||||
pk := typeutil.GetPK(existIDs, int64(i))
|
||||
existPKToIndex[pk] = i
|
||||
}
|
||||
|
||||
baseIdx := 0
|
||||
idxComputer := typeutil.NewFieldDataIdxComputer(existFieldData)
|
||||
for _, idx := range updateIdxInUpsert {
|
||||
typeutil.AppendIDs(it.deletePKs, upsertIDs, idx)
|
||||
oldPK := typeutil.GetPK(upsertIDs, int64(idx))
|
||||
existIndex, ok := existPKToIndex[oldPK]
|
||||
existIndices := make([]int64, len(updateIdxInUpsert))
|
||||
for i, upsertIdx := range updateIdxInUpsert {
|
||||
typeutil.AppendIDs(it.deletePKs, upsertIDs, upsertIdx)
|
||||
oldPK := typeutil.GetPK(upsertIDs, int64(upsertIdx))
|
||||
idx, ok := existPKToIndex[oldPK]
|
||||
if !ok {
|
||||
return merr.WrapErrParameterInvalidMsg("primary key not found in exist data mapping")
|
||||
return merr.WrapErrParameterInvalidMsg(fmt.Sprintf("upsert pk %v not found in query result", oldPK))
|
||||
}
|
||||
fieldIdxs := idxComputer.Compute(int64(existIndex))
|
||||
typeutil.AppendFieldData(it.insertFieldData, existFieldData, int64(existIndex), fieldIdxs...)
|
||||
err := typeutil.UpdateFieldData(it.insertFieldData, it.upsertMsg.InsertMsg.GetFieldsData(), int64(baseIdx), int64(idx))
|
||||
baseIdx += 1
|
||||
existIndices[i] = int64(idx)
|
||||
}
|
||||
|
||||
for fieldIdx, existField := range existFieldData {
|
||||
fieldSchema, err := it.schema.schemaHelper.GetFieldFromName(existField.GetFieldName())
|
||||
if err != nil {
|
||||
log.Info("update field data failed", zap.Error(err))
|
||||
log.Info("get field schema failed", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
// Collect merged validity for nullable vector fields
|
||||
for _, ctx := range nullableVectorContexts {
|
||||
if ctx.hasUpsertData {
|
||||
ctx.mergedValid = append(ctx.mergedValid, ctx.upsertField.GetValidData()[idx])
|
||||
} else {
|
||||
ctx.mergedValid = append(ctx.mergedValid, ctx.existField.GetValidData()[existIndex])
|
||||
// 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
|
||||
if fieldSchema.GetDefaultValue() != nil && len(existField.GetValidData()) == 0 {
|
||||
existField.ValidData = make([]bool, existIDsLen)
|
||||
for i := range existField.ValidData {
|
||||
existField.ValidData[i] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rebuild nullable vector fields
|
||||
for i, fieldData := range it.insertFieldData {
|
||||
ctx, ok := nullableVectorContexts[fieldData.FieldId]
|
||||
if !ok {
|
||||
continue
|
||||
dstField := it.insertFieldData[fieldIdx]
|
||||
upsertField := upsertFieldMap[existField.FieldId]
|
||||
isNullableVector := len(existField.GetValidData()) > 0 && typeutil.IsVectorType(existField.GetType())
|
||||
existComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{existField})
|
||||
existSrcIndices := make([]int64, len(updateIdxInUpsert))
|
||||
for i, existIdx := range existIndices {
|
||||
existSrcIndices[i] = existComputer.Compute(existIdx)[0]
|
||||
}
|
||||
|
||||
newFieldData := rebuildNullableVectorFieldData(ctx, updateIdxInUpsert, upsertIDs, existPKToIndex)
|
||||
if newFieldData != nil {
|
||||
it.insertFieldData[i] = newFieldData
|
||||
if upsertField != nil {
|
||||
upsertComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{upsertField})
|
||||
upsertSrcIndices := make([]int64, len(updateIdxInUpsert))
|
||||
for i, upsertIdx := range updateIdxInUpsert {
|
||||
upsertSrcIndices[i] = upsertComputer.Compute(int64(upsertIdx))[0]
|
||||
}
|
||||
if isNullableVector {
|
||||
// For nullable vector: only copy data for non-null rows
|
||||
upsertRowIndices := make([]int64, len(updateIdxInUpsert))
|
||||
validDataIndices := make([]int64, 0, len(updateIdxInUpsert))
|
||||
for i, upsertIdx := range updateIdxInUpsert {
|
||||
upsertRowIndices[i] = int64(upsertIdx)
|
||||
if upsertField.ValidData[upsertIdx] {
|
||||
validDataIndices = append(validDataIndices, upsertSrcIndices[i])
|
||||
}
|
||||
}
|
||||
typeutil.AppendFieldDataByColumn(dstField, upsertField, validDataIndices, upsertRowIndices)
|
||||
} else {
|
||||
typeutil.AppendFieldDataByColumn(dstField, existField, existSrcIndices)
|
||||
dstIndices := make([]int64, len(updateIdxInUpsert))
|
||||
for i := range dstIndices {
|
||||
dstIndices[i] = int64(i)
|
||||
}
|
||||
if err := typeutil.UpdateFieldDataByColumn(dstField, upsertField, dstIndices, upsertSrcIndices); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if isNullableVector {
|
||||
validDataIndices := make([]int64, 0, len(existIndices))
|
||||
for i, existIdx := range existIndices {
|
||||
if existField.ValidData[int(existIdx)] {
|
||||
validDataIndices = append(validDataIndices, existSrcIndices[i])
|
||||
}
|
||||
}
|
||||
typeutil.AppendFieldDataByColumn(dstField, existField, validDataIndices, existIndices)
|
||||
} else {
|
||||
typeutil.AppendFieldDataByColumn(dstField, existField, existSrcIndices)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -483,32 +487,47 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error {
|
||||
insertWithNullField = append(insertWithNullField, fieldData)
|
||||
}
|
||||
}
|
||||
vectorIdxMap := make([][]int64, len(insertIdxInUpsert))
|
||||
for rowIdx, offset := range insertIdxInUpsert {
|
||||
vectorIdxMap[rowIdx] = make([]int64, len(insertWithNullField))
|
||||
for fieldIdx := range insertWithNullField {
|
||||
vectorIdxMap[rowIdx][fieldIdx] = int64(offset)
|
||||
}
|
||||
}
|
||||
for fieldIdx, fieldData := range insertWithNullField {
|
||||
validData := fieldData.GetValidData()
|
||||
if len(validData) > 0 && typeutil.IsVectorType(fieldData.Type) {
|
||||
dataIdx := int64(0)
|
||||
rowIdx := 0
|
||||
for i := 0; i < len(validData) && rowIdx < len(insertIdxInUpsert); i++ {
|
||||
if i == insertIdxInUpsert[rowIdx] {
|
||||
vectorIdxMap[rowIdx][fieldIdx] = dataIdx
|
||||
rowIdx++
|
||||
}
|
||||
if validData[i] {
|
||||
dataIdx++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build mapping from FieldId to index in it.insertFieldData
|
||||
insertFieldDataMap := make(map[int64]int, len(it.insertFieldData))
|
||||
for i, fd := range it.insertFieldData {
|
||||
insertFieldDataMap[fd.FieldId] = i
|
||||
}
|
||||
|
||||
for rowIdx, idx := range insertIdxInUpsert {
|
||||
typeutil.AppendFieldData(it.insertFieldData, insertWithNullField, int64(idx), vectorIdxMap[rowIdx]...)
|
||||
// Process insert data by column (field), similar to update path
|
||||
for _, srcField := range insertWithNullField {
|
||||
isNullableVector := len(srcField.GetValidData()) > 0 && typeutil.IsVectorType(srcField.GetType())
|
||||
srcComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{srcField})
|
||||
|
||||
// Find or create destination field in it.insertFieldData
|
||||
var dstField *schemapb.FieldData
|
||||
if idx, ok := insertFieldDataMap[srcField.FieldId]; ok {
|
||||
dstField = it.insertFieldData[idx]
|
||||
} else {
|
||||
// New field not in existFieldData, create empty field data and append
|
||||
newFields := typeutil.PrepareResultFieldData([]*schemapb.FieldData{srcField}, int64(len(insertIdxInUpsert)))
|
||||
dstField = newFields[0]
|
||||
it.insertFieldData = append(it.insertFieldData, dstField)
|
||||
}
|
||||
|
||||
if isNullableVector {
|
||||
rowIndices := make([]int64, len(insertIdxInUpsert))
|
||||
validDataIndices := make([]int64, 0, len(insertIdxInUpsert))
|
||||
for i, upsertIdx := range insertIdxInUpsert {
|
||||
rowIndices[i] = int64(upsertIdx)
|
||||
srcIdx := srcComputer.Compute(int64(upsertIdx))[0]
|
||||
if srcField.ValidData[upsertIdx] {
|
||||
validDataIndices = append(validDataIndices, srcIdx)
|
||||
}
|
||||
}
|
||||
typeutil.AppendFieldDataByColumn(dstField, srcField, validDataIndices, rowIndices)
|
||||
} else {
|
||||
srcIndices := make([]int64, len(insertIdxInUpsert))
|
||||
for i, upsertIdx := range insertIdxInUpsert {
|
||||
srcIndices[i] = srcComputer.Compute(int64(upsertIdx))[0]
|
||||
}
|
||||
typeutil.AppendFieldDataByColumn(dstField, srcField, srcIndices)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -524,196 +543,6 @@ func (it *upsertTask) queryPreExecute(ctx context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildNullableVectorIdxMap(validData []bool) []int64 {
|
||||
if len(validData) == 0 {
|
||||
return nil
|
||||
}
|
||||
idxMap := make([]int64, len(validData))
|
||||
dataIdx := int64(0)
|
||||
for i, valid := range validData {
|
||||
if valid {
|
||||
idxMap[i] = dataIdx
|
||||
dataIdx++
|
||||
} else {
|
||||
idxMap[i] = -1
|
||||
}
|
||||
}
|
||||
return idxMap
|
||||
}
|
||||
|
||||
type nullableVectorMergeContext struct {
|
||||
existIdxMap []int64 // row index -> data index for exist data
|
||||
upsertIdxMap []int64 // row index -> data index for upsert data
|
||||
mergedValid []bool // merged validity for the output
|
||||
existField *schemapb.FieldData
|
||||
upsertField *schemapb.FieldData
|
||||
hasUpsertData bool // whether upsert request contains this field
|
||||
}
|
||||
|
||||
func rebuildNullableVectorFieldData(
|
||||
ctx *nullableVectorMergeContext,
|
||||
updateIdxInUpsert []int,
|
||||
upsertIDs *schemapb.IDs,
|
||||
existPKToIndex map[interface{}]int,
|
||||
) *schemapb.FieldData {
|
||||
var sourceField *schemapb.FieldData
|
||||
var sourceIdxMap []int64
|
||||
|
||||
if ctx.hasUpsertData {
|
||||
sourceField = ctx.upsertField
|
||||
sourceIdxMap = ctx.upsertIdxMap
|
||||
} else {
|
||||
sourceField = ctx.existField
|
||||
sourceIdxMap = ctx.existIdxMap
|
||||
}
|
||||
|
||||
if sourceField == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
newFieldData := prepareNullableVectorFieldData(sourceField, int64(len(updateIdxInUpsert)))
|
||||
newFieldData.ValidData = ctx.mergedValid
|
||||
|
||||
// Append vectors from source using the correct indices
|
||||
for rowIdx, idx := range updateIdxInUpsert {
|
||||
var sourceRowIdx int
|
||||
if ctx.hasUpsertData {
|
||||
sourceRowIdx = idx // upsertMsg uses updateIdxInUpsert
|
||||
} else {
|
||||
// existFieldData uses existIndex
|
||||
oldPK := typeutil.GetPK(upsertIDs, int64(idx))
|
||||
sourceRowIdx = existPKToIndex[oldPK]
|
||||
}
|
||||
|
||||
// Check if this row is valid (has vector data)
|
||||
if rowIdx < len(ctx.mergedValid) && ctx.mergedValid[rowIdx] {
|
||||
dataIdx := int64(sourceRowIdx)
|
||||
if len(sourceIdxMap) > 0 && sourceRowIdx < len(sourceIdxMap) {
|
||||
dataIdx = sourceIdxMap[sourceRowIdx]
|
||||
}
|
||||
if dataIdx >= 0 {
|
||||
appendSingleVector(newFieldData, sourceField, dataIdx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return newFieldData
|
||||
}
|
||||
|
||||
func prepareNullableVectorFieldData(sample *schemapb.FieldData, capacity int64) *schemapb.FieldData {
|
||||
fd := &schemapb.FieldData{
|
||||
Type: sample.Type,
|
||||
FieldName: sample.FieldName,
|
||||
FieldId: sample.FieldId,
|
||||
IsDynamic: sample.IsDynamic,
|
||||
ValidData: make([]bool, 0, capacity),
|
||||
}
|
||||
|
||||
vectorField := sample.GetVectors()
|
||||
if vectorField == nil {
|
||||
return fd
|
||||
}
|
||||
dim := vectorField.GetDim()
|
||||
vectors := &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
},
|
||||
}
|
||||
|
||||
switch vectorField.Data.(type) {
|
||||
case *schemapb.VectorField_FloatVector:
|
||||
vectors.Vectors.Data = &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{
|
||||
Data: make([]float32, 0, dim*capacity),
|
||||
},
|
||||
}
|
||||
case *schemapb.VectorField_Float16Vector:
|
||||
vectors.Vectors.Data = &schemapb.VectorField_Float16Vector{
|
||||
Float16Vector: make([]byte, 0, dim*2*capacity),
|
||||
}
|
||||
case *schemapb.VectorField_Bfloat16Vector:
|
||||
vectors.Vectors.Data = &schemapb.VectorField_Bfloat16Vector{
|
||||
Bfloat16Vector: make([]byte, 0, dim*2*capacity),
|
||||
}
|
||||
case *schemapb.VectorField_BinaryVector:
|
||||
vectors.Vectors.Data = &schemapb.VectorField_BinaryVector{
|
||||
BinaryVector: make([]byte, 0, dim/8*capacity),
|
||||
}
|
||||
case *schemapb.VectorField_Int8Vector:
|
||||
vectors.Vectors.Data = &schemapb.VectorField_Int8Vector{
|
||||
Int8Vector: make([]byte, 0, dim*capacity),
|
||||
}
|
||||
case *schemapb.VectorField_SparseFloatVector:
|
||||
vectors.Vectors.Data = &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Contents: make([][]byte, 0, capacity),
|
||||
Dim: vectorField.GetSparseFloatVector().GetDim(),
|
||||
},
|
||||
}
|
||||
}
|
||||
fd.Field = vectors
|
||||
return fd
|
||||
}
|
||||
|
||||
func appendSingleVector(target *schemapb.FieldData, source *schemapb.FieldData, dataIdx int64) {
|
||||
targetVectors := target.GetVectors()
|
||||
sourceVectors := source.GetVectors()
|
||||
if targetVectors == nil || sourceVectors == nil {
|
||||
return
|
||||
}
|
||||
dim := sourceVectors.GetDim()
|
||||
|
||||
switch sv := sourceVectors.Data.(type) {
|
||||
case *schemapb.VectorField_FloatVector:
|
||||
tv := targetVectors.Data.(*schemapb.VectorField_FloatVector)
|
||||
start := dataIdx * dim
|
||||
end := start + dim
|
||||
if end <= int64(len(sv.FloatVector.Data)) {
|
||||
tv.FloatVector.Data = append(tv.FloatVector.Data, sv.FloatVector.Data[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_Float16Vector:
|
||||
tv := targetVectors.Data.(*schemapb.VectorField_Float16Vector)
|
||||
unitSize := dim * 2
|
||||
start := dataIdx * unitSize
|
||||
end := start + unitSize
|
||||
if end <= int64(len(sv.Float16Vector)) {
|
||||
tv.Float16Vector = append(tv.Float16Vector, sv.Float16Vector[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_Bfloat16Vector:
|
||||
tv := targetVectors.Data.(*schemapb.VectorField_Bfloat16Vector)
|
||||
unitSize := dim * 2
|
||||
start := dataIdx * unitSize
|
||||
end := start + unitSize
|
||||
if end <= int64(len(sv.Bfloat16Vector)) {
|
||||
tv.Bfloat16Vector = append(tv.Bfloat16Vector, sv.Bfloat16Vector[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_BinaryVector:
|
||||
tv := targetVectors.Data.(*schemapb.VectorField_BinaryVector)
|
||||
unitSize := dim / 8
|
||||
start := dataIdx * unitSize
|
||||
end := start + unitSize
|
||||
if end <= int64(len(sv.BinaryVector)) {
|
||||
tv.BinaryVector = append(tv.BinaryVector, sv.BinaryVector[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_Int8Vector:
|
||||
tv := targetVectors.Data.(*schemapb.VectorField_Int8Vector)
|
||||
start := dataIdx * dim
|
||||
end := start + dim
|
||||
if end <= int64(len(sv.Int8Vector)) {
|
||||
tv.Int8Vector = append(tv.Int8Vector, sv.Int8Vector[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_SparseFloatVector:
|
||||
tv := targetVectors.Data.(*schemapb.VectorField_SparseFloatVector)
|
||||
if dataIdx < int64(len(sv.SparseFloatVector.Contents)) {
|
||||
tv.SparseFloatVector.Contents = append(tv.SparseFloatVector.Contents, sv.SparseFloatVector.Contents[dataIdx])
|
||||
// Update dimension if necessary
|
||||
if sv.SparseFloatVector.Dim > tv.SparseFloatVector.Dim {
|
||||
tv.SparseFloatVector.Dim = sv.SparseFloatVector.Dim
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ToCompressedFormatNullable converts the field data from full format nullable to compressed format nullable
|
||||
func ToCompressedFormatNullable(field *schemapb.FieldData) error {
|
||||
if getValidNumber(field.GetValidData()) == len(field.GetValidData()) {
|
||||
@@ -1073,8 +902,115 @@ func GenNullableFieldData(field *schemapb.FieldSchema, upsertIDSize int) (*schem
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
// Nullable vector types
|
||||
case schemapb.DataType_FloatVector:
|
||||
dim, err := typeutil.GetDim(field)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &schemapb.FieldData{
|
||||
FieldId: field.FieldID,
|
||||
FieldName: field.Name,
|
||||
Type: field.DataType,
|
||||
ValidData: make([]bool, upsertIDSize), // all false = all null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{}}},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
case schemapb.DataType_Float16Vector:
|
||||
dim, err := typeutil.GetDim(field)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &schemapb.FieldData{
|
||||
FieldId: field.FieldID,
|
||||
FieldName: field.Name,
|
||||
Type: field.DataType,
|
||||
ValidData: make([]bool, upsertIDSize), // all false = all null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{}},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
case schemapb.DataType_BFloat16Vector:
|
||||
dim, err := typeutil.GetDim(field)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &schemapb.FieldData{
|
||||
FieldId: field.FieldID,
|
||||
FieldName: field.Name,
|
||||
Type: field.DataType,
|
||||
ValidData: make([]bool, upsertIDSize), // all false = all null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{}},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
case schemapb.DataType_BinaryVector:
|
||||
dim, err := typeutil.GetDim(field)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &schemapb.FieldData{
|
||||
FieldId: field.FieldID,
|
||||
FieldName: field.Name,
|
||||
Type: field.DataType,
|
||||
ValidData: make([]bool, upsertIDSize), // all false = all null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{}},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
case schemapb.DataType_SparseFloatVector:
|
||||
return &schemapb.FieldData{
|
||||
FieldId: field.FieldID,
|
||||
FieldName: field.Name,
|
||||
Type: field.DataType,
|
||||
ValidData: make([]bool, upsertIDSize), // all false = all null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Contents: [][]byte{},
|
||||
}},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
case schemapb.DataType_Int8Vector:
|
||||
dim, err := typeutil.GetDim(field)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &schemapb.FieldData{
|
||||
FieldId: field.FieldID,
|
||||
FieldName: field.Name,
|
||||
Type: field.DataType,
|
||||
ValidData: make([]bool, upsertIDSize), // all false = all null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{}},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
default:
|
||||
return nil, merr.WrapErrParameterInvalidMsg(fmt.Sprintf("undefined scalar data type:%s", field.DataType.String()))
|
||||
return nil, merr.WrapErrParameterInvalidMsg(fmt.Sprintf("undefined data type:%s", field.DataType.String()))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+255
-349
@@ -34,6 +34,7 @@ import (
|
||||
"github.com/milvus-io/milvus/internal/proxy/shardclient"
|
||||
"github.com/milvus-io/milvus/internal/util/function/embedding"
|
||||
"github.com/milvus-io/milvus/internal/util/segcore"
|
||||
"github.com/milvus-io/milvus/pkg/v2/common"
|
||||
"github.com/milvus-io/milvus/pkg/v2/mq/msgstream"
|
||||
"github.com/milvus-io/milvus/pkg/v2/proto/planpb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/proto/rootcoordpb"
|
||||
@@ -1826,382 +1827,287 @@ func TestUpsertTask_queryPreExecute_EmptyDataArray(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestBuildNullableVectorIdxMap(t *testing.T) {
|
||||
t.Run("empty validData", func(t *testing.T) {
|
||||
result := buildNullableVectorIdxMap(nil)
|
||||
assert.Nil(t, result)
|
||||
|
||||
result = buildNullableVectorIdxMap([]bool{})
|
||||
assert.Nil(t, result)
|
||||
})
|
||||
|
||||
t.Run("all valid", func(t *testing.T) {
|
||||
validData := []bool{true, true, true, true}
|
||||
result := buildNullableVectorIdxMap(validData)
|
||||
assert.Equal(t, []int64{0, 1, 2, 3}, result)
|
||||
})
|
||||
|
||||
t.Run("all null", func(t *testing.T) {
|
||||
validData := []bool{false, false, false, false}
|
||||
result := buildNullableVectorIdxMap(validData)
|
||||
assert.Equal(t, []int64{-1, -1, -1, -1}, result)
|
||||
})
|
||||
|
||||
t.Run("mixed valid and null", func(t *testing.T) {
|
||||
validData := []bool{true, false, true, false, true}
|
||||
result := buildNullableVectorIdxMap(validData)
|
||||
// dataIdx: 0 for row 0, -1 for row 1, 1 for row 2, -1 for row 3, 2 for row 4
|
||||
assert.Equal(t, []int64{0, -1, 1, -1, 2}, result)
|
||||
})
|
||||
|
||||
t.Run("alternating pattern", func(t *testing.T) {
|
||||
validData := []bool{false, true, false, true, false, true}
|
||||
result := buildNullableVectorIdxMap(validData)
|
||||
assert.Equal(t, []int64{-1, 0, -1, 1, -1, 2}, result)
|
||||
})
|
||||
}
|
||||
|
||||
func TestPrepareNullableVectorFieldData(t *testing.T) {
|
||||
func TestUpsertTask_queryPreExecute_NullableFields(t *testing.T) {
|
||||
dim := int64(4)
|
||||
|
||||
t.Run("FloatVector", func(t *testing.T) {
|
||||
sample := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_FloatVector,
|
||||
FieldName: "float_vec",
|
||||
FieldId: 100,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}},
|
||||
schema := newSchemaInfo(&schemapb.CollectionSchema{
|
||||
Name: "test_nullable_vec",
|
||||
Fields: []*schemapb.FieldSchema{
|
||||
{FieldID: 100, Name: "id", IsPrimaryKey: true, DataType: schemapb.DataType_Int64},
|
||||
{FieldID: 101, Name: "vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}},
|
||||
{FieldID: 102, Name: "nullable_vec", DataType: schemapb.DataType_FloatVector, Nullable: true, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}},
|
||||
},
|
||||
})
|
||||
|
||||
// Generate vector data: [pk, pk, pk, pk]
|
||||
genVec := func(pk int64) []float32 {
|
||||
return []float32{float32(pk), float32(pk), float32(pk), float32(pk)}
|
||||
}
|
||||
|
||||
// Create all_columns upsert data (includes nullable_vec)
|
||||
// nullable_vec = [pk+100, pk+100, pk+100, pk+100], ValidData = all true
|
||||
createAllCols := func(pks []int64) []*schemapb.FieldData {
|
||||
var ids []int64
|
||||
var vecData, nullableData []float32
|
||||
var validData []bool
|
||||
for _, pk := range pks {
|
||||
ids = append(ids, pk)
|
||||
vecData = append(vecData, genVec(pk)...)
|
||||
nullableData = append(nullableData, genVec(pk+100)...)
|
||||
validData = append(validData, true)
|
||||
}
|
||||
return []*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: ids}}}},
|
||||
},
|
||||
{
|
||||
FieldName: "vector", FieldId: 101, Type: schemapb.DataType_FloatVector,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: vecData}}}},
|
||||
},
|
||||
{
|
||||
FieldName: "nullable_vec", FieldId: 102, Type: schemapb.DataType_FloatVector, ValidData: validData,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: nullableData}}}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Create partial_columns upsert data (excludes nullable_vec)
|
||||
createPartialCols := func(pks []int64) []*schemapb.FieldData {
|
||||
var ids []int64
|
||||
var vecData []float32
|
||||
for _, pk := range pks {
|
||||
ids = append(ids, pk)
|
||||
vecData = append(vecData, genVec(pk)...)
|
||||
}
|
||||
return []*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: ids}}}},
|
||||
},
|
||||
{
|
||||
FieldName: "vector", FieldId: 101, Type: schemapb.DataType_FloatVector,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: vecData}}}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Create mock query result
|
||||
// existing nullable_vec = [pk+300, pk+300, pk+300, pk+300], ValidData = all true
|
||||
queryResult := func(pks []int64) *milvuspb.QueryResults {
|
||||
var ids []int64
|
||||
var vecData, nullableData []float32
|
||||
var validData []bool
|
||||
for _, pk := range pks {
|
||||
ids = append(ids, pk)
|
||||
vecData = append(vecData, genVec(pk+200)...)
|
||||
nullableData = append(nullableData, genVec(pk+300)...)
|
||||
validData = append(validData, true)
|
||||
}
|
||||
return &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: ids}}}},
|
||||
},
|
||||
{
|
||||
FieldName: "vector", FieldId: 101, Type: schemapb.DataType_FloatVector,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: vecData}}}},
|
||||
},
|
||||
{
|
||||
FieldName: "nullable_vec", FieldId: 102, Type: schemapb.DataType_FloatVector, ValidData: validData,
|
||||
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: nullableData}}}},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := prepareNullableVectorFieldData(sample, 10)
|
||||
assert.Equal(t, schemapb.DataType_FloatVector, result.Type)
|
||||
assert.Equal(t, "float_vec", result.FieldName)
|
||||
assert.Equal(t, int64(100), result.FieldId)
|
||||
assert.Equal(t, dim, result.GetVectors().GetDim())
|
||||
assert.Empty(t, result.GetVectors().GetFloatVector().Data)
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("Float16Vector", func(t *testing.T) {
|
||||
sample := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_Float16Vector,
|
||||
FieldName: "fp16_vec",
|
||||
FieldId: 101,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
|
||||
runUpsert := func(upsertData []*schemapb.FieldData, mockResult *milvuspb.QueryResults) *upsertTask {
|
||||
numRows := uint32(len(upsertData[0].GetScalars().GetLongData().GetData()))
|
||||
task := &upsertTask{
|
||||
ctx: context.Background(),
|
||||
schema: schema,
|
||||
req: &milvuspb.UpsertRequest{FieldsData: upsertData, NumRows: numRows},
|
||||
upsertMsg: &msgstream.UpsertMsg{InsertMsg: &msgstream.InsertMsg{
|
||||
InsertRequest: &msgpb.InsertRequest{
|
||||
FieldsData: upsertData,
|
||||
NumRows: uint64(numRows),
|
||||
Version: msgpb.InsertDataVersion_ColumnBased, // Required, otherwise NRows() returns 0
|
||||
},
|
||||
},
|
||||
}},
|
||||
node: &Proxy{},
|
||||
}
|
||||
result := prepareNullableVectorFieldData(sample, 10)
|
||||
assert.Equal(t, schemapb.DataType_Float16Vector, result.Type)
|
||||
assert.Equal(t, dim, result.GetVectors().GetDim())
|
||||
assert.Empty(t, result.GetVectors().GetFloat16Vector())
|
||||
})
|
||||
mock := mockey.Mock(retrieveByPKs).Return(mockResult, segcore.StorageCost{}, nil).Build()
|
||||
defer mock.UnPatch()
|
||||
err := task.queryPreExecute(context.Background())
|
||||
assert.NoError(t, err)
|
||||
return task
|
||||
}
|
||||
|
||||
t.Run("BFloat16Vector", func(t *testing.T) {
|
||||
sample := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_BFloat16Vector,
|
||||
FieldName: "bf16_vec",
|
||||
FieldId: 102,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := prepareNullableVectorFieldData(sample, 10)
|
||||
assert.Equal(t, schemapb.DataType_BFloat16Vector, result.Type)
|
||||
assert.Equal(t, dim, result.GetVectors().GetDim())
|
||||
assert.Empty(t, result.GetVectors().GetBfloat16Vector())
|
||||
})
|
||||
// Step 1a: Empty data, upsert pk1(partial) -> insert, nullable_vec=null
|
||||
task1a := runUpsert(createPartialCols([]int64{1}), queryResult(nil))
|
||||
assert.Empty(t, task1a.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{1}, task1a.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
assert.Equal(t, []float32{1, 1, 1, 1}, task1a.insertFieldData[1].GetVectors().GetFloatVector().GetData())
|
||||
assert.Equal(t, []bool{false}, task1a.insertFieldData[2].ValidData)
|
||||
assert.Empty(t, task1a.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
|
||||
t.Run("BinaryVector", func(t *testing.T) {
|
||||
sample := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_BinaryVector,
|
||||
FieldName: "binary_vec",
|
||||
FieldId: 103,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: int64(32), // binary vector dim is in bits
|
||||
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{1, 2, 3, 4}},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := prepareNullableVectorFieldData(sample, 10)
|
||||
assert.Equal(t, schemapb.DataType_BinaryVector, result.Type)
|
||||
assert.Equal(t, int64(32), result.GetVectors().GetDim())
|
||||
assert.Empty(t, result.GetVectors().GetBinaryVector())
|
||||
})
|
||||
// Step 1b: Empty data, upsert pk2(all) -> insert, nullable_vec=[102,...]
|
||||
task1b := runUpsert(createAllCols([]int64{2}), queryResult(nil))
|
||||
assert.Empty(t, task1b.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{2}, task1b.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
assert.Equal(t, []float32{2, 2, 2, 2}, task1b.insertFieldData[1].GetVectors().GetFloatVector().GetData())
|
||||
assert.Equal(t, []float32{102, 102, 102, 102}, task1b.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
|
||||
t.Run("SparseFloatVector", func(t *testing.T) {
|
||||
sample := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_SparseFloatVector,
|
||||
FieldName: "sparse_vec",
|
||||
FieldId: 104,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 100,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 100,
|
||||
Contents: [][]byte{{1, 2, 3}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := prepareNullableVectorFieldData(sample, 10)
|
||||
assert.Equal(t, schemapb.DataType_SparseFloatVector, result.Type)
|
||||
assert.Empty(t, result.GetVectors().GetSparseFloatVector().Contents)
|
||||
})
|
||||
// Step 2a: pk1 exists, upsert pk1(all) -> update, nullable_vec=[101,...] (from upsert)
|
||||
task2a := runUpsert(createAllCols([]int64{1}), queryResult([]int64{1}))
|
||||
assert.Equal(t, []int64{1}, task2a.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{1}, task2a.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
assert.Equal(t, []float32{1, 1, 1, 1}, task2a.insertFieldData[1].GetVectors().GetFloatVector().GetData())
|
||||
assert.Equal(t, []float32{101, 101, 101, 101}, task2a.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
|
||||
t.Run("nil vectors", func(t *testing.T) {
|
||||
sample := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_FloatVector,
|
||||
FieldName: "empty_vec",
|
||||
FieldId: 105,
|
||||
}
|
||||
result := prepareNullableVectorFieldData(sample, 10)
|
||||
assert.Equal(t, schemapb.DataType_FloatVector, result.Type)
|
||||
assert.Nil(t, result.GetVectors())
|
||||
})
|
||||
// Step 2b: pk2 exists, upsert pk2(partial) -> update, nullable_vec=[302,...] (from existing)
|
||||
task2b := runUpsert(createPartialCols([]int64{2}), queryResult([]int64{2}))
|
||||
assert.Equal(t, []int64{2}, task2b.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{2}, task2b.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
assert.Equal(t, []float32{2, 2, 2, 2}, task2b.insertFieldData[1].GetVectors().GetFloatVector().GetData())
|
||||
assert.Equal(t, []float32{302, 302, 302, 302}, task2b.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
|
||||
// Step 3a: Empty data, upsert pk3(partial) -> insert, nullable_vec=null
|
||||
task3a := runUpsert(createPartialCols([]int64{3}), queryResult(nil))
|
||||
assert.Empty(t, task3a.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{3}, task3a.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
assert.Equal(t, []bool{false}, task3a.insertFieldData[2].ValidData)
|
||||
assert.Empty(t, task3a.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
|
||||
// Step 3b: Empty data, upsert pk4(all) -> insert, nullable_vec=[104,...]
|
||||
task3b := runUpsert(createAllCols([]int64{4}), queryResult(nil))
|
||||
assert.Empty(t, task3b.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{4}, task3b.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
assert.Equal(t, []float32{104, 104, 104, 104}, task3b.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
|
||||
// Step 4a: pk3,pk4 exist, upsert pk3,pk4,pk5,pk6(all) -> pk3,pk4 update, pk5,pk6 insert
|
||||
task4a := runUpsert(createAllCols([]int64{3, 4, 5, 6}), queryResult([]int64{3, 4}))
|
||||
assert.Equal(t, []int64{3, 4}, task4a.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{3, 4, 5, 6}, task4a.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
assert.Equal(t, []float32{3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6}, task4a.insertFieldData[1].GetVectors().GetFloatVector().GetData())
|
||||
assert.Equal(t, []float32{103, 103, 103, 103, 104, 104, 104, 104, 105, 105, 105, 105, 106, 106, 106, 106}, task4a.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
|
||||
// Step 4b: pk3,pk4 exist, upsert pk3,pk4,pk5,pk6(partial) -> pk3,pk4 update (use existing), pk5,pk6 insert (null)
|
||||
task4b := runUpsert(createPartialCols([]int64{3, 4, 5, 6}), queryResult([]int64{3, 4}))
|
||||
assert.Equal(t, []int64{3, 4}, task4b.deletePKs.GetIntId().GetData())
|
||||
assert.Equal(t, []int64{3, 4, 5, 6}, task4b.insertFieldData[0].GetScalars().GetLongData().GetData())
|
||||
// Update rows pk3,pk4: nullable_vec from existing data (ValidData=true)
|
||||
// Insert rows pk5,pk6: nullable_vec generated by GenNullableFieldData (null, ValidData=false)
|
||||
// ValidData has 4 elements, FloatVector only contains data for ValidData=true rows
|
||||
assert.Equal(t, []bool{true, true, false, false}, task4b.insertFieldData[2].ValidData)
|
||||
assert.Equal(t, []float32{303, 303, 303, 303, 304, 304, 304, 304}, task4b.insertFieldData[2].GetVectors().GetFloatVector().GetData())
|
||||
}
|
||||
|
||||
func TestAppendSingleVector(t *testing.T) {
|
||||
t.Run("FloatVector", func(t *testing.T) {
|
||||
dim := int64(4)
|
||||
source := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6, 7, 8}}, // 2 vectors
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
target := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{}}},
|
||||
},
|
||||
},
|
||||
func TestUpsertTask_GenNullableFieldData(t *testing.T) {
|
||||
upsertIDSize := 5
|
||||
|
||||
t.Run("scalar_types", func(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
dataType schemapb.DataType
|
||||
}{
|
||||
{"Bool", schemapb.DataType_Bool},
|
||||
{"Int32", schemapb.DataType_Int32},
|
||||
{"Int64", schemapb.DataType_Int64},
|
||||
{"Float", schemapb.DataType_Float},
|
||||
{"Double", schemapb.DataType_Double},
|
||||
{"VarChar", schemapb.DataType_VarChar},
|
||||
{"JSON", schemapb.DataType_JSON},
|
||||
{"Array", schemapb.DataType_Array},
|
||||
{"Timestamptz", schemapb.DataType_Timestamptz},
|
||||
{"Geometry", schemapb.DataType_Geometry},
|
||||
}
|
||||
|
||||
appendSingleVector(target, source, 0)
|
||||
assert.Equal(t, []float32{1, 2, 3, 4}, target.GetVectors().GetFloatVector().Data)
|
||||
|
||||
appendSingleVector(target, source, 1)
|
||||
assert.Equal(t, []float32{1, 2, 3, 4, 5, 6, 7, 8}, target.GetVectors().GetFloatVector().Data)
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{
|
||||
FieldID: 100,
|
||||
Name: "test_field",
|
||||
DataType: tc.dataType,
|
||||
Nullable: true,
|
||||
}
|
||||
result, err := GenNullableFieldData(field, upsertIDSize)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, field.FieldID, result.FieldId)
|
||||
assert.Equal(t, field.Name, result.FieldName)
|
||||
assert.Equal(t, tc.dataType, result.Type)
|
||||
assert.Equal(t, upsertIDSize, len(result.ValidData))
|
||||
// All ValidData should be false (null)
|
||||
for _, v := range result.ValidData {
|
||||
assert.False(t, v)
|
||||
}
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("Float16Vector", func(t *testing.T) {
|
||||
dim := int64(2)
|
||||
source := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}}, // 2 vectors (dim*2 bytes each)
|
||||
},
|
||||
},
|
||||
}
|
||||
target := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{}},
|
||||
},
|
||||
},
|
||||
t.Run("vector_types", func(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
dataType schemapb.DataType
|
||||
}{
|
||||
{"FloatVector", schemapb.DataType_FloatVector},
|
||||
{"Float16Vector", schemapb.DataType_Float16Vector},
|
||||
{"BFloat16Vector", schemapb.DataType_BFloat16Vector},
|
||||
{"BinaryVector", schemapb.DataType_BinaryVector},
|
||||
{"Int8Vector", schemapb.DataType_Int8Vector},
|
||||
}
|
||||
|
||||
appendSingleVector(target, source, 0)
|
||||
assert.Equal(t, []byte{1, 2, 3, 4}, target.GetVectors().GetFloat16Vector())
|
||||
|
||||
appendSingleVector(target, source, 1)
|
||||
assert.Equal(t, []byte{1, 2, 3, 4, 5, 6, 7, 8}, target.GetVectors().GetFloat16Vector())
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{
|
||||
FieldID: 100,
|
||||
Name: "test_vector",
|
||||
DataType: tc.dataType,
|
||||
Nullable: true,
|
||||
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "128"}},
|
||||
}
|
||||
result, err := GenNullableFieldData(field, upsertIDSize)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, field.FieldID, result.FieldId)
|
||||
assert.Equal(t, field.Name, result.FieldName)
|
||||
assert.Equal(t, tc.dataType, result.Type)
|
||||
assert.Equal(t, upsertIDSize, len(result.ValidData))
|
||||
// All ValidData should be false (null)
|
||||
for _, v := range result.ValidData {
|
||||
assert.False(t, v)
|
||||
}
|
||||
assert.NotNil(t, result.GetVectors())
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("BinaryVector", func(t *testing.T) {
|
||||
dim := int64(16) // 16 bits = 2 bytes per vector
|
||||
source := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{1, 2, 3, 4}}, // 2 vectors
|
||||
},
|
||||
},
|
||||
t.Run("sparse_float_vector", func(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{
|
||||
FieldID: 100,
|
||||
Name: "test_sparse",
|
||||
DataType: schemapb.DataType_SparseFloatVector,
|
||||
Nullable: true,
|
||||
}
|
||||
target := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{}},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
appendSingleVector(target, source, 0)
|
||||
assert.Equal(t, []byte{1, 2}, target.GetVectors().GetBinaryVector())
|
||||
|
||||
appendSingleVector(target, source, 1)
|
||||
assert.Equal(t, []byte{1, 2, 3, 4}, target.GetVectors().GetBinaryVector())
|
||||
})
|
||||
|
||||
t.Run("SparseFloatVector", func(t *testing.T) {
|
||||
source := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 100,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 100,
|
||||
Contents: [][]byte{{1, 2}, {3, 4}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
target := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 100,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 0,
|
||||
Contents: [][]byte{},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
appendSingleVector(target, source, 0)
|
||||
assert.Equal(t, [][]byte{{1, 2}}, target.GetVectors().GetSparseFloatVector().Contents)
|
||||
|
||||
appendSingleVector(target, source, 1)
|
||||
assert.Equal(t, [][]byte{{1, 2}, {3, 4}}, target.GetVectors().GetSparseFloatVector().Contents)
|
||||
})
|
||||
|
||||
t.Run("nil vectors", func(t *testing.T) {
|
||||
source := &schemapb.FieldData{}
|
||||
target := &schemapb.FieldData{}
|
||||
// Should not panic
|
||||
appendSingleVector(target, source, 0)
|
||||
})
|
||||
}
|
||||
|
||||
func TestRebuildNullableVectorFieldData(t *testing.T) {
|
||||
dim := int64(4)
|
||||
|
||||
t.Run("with upsert data - all valid", func(t *testing.T) {
|
||||
upsertField := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_FloatVector,
|
||||
FieldName: "vec",
|
||||
FieldId: 100,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}}},
|
||||
},
|
||||
},
|
||||
ValidData: []bool{true, true, true},
|
||||
}
|
||||
|
||||
ctx := &nullableVectorMergeContext{
|
||||
upsertIdxMap: []int64{0, 1, 2},
|
||||
upsertField: upsertField,
|
||||
hasUpsertData: true,
|
||||
mergedValid: []bool{true, true},
|
||||
}
|
||||
|
||||
updateIdxInUpsert := []int{0, 2}
|
||||
upsertIDs := &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}}
|
||||
existPKToIndex := map[interface{}]int{}
|
||||
|
||||
result := rebuildNullableVectorFieldData(ctx, updateIdxInUpsert, upsertIDs, existPKToIndex)
|
||||
result, err := GenNullableFieldData(field, upsertIDSize)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, []bool{true, true}, result.ValidData)
|
||||
// Should have 2 vectors: index 0 and index 2 from source
|
||||
assert.Equal(t, []float32{1, 2, 3, 4, 9, 10, 11, 12}, result.GetVectors().GetFloatVector().Data)
|
||||
assert.Equal(t, upsertIDSize, len(result.ValidData))
|
||||
assert.NotNil(t, result.GetVectors().GetSparseFloatVector())
|
||||
})
|
||||
|
||||
t.Run("with upsert data - mixed valid/null", func(t *testing.T) {
|
||||
upsertField := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_FloatVector,
|
||||
FieldName: "vec",
|
||||
FieldId: 100,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6, 7, 8}}}, // only 2 vectors (compressed)
|
||||
},
|
||||
},
|
||||
ValidData: []bool{true, false, true}, // row 0 valid, row 1 null, row 2 valid
|
||||
t.Run("unsupported_type", func(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{
|
||||
FieldID: 100,
|
||||
Name: "test_unsupported",
|
||||
DataType: schemapb.DataType_None,
|
||||
Nullable: true,
|
||||
}
|
||||
|
||||
ctx := &nullableVectorMergeContext{
|
||||
upsertIdxMap: []int64{0, -1, 1}, // row 0 -> data 0, row 1 -> null, row 2 -> data 1
|
||||
upsertField: upsertField,
|
||||
hasUpsertData: true,
|
||||
mergedValid: []bool{true, false}, // result: row 0 valid, row 1 null
|
||||
}
|
||||
|
||||
updateIdxInUpsert := []int{0, 1}
|
||||
upsertIDs := &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}}
|
||||
existPKToIndex := map[interface{}]int{}
|
||||
|
||||
result := rebuildNullableVectorFieldData(ctx, updateIdxInUpsert, upsertIDs, existPKToIndex)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, []bool{true, false}, result.ValidData)
|
||||
// Should have only 1 vector (row 0), row 1 is null
|
||||
assert.Equal(t, []float32{1, 2, 3, 4}, result.GetVectors().GetFloatVector().Data)
|
||||
})
|
||||
|
||||
t.Run("without upsert data - use exist data", func(t *testing.T) {
|
||||
existField := &schemapb.FieldData{
|
||||
Type: schemapb.DataType_FloatVector,
|
||||
FieldName: "vec",
|
||||
FieldId: 100,
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6, 7, 8}}},
|
||||
},
|
||||
},
|
||||
ValidData: []bool{true, true},
|
||||
}
|
||||
|
||||
ctx := &nullableVectorMergeContext{
|
||||
existIdxMap: []int64{0, 1},
|
||||
existField: existField,
|
||||
hasUpsertData: false,
|
||||
mergedValid: []bool{true, true},
|
||||
}
|
||||
|
||||
updateIdxInUpsert := []int{0, 1}
|
||||
upsertIDs := &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{10, 20}}}}
|
||||
existPKToIndex := map[interface{}]int{int64(10): 0, int64(20): 1}
|
||||
|
||||
result := rebuildNullableVectorFieldData(ctx, updateIdxInUpsert, upsertIDs, existPKToIndex)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, []bool{true, true}, result.ValidData)
|
||||
assert.Equal(t, []float32{1, 2, 3, 4, 5, 6, 7, 8}, result.GetVectors().GetFloatVector().Data)
|
||||
})
|
||||
|
||||
t.Run("nil source field", func(t *testing.T) {
|
||||
ctx := &nullableVectorMergeContext{
|
||||
hasUpsertData: false,
|
||||
existField: nil,
|
||||
mergedValid: []bool{true},
|
||||
}
|
||||
|
||||
result := rebuildNullableVectorFieldData(ctx, []int{0}, nil, nil)
|
||||
result, err := GenNullableFieldData(field, upsertIDSize)
|
||||
assert.Error(t, err)
|
||||
assert.Nil(t, result)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -57,6 +57,9 @@ type rerankInputs struct {
|
||||
|
||||
// There is only fieldId in schemapb.SearchResultData, but no fieldName
|
||||
inputFieldIds []int64
|
||||
|
||||
// idxComputers for computing correct field indices for nullable vectors
|
||||
idxComputers []*typeutil.FieldDataIdxComputer
|
||||
}
|
||||
|
||||
func organizeFieldIdData(multipSearchResultData []*schemapb.SearchResultData, inputFieldIds []int64) ([]map[int64]*schemapb.FieldData, error) {
|
||||
@@ -119,14 +122,21 @@ func newRerankInputs(multipSearchResultData []*schemapb.SearchResultData, inputF
|
||||
start += size
|
||||
}
|
||||
}
|
||||
idxComputers := make([]*typeutil.FieldDataIdxComputer, len(multipSearchResultData))
|
||||
for i, srd := range multipSearchResultData {
|
||||
if srd != nil {
|
||||
idxComputers[i] = typeutil.NewFieldDataIdxComputer(srd.GetFieldsData())
|
||||
}
|
||||
}
|
||||
|
||||
if isGrouping {
|
||||
idGroup, err := genIdGroupingMap(multipSearchResultData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &rerankInputs{cols, idGroup, nq, multipSearchResultData, inputFieldIds}, nil
|
||||
return &rerankInputs{cols, idGroup, nq, multipSearchResultData, inputFieldIds, idxComputers}, nil
|
||||
}
|
||||
return &rerankInputs{cols, nil, nq, multipSearchResultData, inputFieldIds}, nil
|
||||
return &rerankInputs{cols, nil, nq, multipSearchResultData, inputFieldIds, idxComputers}, nil
|
||||
}
|
||||
|
||||
func (inputs *rerankInputs) numOfQueries() int64 {
|
||||
@@ -165,7 +175,13 @@ func appendResult[T PKType](inputs *rerankInputs, outputs *rerankOutputs, idScor
|
||||
if len(inputs.fieldData) > 0 && len(outputs.searchResultData.FieldsData) > 0 {
|
||||
for idx := range ids {
|
||||
loc := idScores.locations[idx]
|
||||
typeutil.AppendFieldData(outputs.searchResultData.FieldsData, inputs.fieldData[loc.batchIdx].GetFieldsData(), int64(loc.offset))
|
||||
fieldsData := inputs.fieldData[loc.batchIdx].GetFieldsData()
|
||||
rowIdx := int64(loc.offset)
|
||||
var fieldIdxs []int64
|
||||
if inputs.idxComputers[loc.batchIdx] != nil {
|
||||
fieldIdxs = inputs.idxComputers[loc.batchIdx].Compute(rowIdx)
|
||||
}
|
||||
typeutil.AppendFieldData(outputs.searchResultData.FieldsData, fieldsData, rowIdx, fieldIdxs...)
|
||||
}
|
||||
}
|
||||
switch any(ids).(type) {
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/metric"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
|
||||
)
|
||||
|
||||
func TestUtil(t *testing.T) {
|
||||
@@ -338,3 +339,158 @@ func (s *UtilSuite) TestAppendResultWithEmptyFieldsData() {
|
||||
// FieldsData should still be empty
|
||||
s.Equal(0, len(outputs.searchResultData.FieldsData))
|
||||
}
|
||||
|
||||
// TestAppendResultWithNullableSparseVector tests appendResult with nullable sparse vector fields.
|
||||
// This ensures that idxComputers correctly maps row indices to data indices for nullable vectors.
|
||||
func (s *UtilSuite) TestAppendResultWithNullableSparseVector() {
|
||||
// Create sparse vector data: 3 rows where row 1 (middle) is null
|
||||
// ValidData: [true, false, true] means rows 0 and 2 have data, row 1 is null
|
||||
// Contents only has 2 entries (for the non-null rows at indices 0 and 2)
|
||||
sparseContent0 := []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x3f} // sparse vector data for row 0
|
||||
sparseContent2 := []byte{0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40} // sparse vector data for row 2
|
||||
|
||||
inputs := &rerankInputs{
|
||||
fieldData: []*schemapb.SearchResultData{
|
||||
{
|
||||
Ids: &schemapb.IDs{
|
||||
IdField: &schemapb.IDs_IntId{
|
||||
IntId: &schemapb.LongArray{
|
||||
Data: []int64{1, 2, 3}, // 3 rows
|
||||
},
|
||||
},
|
||||
},
|
||||
FieldsData: []*schemapb.FieldData{
|
||||
{
|
||||
Type: schemapb.DataType_SparseFloatVector,
|
||||
FieldName: "sparse_vec",
|
||||
FieldId: 101,
|
||||
ValidData: []bool{true, false, true}, // Row 1 is null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 700,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 700,
|
||||
Contents: [][]byte{sparseContent0, sparseContent2}, // Only 2 entries
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
nq: 1,
|
||||
}
|
||||
|
||||
// Initialize idxComputers for the inputs
|
||||
inputs.idxComputers = make([]*typeutil.FieldDataIdxComputer, len(inputs.fieldData))
|
||||
for i, srd := range inputs.fieldData {
|
||||
if srd != nil {
|
||||
inputs.idxComputers[i] = typeutil.NewFieldDataIdxComputer(srd.GetFieldsData())
|
||||
}
|
||||
}
|
||||
|
||||
searchParams := &SearchParams{limit: 10}
|
||||
outputs := newRerankOutputs(inputs, searchParams)
|
||||
|
||||
// Select rows in reverse order: row 2 first, then row 0
|
||||
// Row 2 is valid (should get sparseContent2), Row 0 is valid (should get sparseContent0)
|
||||
idScores := &IDScores[int64]{
|
||||
ids: []int64{3, 1}, // IDs for rows 2 and 0
|
||||
scores: []float32{0.9, 0.8}, // scores
|
||||
locations: []IDLoc{{batchIdx: 0, offset: 2}, {batchIdx: 0, offset: 0}}, // row indices
|
||||
}
|
||||
|
||||
// This should NOT panic - the bug was that without idxComputers,
|
||||
// it would use row index directly to access Contents, causing index out of range
|
||||
s.NotPanics(func() {
|
||||
appendResult(inputs, outputs, idScores)
|
||||
})
|
||||
|
||||
// Verify results
|
||||
s.Equal(int64(2), outputs.searchResultData.Topks[0])
|
||||
s.Equal([]int64{3, 1}, outputs.searchResultData.Ids.GetIntId().Data)
|
||||
|
||||
// Verify sparse vector data is correctly copied
|
||||
resultSparse := outputs.searchResultData.FieldsData[0]
|
||||
s.Equal(schemapb.DataType_SparseFloatVector, resultSparse.GetType())
|
||||
|
||||
// ValidData should reflect the selected rows: both row 2 and row 0 are valid
|
||||
s.Equal([]bool{true, true}, resultSparse.GetValidData())
|
||||
|
||||
// Contents should have 2 entries in the correct order
|
||||
contents := resultSparse.GetVectors().GetSparseFloatVector().GetContents()
|
||||
s.Len(contents, 2)
|
||||
s.Equal(sparseContent2, contents[0]) // Row 2's data comes first
|
||||
s.Equal(sparseContent0, contents[1]) // Row 0's data comes second
|
||||
}
|
||||
|
||||
// TestAppendResultWithNullableSparseVectorSelectingNullRow tests appendResult when selecting a null row.
|
||||
func (s *UtilSuite) TestAppendResultWithNullableSparseVectorSelectingNullRow() {
|
||||
sparseContent0 := []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x3f}
|
||||
|
||||
inputs := &rerankInputs{
|
||||
fieldData: []*schemapb.SearchResultData{
|
||||
{
|
||||
Ids: &schemapb.IDs{
|
||||
IdField: &schemapb.IDs_IntId{
|
||||
IntId: &schemapb.LongArray{
|
||||
Data: []int64{1, 2, 3},
|
||||
},
|
||||
},
|
||||
},
|
||||
FieldsData: []*schemapb.FieldData{
|
||||
{
|
||||
Type: schemapb.DataType_SparseFloatVector,
|
||||
FieldName: "sparse_vec",
|
||||
FieldId: 101,
|
||||
ValidData: []bool{true, false, false}, // Only row 0 has data
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 700,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 700,
|
||||
Contents: [][]byte{sparseContent0}, // Only 1 entry
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
nq: 1,
|
||||
}
|
||||
|
||||
inputs.idxComputers = make([]*typeutil.FieldDataIdxComputer, len(inputs.fieldData))
|
||||
for i, srd := range inputs.fieldData {
|
||||
if srd != nil {
|
||||
inputs.idxComputers[i] = typeutil.NewFieldDataIdxComputer(srd.GetFieldsData())
|
||||
}
|
||||
}
|
||||
|
||||
searchParams := &SearchParams{limit: 10}
|
||||
outputs := newRerankOutputs(inputs, searchParams)
|
||||
|
||||
// Select row 1 (null) and row 0 (valid)
|
||||
idScores := &IDScores[int64]{
|
||||
ids: []int64{2, 1},
|
||||
scores: []float32{0.9, 0.8},
|
||||
locations: []IDLoc{{batchIdx: 0, offset: 1}, {batchIdx: 0, offset: 0}},
|
||||
}
|
||||
|
||||
s.NotPanics(func() {
|
||||
appendResult(inputs, outputs, idScores)
|
||||
})
|
||||
|
||||
// Verify ValidData: row 1 is null, row 0 is valid
|
||||
resultSparse := outputs.searchResultData.FieldsData[0]
|
||||
s.Equal([]bool{false, true}, resultSparse.GetValidData())
|
||||
|
||||
// Contents should have 1 entry (only for the valid row)
|
||||
contents := resultSparse.GetVectors().GetSparseFloatVector().GetContents()
|
||||
s.Len(contents, 1)
|
||||
s.Equal(sparseContent0, contents[0])
|
||||
}
|
||||
|
||||
@@ -786,6 +786,12 @@ func PrepareResultFieldData(sample []*schemapb.FieldData, topK int64) []*schemap
|
||||
Data: make([][]byte, 0, topK),
|
||||
},
|
||||
}
|
||||
case *schemapb.ScalarField_GeometryWktData:
|
||||
scalar.Scalars.Data = &schemapb.ScalarField_GeometryWktData{
|
||||
GeometryWktData: &schemapb.GeometryWktArray{
|
||||
Data: make([]string, 0, topK),
|
||||
},
|
||||
}
|
||||
case *schemapb.ScalarField_ArrayData:
|
||||
scalar.Scalars.Data = &schemapb.ScalarField_ArrayData{
|
||||
ArrayData: &schemapb.ArrayArray{
|
||||
@@ -1187,6 +1193,232 @@ func AppendFieldData(dst, src []*schemapb.FieldData, idx int64, fieldIdxs ...int
|
||||
return
|
||||
}
|
||||
|
||||
func AppendFieldDataByColumn(dst, src *schemapb.FieldData, dataIndices []int64, rowIndices ...[]int64) {
|
||||
if dst == nil || src == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle ValidData: use rowIndices if provided, otherwise use dataIndices
|
||||
if len(src.GetValidData()) > 0 {
|
||||
validIndices := dataIndices
|
||||
if len(rowIndices) > 0 {
|
||||
validIndices = rowIndices[0]
|
||||
}
|
||||
if dst.ValidData == nil {
|
||||
dst.ValidData = make([]bool, 0, len(validIndices))
|
||||
}
|
||||
for _, idx := range validIndices {
|
||||
dst.ValidData = append(dst.ValidData, src.ValidData[int(idx)])
|
||||
}
|
||||
}
|
||||
|
||||
// Return early if no data to copy
|
||||
if len(dataIndices) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
switch srcField := src.Field.(type) {
|
||||
case *schemapb.FieldData_Scalars:
|
||||
if dst.GetScalars() == nil {
|
||||
dst.Field = &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{},
|
||||
}
|
||||
}
|
||||
dstScalar := dst.GetScalars()
|
||||
switch srcScalar := srcField.Scalars.Data.(type) {
|
||||
case *schemapb.ScalarField_BoolData:
|
||||
if dstScalar.GetBoolData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_BoolData{
|
||||
BoolData: &schemapb.BoolArray{Data: make([]bool, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetBoolData().Data = append(dstScalar.GetBoolData().Data, srcScalar.BoolData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_IntData:
|
||||
if dstScalar.GetIntData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: make([]int32, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetIntData().Data = append(dstScalar.GetIntData().Data, srcScalar.IntData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_LongData:
|
||||
if dstScalar.GetLongData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_LongData{
|
||||
LongData: &schemapb.LongArray{Data: make([]int64, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetLongData().Data = append(dstScalar.GetLongData().Data, srcScalar.LongData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_FloatData:
|
||||
if dstScalar.GetFloatData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_FloatData{
|
||||
FloatData: &schemapb.FloatArray{Data: make([]float32, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetFloatData().Data = append(dstScalar.GetFloatData().Data, srcScalar.FloatData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_DoubleData:
|
||||
if dstScalar.GetDoubleData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_DoubleData{
|
||||
DoubleData: &schemapb.DoubleArray{Data: make([]float64, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetDoubleData().Data = append(dstScalar.GetDoubleData().Data, srcScalar.DoubleData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_StringData:
|
||||
if dstScalar.GetStringData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_StringData{
|
||||
StringData: &schemapb.StringArray{Data: make([]string, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetStringData().Data = append(dstScalar.GetStringData().Data, srcScalar.StringData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_ArrayData:
|
||||
if dstScalar.GetArrayData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_ArrayData{
|
||||
ArrayData: &schemapb.ArrayArray{
|
||||
Data: make([]*schemapb.ScalarField, 0, len(dataIndices)),
|
||||
ElementType: srcScalar.ArrayData.ElementType,
|
||||
},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetArrayData().Data = append(dstScalar.GetArrayData().Data, srcScalar.ArrayData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_JsonData:
|
||||
if dstScalar.GetJsonData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_JsonData{
|
||||
JsonData: &schemapb.JSONArray{Data: make([][]byte, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetJsonData().Data = append(dstScalar.GetJsonData().Data, srcScalar.JsonData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_GeometryData:
|
||||
if dstScalar.GetGeometryData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_GeometryData{
|
||||
GeometryData: &schemapb.GeometryArray{Data: make([][]byte, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetGeometryData().Data = append(dstScalar.GetGeometryData().Data, srcScalar.GeometryData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_GeometryWktData:
|
||||
if dstScalar.GetGeometryWktData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_GeometryWktData{
|
||||
GeometryWktData: &schemapb.GeometryWktArray{Data: make([]string, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetGeometryWktData().Data = append(dstScalar.GetGeometryWktData().Data, srcScalar.GeometryWktData.Data[idx])
|
||||
}
|
||||
case *schemapb.ScalarField_TimestamptzData:
|
||||
if dstScalar.GetTimestamptzData() == nil {
|
||||
dstScalar.Data = &schemapb.ScalarField_TimestamptzData{
|
||||
TimestamptzData: &schemapb.TimestamptzArray{Data: make([]int64, 0, len(dataIndices))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstScalar.GetTimestamptzData().Data = append(dstScalar.GetTimestamptzData().Data, srcScalar.TimestamptzData.Data[idx])
|
||||
}
|
||||
}
|
||||
case *schemapb.FieldData_Vectors:
|
||||
dim := srcField.Vectors.Dim
|
||||
if dst.GetVectors() == nil {
|
||||
dst.Field = &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{Dim: dim},
|
||||
}
|
||||
}
|
||||
dstVector := dst.GetVectors()
|
||||
switch srcVector := srcField.Vectors.Data.(type) {
|
||||
case *schemapb.VectorField_BinaryVector:
|
||||
if dstVector.GetBinaryVector() == nil {
|
||||
dstVector.Data = &schemapb.VectorField_BinaryVector{
|
||||
BinaryVector: make([]byte, 0, len(dataIndices)*int(dim/8)),
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
start := idx * (dim / 8)
|
||||
end := (idx + 1) * (dim / 8)
|
||||
dstVector.Data.(*schemapb.VectorField_BinaryVector).BinaryVector = append(
|
||||
dstVector.Data.(*schemapb.VectorField_BinaryVector).BinaryVector,
|
||||
srcVector.BinaryVector[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_FloatVector:
|
||||
if dstVector.GetFloatVector() == nil {
|
||||
dstVector.Data = &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: make([]float32, 0, len(dataIndices)*int(dim))},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
start := idx * dim
|
||||
end := (idx + 1) * dim
|
||||
dstVector.GetFloatVector().Data = append(dstVector.GetFloatVector().Data, srcVector.FloatVector.Data[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_Float16Vector:
|
||||
if dstVector.GetFloat16Vector() == nil {
|
||||
dstVector.Data = &schemapb.VectorField_Float16Vector{
|
||||
Float16Vector: make([]byte, 0, len(dataIndices)*int(dim*2)),
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
start := idx * (dim * 2)
|
||||
end := (idx + 1) * (dim * 2)
|
||||
dstVector.Data.(*schemapb.VectorField_Float16Vector).Float16Vector = append(
|
||||
dstVector.Data.(*schemapb.VectorField_Float16Vector).Float16Vector,
|
||||
srcVector.Float16Vector[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_Bfloat16Vector:
|
||||
if dstVector.GetBfloat16Vector() == nil {
|
||||
dstVector.Data = &schemapb.VectorField_Bfloat16Vector{
|
||||
Bfloat16Vector: make([]byte, 0, len(dataIndices)*int(dim*2)),
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
start := idx * (dim * 2)
|
||||
end := (idx + 1) * (dim * 2)
|
||||
dstVector.Data.(*schemapb.VectorField_Bfloat16Vector).Bfloat16Vector = append(
|
||||
dstVector.Data.(*schemapb.VectorField_Bfloat16Vector).Bfloat16Vector,
|
||||
srcVector.Bfloat16Vector[start:end]...)
|
||||
}
|
||||
case *schemapb.VectorField_SparseFloatVector:
|
||||
if dstVector.GetSparseFloatVector() == nil {
|
||||
dstVector.Data = &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: srcVector.SparseFloatVector.Dim,
|
||||
Contents: make([][]byte, 0, len(dataIndices)),
|
||||
},
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
dstVector.GetSparseFloatVector().Contents = append(
|
||||
dstVector.GetSparseFloatVector().Contents,
|
||||
srcVector.SparseFloatVector.Contents[idx])
|
||||
}
|
||||
case *schemapb.VectorField_Int8Vector:
|
||||
if dstVector.GetInt8Vector() == nil {
|
||||
dstVector.Data = &schemapb.VectorField_Int8Vector{
|
||||
Int8Vector: make([]byte, 0, len(dataIndices)*int(dim)),
|
||||
}
|
||||
}
|
||||
for _, idx := range dataIndices {
|
||||
start := idx * dim
|
||||
end := (idx + 1) * dim
|
||||
dstVector.Data.(*schemapb.VectorField_Int8Vector).Int8Vector = append(
|
||||
dstVector.Data.(*schemapb.VectorField_Int8Vector).Int8Vector,
|
||||
srcVector.Int8Vector[start:end]...)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// DeleteFieldData delete fields data appended last time
|
||||
func DeleteFieldData(dst []*schemapb.FieldData) {
|
||||
for i, fieldData := range dst {
|
||||
@@ -1466,6 +1698,180 @@ func UpdateFieldData(base, update []*schemapb.FieldData, baseIdx, updateIdx int6
|
||||
return nil
|
||||
}
|
||||
|
||||
func UpdateFieldDataByColumn(base, update *schemapb.FieldData, baseIndices, updateIndices []int64) error {
|
||||
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))
|
||||
}
|
||||
|
||||
// Handle ValidData
|
||||
if len(update.GetValidData()) > 0 && len(base.GetValidData()) > 0 {
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
base.ValidData[baseIdx] = update.ValidData[updateIdx]
|
||||
}
|
||||
}
|
||||
|
||||
switch baseField := base.Field.(type) {
|
||||
case *schemapb.FieldData_Scalars:
|
||||
updateField := update.Field.(*schemapb.FieldData_Scalars)
|
||||
baseScalar := baseField.Scalars
|
||||
updateScalar := updateField.Scalars
|
||||
|
||||
switch baseScalar.Data.(type) {
|
||||
case *schemapb.ScalarField_BoolData:
|
||||
baseData := baseScalar.GetBoolData().Data
|
||||
updateData := updateScalar.GetBoolData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_IntData:
|
||||
baseData := baseScalar.GetIntData().Data
|
||||
updateData := updateScalar.GetIntData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_LongData:
|
||||
baseData := baseScalar.GetLongData().Data
|
||||
updateData := updateScalar.GetLongData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_FloatData:
|
||||
baseData := baseScalar.GetFloatData().Data
|
||||
updateData := updateScalar.GetFloatData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_DoubleData:
|
||||
baseData := baseScalar.GetDoubleData().Data
|
||||
updateData := updateScalar.GetDoubleData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_StringData:
|
||||
baseData := baseScalar.GetStringData().Data
|
||||
updateData := updateScalar.GetStringData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_ArrayData:
|
||||
baseData := baseScalar.GetArrayData().Data
|
||||
updateData := updateScalar.GetArrayData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_JsonData:
|
||||
baseData := baseScalar.GetJsonData().Data
|
||||
updateData := updateScalar.GetJsonData().Data
|
||||
if base.GetIsDynamic() {
|
||||
// dynamic field is a json with only 1 level nested struct,
|
||||
// so we need to unmarshal and iterate updateData's key value, and update the baseData's key value
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
var baseMap map[string]interface{}
|
||||
var updateMap map[string]interface{}
|
||||
// unmarshal base and update
|
||||
if err := json.Unmarshal(baseData[baseIdx], &baseMap); err != nil {
|
||||
return fmt.Errorf("failed to unmarshal base json: %v", err)
|
||||
}
|
||||
if err := json.Unmarshal(updateData[updateIdx], &updateMap); err != nil {
|
||||
return fmt.Errorf("failed to unmarshal update json: %v", err)
|
||||
}
|
||||
// merge
|
||||
for k, v := range updateMap {
|
||||
baseMap[k] = v
|
||||
}
|
||||
// marshal back
|
||||
newJSON, err := json.Marshal(baseMap)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal merged json: %v", err)
|
||||
}
|
||||
baseData[baseIdx] = newJSON
|
||||
}
|
||||
} else {
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
}
|
||||
case *schemapb.ScalarField_GeometryData:
|
||||
baseData := baseScalar.GetGeometryData().Data
|
||||
updateData := updateScalar.GetGeometryData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_GeometryWktData:
|
||||
baseData := baseScalar.GetGeometryWktData().Data
|
||||
updateData := updateScalar.GetGeometryWktData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.ScalarField_TimestamptzData:
|
||||
baseData := baseScalar.GetTimestamptzData().Data
|
||||
updateData := updateScalar.GetTimestamptzData().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
}
|
||||
case *schemapb.FieldData_Vectors:
|
||||
updateField := update.Field.(*schemapb.FieldData_Vectors)
|
||||
baseVector := baseField.Vectors
|
||||
updateVector := updateField.Vectors
|
||||
dim := baseVector.Dim
|
||||
|
||||
switch baseVector.Data.(type) {
|
||||
case *schemapb.VectorField_BinaryVector:
|
||||
baseData := baseVector.GetBinaryVector()
|
||||
updateData := updateVector.GetBinaryVector()
|
||||
elemSize := dim / 8
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
copy(baseData[baseIdx*elemSize:(baseIdx+1)*elemSize], updateData[updateIdx*elemSize:(updateIdx+1)*elemSize])
|
||||
}
|
||||
case *schemapb.VectorField_FloatVector:
|
||||
baseData := baseVector.GetFloatVector().Data
|
||||
updateData := updateVector.GetFloatVector().Data
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
copy(baseData[baseIdx*dim:(baseIdx+1)*dim], updateData[updateIdx*dim:(updateIdx+1)*dim])
|
||||
}
|
||||
case *schemapb.VectorField_Float16Vector:
|
||||
baseData := baseVector.GetFloat16Vector()
|
||||
updateData := updateVector.GetFloat16Vector()
|
||||
elemSize := dim * 2
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
copy(baseData[baseIdx*elemSize:(baseIdx+1)*elemSize], updateData[updateIdx*elemSize:(updateIdx+1)*elemSize])
|
||||
}
|
||||
case *schemapb.VectorField_Bfloat16Vector:
|
||||
baseData := baseVector.GetBfloat16Vector()
|
||||
updateData := updateVector.GetBfloat16Vector()
|
||||
elemSize := dim * 2
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
copy(baseData[baseIdx*elemSize:(baseIdx+1)*elemSize], updateData[updateIdx*elemSize:(updateIdx+1)*elemSize])
|
||||
}
|
||||
case *schemapb.VectorField_SparseFloatVector:
|
||||
baseData := baseVector.GetSparseFloatVector().Contents
|
||||
updateData := updateVector.GetSparseFloatVector().Contents
|
||||
for i, baseIdx := range baseIndices {
|
||||
baseData[baseIdx] = updateData[updateIndices[i]]
|
||||
}
|
||||
case *schemapb.VectorField_Int8Vector:
|
||||
baseData := baseVector.GetInt8Vector()
|
||||
updateData := updateVector.GetInt8Vector()
|
||||
for i, baseIdx := range baseIndices {
|
||||
updateIdx := updateIndices[i]
|
||||
copy(baseData[baseIdx*dim:(baseIdx+1)*dim], updateData[updateIdx*dim:(updateIdx+1)*dim])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// MergeFieldData appends fields data to dst
|
||||
func MergeFieldData(dst []*schemapb.FieldData, src []*schemapb.FieldData) error {
|
||||
fieldID2Data := make(map[int64]*schemapb.FieldData)
|
||||
|
||||
@@ -5203,3 +5203,349 @@ func TestValidateExternalCollectionSchema(t *testing.T) {
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestAppendFieldDataByColumn(t *testing.T) {
|
||||
t.Run("nil dst or src", func(t *testing.T) {
|
||||
src := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
// nil dst
|
||||
AppendFieldDataByColumn(nil, src, []int64{0, 1})
|
||||
// nil src
|
||||
AppendFieldDataByColumn(dst, nil, []int64{0, 1})
|
||||
// empty dataIndices
|
||||
AppendFieldDataByColumn(dst, src, []int64{})
|
||||
assert.Nil(t, dst.GetScalars())
|
||||
})
|
||||
|
||||
t.Run("scalar int32", func(t *testing.T) {
|
||||
src := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{10, 20, 30, 40, 50}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
AppendFieldDataByColumn(dst, src, []int64{1, 3})
|
||||
assert.Equal(t, []int32{20, 40}, dst.GetScalars().GetIntData().Data)
|
||||
})
|
||||
|
||||
t.Run("scalar with ValidData", func(t *testing.T) {
|
||||
src := &schemapb.FieldData{
|
||||
ValidData: []bool{true, false, true, false, true},
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{
|
||||
LongData: &schemapb.LongArray{Data: []int64{100, 200, 300, 400, 500}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
AppendFieldDataByColumn(dst, src, []int64{0, 2, 4})
|
||||
assert.Equal(t, []int64{100, 300, 500}, dst.GetScalars().GetLongData().Data)
|
||||
assert.Equal(t, []bool{true, true, true}, dst.ValidData)
|
||||
})
|
||||
|
||||
t.Run("string data", func(t *testing.T) {
|
||||
src := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_StringData{
|
||||
StringData: &schemapb.StringArray{Data: []string{"a", "b", "c", "d"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
AppendFieldDataByColumn(dst, src, []int64{1, 2})
|
||||
assert.Equal(t, []string{"b", "c"}, dst.GetScalars().GetStringData().Data)
|
||||
})
|
||||
|
||||
t.Run("json data", func(t *testing.T) {
|
||||
src := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_JsonData{
|
||||
JsonData: &schemapb.JSONArray{Data: [][]byte{[]byte(`{"a":1}`), []byte(`{"b":2}`)}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
AppendFieldDataByColumn(dst, src, []int64{0})
|
||||
assert.Equal(t, [][]byte{[]byte(`{"a":1}`)}, dst.GetScalars().GetJsonData().Data)
|
||||
})
|
||||
|
||||
t.Run("float vector", func(t *testing.T) {
|
||||
dim := int64(4)
|
||||
src := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: []float32{
|
||||
1, 2, 3, 4,
|
||||
5, 6, 7, 8,
|
||||
9, 10, 11, 12,
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
AppendFieldDataByColumn(dst, src, []int64{0, 2})
|
||||
assert.Equal(t, []float32{1, 2, 3, 4, 9, 10, 11, 12}, dst.GetVectors().GetFloatVector().Data)
|
||||
assert.Equal(t, dim, dst.GetVectors().GetDim())
|
||||
})
|
||||
|
||||
t.Run("float vector with ValidData - all null case", func(t *testing.T) {
|
||||
dim := int64(4)
|
||||
src := &schemapb.FieldData{
|
||||
ValidData: []bool{false, false, false}, // all null
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: []float32{}}, // COMPRESSED: no data
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
// dataIndices is empty (no non-null data), but rowIndices has all rows
|
||||
dataIndices := []int64{}
|
||||
rowIndices := []int64{0, 1, 2}
|
||||
AppendFieldDataByColumn(dst, src, dataIndices, rowIndices)
|
||||
// ValidData should be copied even though data is empty
|
||||
assert.Equal(t, []bool{false, false, false}, dst.ValidData)
|
||||
// No vector data copied
|
||||
assert.Nil(t, dst.GetVectors())
|
||||
})
|
||||
|
||||
t.Run("sparse float vector", func(t *testing.T) {
|
||||
src := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: 100,
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 100,
|
||||
Contents: [][]byte{{1, 2}, {3, 4}, {5, 6}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
dst := &schemapb.FieldData{}
|
||||
AppendFieldDataByColumn(dst, src, []int64{0, 2})
|
||||
assert.Equal(t, [][]byte{{1, 2}, {5, 6}}, dst.GetVectors().GetSparseFloatVector().Contents)
|
||||
})
|
||||
}
|
||||
|
||||
func TestUpdateFieldDataByColumn(t *testing.T) {
|
||||
t.Run("nil base or update", func(t *testing.T) {
|
||||
base := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
update := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{10, 20}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// nil base
|
||||
err := UpdateFieldDataByColumn(nil, update, []int64{0}, []int64{0})
|
||||
assert.NoError(t, err)
|
||||
// nil update
|
||||
err = UpdateFieldDataByColumn(base, nil, []int64{0}, []int64{0})
|
||||
assert.NoError(t, err)
|
||||
// empty indices
|
||||
err = UpdateFieldDataByColumn(base, update, []int64{}, []int64{})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, []int32{1, 2, 3}, base.GetScalars().GetIntData().Data)
|
||||
})
|
||||
|
||||
t.Run("indices length mismatch", func(t *testing.T) {
|
||||
base := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
update := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{10, 20}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
err := UpdateFieldDataByColumn(base, update, []int64{0, 1}, []int64{0})
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "length mismatch")
|
||||
})
|
||||
|
||||
t.Run("scalar int32", func(t *testing.T) {
|
||||
base := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{1, 2, 3, 4, 5}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
update := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
IntData: &schemapb.IntArray{Data: []int32{100, 200, 300}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
err := UpdateFieldDataByColumn(base, update, []int64{1, 3}, []int64{0, 2})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, []int32{1, 100, 3, 300, 5}, base.GetScalars().GetIntData().Data)
|
||||
})
|
||||
|
||||
t.Run("scalar with ValidData", func(t *testing.T) {
|
||||
base := &schemapb.FieldData{
|
||||
ValidData: []bool{true, true, true, true},
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{
|
||||
LongData: &schemapb.LongArray{Data: []int64{10, 20, 30, 40}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
update := &schemapb.FieldData{
|
||||
ValidData: []bool{false, true},
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{
|
||||
LongData: &schemapb.LongArray{Data: []int64{100, 200}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
err := UpdateFieldDataByColumn(base, update, []int64{1, 2}, []int64{0, 1})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, []int64{10, 100, 200, 40}, base.GetScalars().GetLongData().Data)
|
||||
assert.Equal(t, []bool{true, false, true, true}, base.ValidData)
|
||||
})
|
||||
|
||||
t.Run("string data", func(t *testing.T) {
|
||||
base := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_StringData{
|
||||
StringData: &schemapb.StringArray{Data: []string{"a", "b", "c"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
update := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Scalars{
|
||||
Scalars: &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_StringData{
|
||||
StringData: &schemapb.StringArray{Data: []string{"X", "Y"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
err := UpdateFieldDataByColumn(base, update, []int64{0, 2}, []int64{0, 1})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, []string{"X", "b", "Y"}, base.GetScalars().GetStringData().Data)
|
||||
})
|
||||
|
||||
t.Run("float vector", func(t *testing.T) {
|
||||
dim := int64(4)
|
||||
base := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: []float32{
|
||||
1, 2, 3, 4,
|
||||
5, 6, 7, 8,
|
||||
9, 10, 11, 12,
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
update := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Dim: dim,
|
||||
Data: &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: []float32{
|
||||
100, 200, 300, 400,
|
||||
500, 600, 700, 800,
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
err := UpdateFieldDataByColumn(base, update, []int64{0, 2}, []int64{0, 1})
|
||||
assert.NoError(t, err)
|
||||
expected := []float32{100, 200, 300, 400, 5, 6, 7, 8, 500, 600, 700, 800}
|
||||
assert.Equal(t, expected, base.GetVectors().GetFloatVector().Data)
|
||||
})
|
||||
|
||||
t.Run("sparse float vector", func(t *testing.T) {
|
||||
base := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 100,
|
||||
Contents: [][]byte{{1}, {2}, {3}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
update := &schemapb.FieldData{
|
||||
Field: &schemapb.FieldData_Vectors{
|
||||
Vectors: &schemapb.VectorField{
|
||||
Data: &schemapb.VectorField_SparseFloatVector{
|
||||
SparseFloatVector: &schemapb.SparseFloatArray{
|
||||
Dim: 100,
|
||||
Contents: [][]byte{{10}, {20}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
err := UpdateFieldDataByColumn(base, update, []int64{0, 2}, []int64{0, 1})
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, [][]byte{{10}, {2}, {20}}, base.GetVectors().GetSparseFloatVector().Contents)
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user