admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2026-01-07 19:31:51 +08:00
Code Issues Packages Projects Releases Wiki Activity
milvus/internal/util/importutilv2/json/reader_test.go
marcelo-cjl 3c2cf2c066
feat: Add nullable vector support in import utility layer (#46142) 
related: #45993 

Add nullable vector support in import utility layer
    
Key changes:

ImportV2 util:
- Add nullable vector types (FloatVector, Float16Vector, BFloat16Vector,
BinaryVector, SparseFloatVector, Int8Vector) to
AppendNullableDefaultFieldsData()
- Add tests for nullable vector field data appending

CSV/JSON/Numpy readers:
- Add nullPercent parameter to test data generation for better null
coverage
- Mark vector fields as nullable in test schemas
- Add test cases for nullable vector field parsing
- Refactor tests to use loop-based approach with 0%, 50%, 100% null
percentages

Parquet field reader:
- Add ReadNullableBinaryData() for nullable
BinaryVector/Float16Vector/BFloat16Vector
- Add ReadNullableFloatVectorData() for nullable FloatVector
- Add ReadNullableSparseFloatVectorData() for nullable SparseFloatVector
- Add ReadNullableInt8VectorData() for nullable Int8Vector
- Add ReadNullableStructData() for generic nullable struct data
- Update Next() to use nullable read methods when field is nullable
- Add null data validation for non-nullable fields

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
- Core invariant: import must preserve per-row alignment and validity
for every field — nullable vector fields are expected to be encoded with
per-row validity masks and all readers/writers must emit arrays aligned
to original input rows (null entries represented explicitly).
- New feature & scope: adds end-to-end nullable-vector support in the
import utility layer — AppendNullableDefaultFieldsData in
internal/datanode/importv2/util.go now appends nil placeholders for
nullable vectors (FloatVector, Float16Vector, BFloat16Vector,
BinaryVector, SparseFloatVector, Int8Vector); parquet reader
(internal/util/importutilv2/parquet/field_reader.go) adds
ReadNullableBinaryData, ReadNullableFloatVectorData,
ReadNullableSparseFloatVectorData, ReadNullableInt8VectorData,
ReadNullableStructData and routes nullable branches to these helpers;
CSV/JSON/Numpy readers and test utilities updated to generate and
validate 0/50/100% null scenarios and mark vector fields as nullable in
test schemas.
- Logic removed / simplified: eliminates ad-hoc "parameter-invalid"
rejections for nullable vectors inside FieldReader.Next by centralizing
nullable handling into ReadNullable* helpers and shared validators
(getArrayDataNullable,
checkNullableVectorAlignWithDim/checkNullableVectorAligned), simplifying
control flow and removing scattered special-case checks.
- No data loss / no regression (concrete code paths): nulls are
preserved end-to-end — AppendNullableDefaultFieldsData explicitly
inserts nil entries per null row (datanode import append path);
ReadNullable*Data helpers return both data and []bool validity masks so
callers in field_reader.go and downstream readers receive exact per-row
validity; testutil.BuildSparseVectorData was extended to accept
validData so sparse vectors are materialized only for valid rows while
null rows are represented as missing. These concrete paths ensure null
rows are represented rather than dropped, preventing data loss or
behavioral regression.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: marcelo-cjl <marcelo.chen@zilliz.com>
2025-12-29 10:51:21 +08:00

525 lines
15 KiB
Go
Raw Blame History

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package json
import (
"context"
"encoding/json"
"math"
"testing"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/storage"
importcommon "github.com/milvus-io/milvus/internal/util/importutilv2/common"
"github.com/milvus-io/milvus/internal/util/nullutil"
"github.com/milvus-io/milvus/internal/util/testutil"
"github.com/milvus-io/milvus/pkg/v2/common"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
func init() {
paramtable.Init()
}
type ReaderSuite struct {
suite.Suite
numRows int
pkDataType schemapb.DataType
vecDataType schemapb.DataType
}
func (suite *ReaderSuite) SetupTest() {
// default suite params
suite.numRows = 100
suite.pkDataType = schemapb.DataType_Int64
suite.vecDataType = schemapb.DataType_FloatVector
}
func (suite *ReaderSuite) run(dataType schemapb.DataType, elemType schemapb.DataType, nullable bool, nullPercent int) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "pk",
IsPrimaryKey: true,
DataType: suite.pkDataType,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: "128",
},
},
},
{
FieldID: 101,
Name: "vec",
DataType: suite.vecDataType,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "8",
},
},
Nullable: nullable,
},
{
FieldID: 102,
Name: dataType.String(),
DataType: dataType,
ElementType: elemType,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: "128",
},
{
Key: common.MaxCapacityKey,
Value: "128",
},
},
Nullable: nullable,
},
},
}
if dataType == schemapb.DataType_VarChar {
// Add a BM25 function if data type is VarChar
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
FieldID: 103,
Name: "sparse",
DataType: schemapb.DataType_SparseFloatVector,
IsFunctionOutput: true,
})
schema.Functions = append(schema.Functions, &schemapb.FunctionSchema{
Id: 1000,
Name: "bm25",
Type: schemapb.FunctionType_BM25,
InputFieldIds: []int64{102},
InputFieldNames: []string{dataType.String()},
OutputFieldIds: []int64{103},
OutputFieldNames: []string{"sparse"},
})
}
insertData, err := testutil.CreateInsertData(schema, suite.numRows, nullPercent)
suite.NoError(err)
rows, err := testutil.CreateInsertDataRowsForJSON(schema, insertData)
suite.NoError(err)
jsonBytes, err := json.Marshal(rows)
suite.NoError(err)
cm := mocks.NewChunkManager(suite.T())
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
r := importcommon.NewMockReader(string(jsonBytes))
return r, nil
})
cm.EXPECT().Size(mock.Anything, "mockPath").Return(128, nil)
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
suite.NoError(err)
suite.NotNil(reader)
defer reader.Close()
size, err := reader.Size()
suite.NoError(err)
suite.Equal(int64(128), size)
size2, err := reader.Size() // size is cached
suite.NoError(err)
suite.Equal(size, size2)
checkFn := func(actualInsertData *storage.InsertData, offsetBegin, expectRows int) {
expectInsertData := insertData
for fieldID, data := range actualInsertData.Data {
suite.Equal(expectRows, data.RowNum())
for i := 0; i < expectRows; i++ {
expect := expectInsertData.Data[fieldID].GetRow(i + offsetBegin)
actual := data.GetRow(i)
suite.Equal(expect, actual)
}
}
}
res, err := reader.Read()
suite.NoError(err)
checkFn(res, 0, suite.numRows)
}
func (suite *ReaderSuite) runWithDefaultValue(dataType schemapb.DataType, elemType schemapb.DataType, oldFormat bool) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "pk",
IsPrimaryKey: true,
DataType: suite.pkDataType,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: "128",
},
},
},
{
FieldID: 101,
Name: "vec",
DataType: suite.vecDataType,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "8",
},
},
},
},
}
// here always set nullable==true just for test, insertData will store validData only nullable==true
// jsonBytes Marshal from rows, if expectInsertData is nulls and set default value
// actualData will be default_value
fieldSchema, err := testutil.CreateFieldWithDefaultValue(dataType, 102, true)
suite.NoError(err)
schema.Fields = append(schema.Fields, fieldSchema)
insertData, err := testutil.CreateInsertData(schema, suite.numRows)
suite.NoError(err)
rows, err := testutil.CreateInsertDataRowsForJSON(schema, insertData)
suite.NoError(err)
var jsonBytes []byte
if oldFormat {
oldRows := make(map[string]any)
oldRows["rows"] = rows
jsonBytes, err = json.Marshal(oldRows)
suite.NoError(err)
} else {
jsonBytes, err = json.Marshal(rows)
suite.NoError(err)
}
cm := mocks.NewChunkManager(suite.T())
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
r := importcommon.NewMockReader(string(jsonBytes))
return r, nil
})
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
suite.NoError(err)
checkFn := func(actualInsertData *storage.InsertData, offsetBegin, expectRows int) {
expectInsertData := insertData
for fieldID, data := range actualInsertData.Data {
suite.Equal(expectRows, data.RowNum())
for i := 0; i < expectRows; i++ {
expect := expectInsertData.Data[fieldID].GetRow(i + offsetBegin)
actual := data.GetRow(i)
if expect == nil {
expect, err = nullutil.GetDefaultValue(fieldSchema)
suite.NoError(err)
}
suite.Equal(expect, actual)
}
}
}
res, err := reader.Read()
suite.NoError(err)
checkFn(res, 0, suite.numRows)
}
func (suite *ReaderSuite) TestReadScalarFields() {
elementTypes := []schemapb.DataType{
schemapb.DataType_Bool,
schemapb.DataType_Int8,
schemapb.DataType_Int16,
schemapb.DataType_Int32,
schemapb.DataType_Int64,
schemapb.DataType_Float,
schemapb.DataType_Double,
schemapb.DataType_VarChar,
}
scalarTypes := append(elementTypes, []schemapb.DataType{schemapb.DataType_JSON, schemapb.DataType_Array}...)
for _, dataType := range scalarTypes {
if dataType == schemapb.DataType_Array {
for _, elementType := range elementTypes {
suite.run(dataType, elementType, false, 0)
for _, nullPercent := range []int{0, 50, 100} {
suite.run(dataType, elementType, true, nullPercent)
}
}
} else {
suite.run(dataType, schemapb.DataType_None, false, 0)
for _, nullPercent := range []int{0, 50, 100} {
suite.run(dataType, schemapb.DataType_None, true, nullPercent)
}
}
}
}
func (suite *ReaderSuite) TestReadScalarFieldsWithDefaultValue() {
scalarTypes := []schemapb.DataType{
schemapb.DataType_Bool,
schemapb.DataType_Int8,
schemapb.DataType_Int16,
schemapb.DataType_Int32,
schemapb.DataType_Int64,
schemapb.DataType_Float,
schemapb.DataType_Double,
schemapb.DataType_VarChar,
}
for _, dataType := range scalarTypes {
suite.runWithDefaultValue(dataType, schemapb.DataType_None, false)
suite.runWithDefaultValue(dataType, schemapb.DataType_None, true)
}
}
func (suite *ReaderSuite) TestStringPK() {
suite.pkDataType = schemapb.DataType_VarChar
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false, 0)
for _, nullPercent := range []int{0, 50, 100} {
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, true, nullPercent)
}
}
func (suite *ReaderSuite) TestVector() {
dataTypes := []schemapb.DataType{
schemapb.DataType_BinaryVector,
schemapb.DataType_FloatVector,
schemapb.DataType_Float16Vector,
schemapb.DataType_BFloat16Vector,
schemapb.DataType_SparseFloatVector,
schemapb.DataType_Int8Vector,
}
for _, dataType := range dataTypes {
suite.vecDataType = dataType
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false, 0)
for _, nullPercent := range []int{0, 50, 100} {
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, true, nullPercent)
}
}
}
func (suite *ReaderSuite) TestDecodeError() {
testDecode := func(jsonContent string, isLinesFormat bool, isOldFormat bool, ioErr error, initErr bool, decodeErr bool) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "pk",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
},
}
cm := mocks.NewChunkManager(suite.T())
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
r := importcommon.NewMockReader(jsonContent)
return r, nil
})
var reader *reader
var err error
if isLinesFormat {
reader, err = NewLinesReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
} else {
reader, err = NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
}
if initErr {
suite.Error(err)
} else {
suite.NoError(err)
}
if err == nil {
reader.isLinesFormat = isLinesFormat
reader.isOldFormat = isOldFormat
_, err = reader.Read()
if decodeErr {
suite.Error(err)
} else {
suite.NoError(err)
}
}
}
testOldFormatDecode := func(jsonContent string, ioErr error, initErr bool, decodeErr bool) {
testDecode(jsonContent, false, true, ioErr, initErr, decodeErr)
}
testListFormatDecode := func(jsonContent string, decodeErr bool) {
testDecode(jsonContent, false, false, nil, false, decodeErr)
}
testLinesDecode := func(jsonContent string, decodeErr bool) {
testDecode(jsonContent, true, false, nil, false, decodeErr)
}
// illegal formats
testOldFormatDecode("", merr.WrapErrImportFailed("error"), true, true)
testOldFormatDecode("", nil, true, true)
testOldFormatDecode("a", nil, true, true)
testOldFormatDecode("2", nil, true, true)
testOldFormatDecode("{", nil, false, true)
testOldFormatDecode("{}", nil, false, true) // return EOF
testOldFormatDecode("{a", nil, false, true)
testOldFormatDecode("{\"a\":2}", nil, false, true)
testOldFormatDecode("{\"rows\"a}", nil, false, true)
testOldFormatDecode("{\"rows\":{}}", nil, false, true)
testOldFormatDecode("{\"rows\":[a]}", nil, false, true)
testOldFormatDecode("{\"rows\":[{\"dummy\": 3}]}", nil, false, true)
testOldFormatDecode("{\"rows\":[{\"pk\": 3}", nil, false, true)
// old format
testOldFormatDecode("{\"rows\":[]}", nil, false, false)
testOldFormatDecode("{\"rows\":[{\"pk\": 3}]}", nil, false, false)
testOldFormatDecode("{\n\"rows\":\n[\n{\"pk\": 3},\n{\"pk\": 4}\n]\n}\n", nil, false, false)
// new format
testListFormatDecode("[]", true) // return EOF
testListFormatDecode("[{\"pk\": 3}]", false)
testListFormatDecode("[\n{\"pk\": 3},\n{\"pk\": 4}\n]", false)
// JSONL/NDJSON format
testLinesDecode("", true) // return EOF
testLinesDecode("{\"pk\": 3}", false)
testLinesDecode("{\"pk\": 3}\n{\"pk\": 4}\n", false)
}
func (suite *ReaderSuite) TestReadCount() {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "pk",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
},
}
rows := make([]any, 0, 100)
for i := 0; i < 100; i++ {
row := make(map[string]any)
row["pk"] = int64(i)
rows = append(rows, row)
}
jsonBytes, err := json.Marshal(rows)
suite.NoError(err)
cm := mocks.NewChunkManager(suite.T())
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
r := importcommon.NewMockReader(string(jsonBytes))
return r, nil
})
// there are 100 rows to be parsed, each read batch is 40 rows
// buffer size is 320 bytes, each read batch is 40 rows
// Read() is called for 3 times, 40 + 40 + 20
reader, err := NewReader(context.Background(), cm, schema, "mockPath", 320)
suite.NoError(err)
data, err := reader.Read()
suite.NoError(err)
suite.Equal(40, data.GetRowNum())
data, err = reader.Read()
suite.NoError(err)
suite.Equal(40, data.GetRowNum())
data, err = reader.Read()
suite.NoError(err)
suite.Equal(20, data.GetRowNum())
}
func TestJsonReader(t *testing.T) {
suite.Run(t, new(ReaderSuite))
}
func (suite *ReaderSuite) TestAllowInsertAutoID_KeepUserPK() {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "pk",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
AutoID: true,
},
{
FieldID: 101,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "8"},
},
},
},
}
// build rows that explicitly include pk and a valid vec of dim 8
rows := make([]any, 0, suite.numRows)
for i := 0; i < suite.numRows; i++ {
row := make(map[string]any)
row["pk"] = int64(i + 1)
vec := make([]float64, 8)
for j := 0; j < 8; j++ {
vec[j] = float64(j)
}
row["vec"] = vec
rows = append(rows, row)
}
jsonBytes, err := json.Marshal(rows)
suite.NoError(err)
// allow_insert_autoid=false, providing PK should error
{
cm := mocks.NewChunkManager(suite.T())
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
r := importcommon.NewMockReader(string(jsonBytes))
return r, nil
})
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
suite.NoError(err)
_, err = reader.Read()
suite.Error(err)
suite.Contains(err.Error(), "is auto-generated, no need to provide")
}
// allow_insert_autoid=true, providing PK should be allowed
{
schema.Properties = []*commonpb.KeyValuePair{{Key: common.AllowInsertAutoIDKey, Value: "true"}}
cm := mocks.NewChunkManager(suite.T())
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
r := importcommon.NewMockReader(string(jsonBytes))
return r, nil
})
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
suite.NoError(err)
_, err = reader.Read()
suite.NoError(err)
}
}
Reference in New Issue View Git Blame Copy Permalink