mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-01-07 19:31:51 +08:00
3c2cf2c066
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>
855 lines
27 KiB
Go
855 lines
27 KiB
Go
package parquet
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math/rand"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/apache/arrow/go/v17/arrow"
|
|
"github.com/apache/arrow/go/v17/arrow/array"
|
|
"github.com/apache/arrow/go/v17/arrow/memory"
|
|
"github.com/apache/arrow/go/v17/parquet"
|
|
"github.com/apache/arrow/go/v17/parquet/file"
|
|
"github.com/apache/arrow/go/v17/parquet/pqarrow"
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"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/storage"
|
|
"github.com/milvus-io/milvus/internal/util/testutil"
|
|
"github.com/milvus-io/milvus/pkg/v2/common"
|
|
"github.com/milvus-io/milvus/pkg/v2/objectstorage"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
|
|
)
|
|
|
|
func init() {
|
|
paramtable.Init()
|
|
}
|
|
|
|
func TestInvalidUTF8(t *testing.T) {
|
|
const (
|
|
fieldID = int64(100)
|
|
numRows = 100
|
|
)
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: fieldID,
|
|
Name: "str",
|
|
DataType: schemapb.DataType_VarChar,
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: "max_length", Value: "256"}},
|
|
},
|
|
},
|
|
}
|
|
|
|
data := make([]string, numRows)
|
|
for i := 0; i < numRows-1; i++ {
|
|
data[i] = randomString(16)
|
|
}
|
|
data[numRows-1] = "\xc3\x28" // invalid utf-8
|
|
|
|
filePath := fmt.Sprintf("/tmp/test_%d_reader.parquet", rand.Int())
|
|
defer os.Remove(filePath)
|
|
wf, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o666)
|
|
assert.NoError(t, err)
|
|
|
|
pqSchema, err := ConvertToArrowSchemaForUT(schema, false)
|
|
assert.NoError(t, err)
|
|
fw, err := pqarrow.NewFileWriter(pqSchema, wf,
|
|
parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(numRows)), pqarrow.DefaultWriterProps())
|
|
assert.NoError(t, err)
|
|
|
|
insertData, err := storage.NewInsertData(schema)
|
|
assert.NoError(t, err)
|
|
err = insertData.Data[fieldID].AppendDataRows(data)
|
|
assert.NoError(t, err)
|
|
|
|
columns, err := testutil.BuildArrayData(schema, insertData, false)
|
|
assert.NoError(t, err)
|
|
|
|
recordBatch := array.NewRecord(pqSchema, columns, numRows)
|
|
err = fw.Write(recordBatch)
|
|
assert.NoError(t, err)
|
|
fw.Close()
|
|
|
|
ctx := context.Background()
|
|
f := storage.NewChunkManagerFactory("local", objectstorage.RootPath(testOutputPath))
|
|
cm, err := f.NewPersistentStorageChunkManager(ctx)
|
|
assert.NoError(t, err)
|
|
reader, err := NewReader(ctx, cm, schema, filePath, 64*1024*1024)
|
|
assert.NoError(t, err)
|
|
|
|
_, err = reader.Read()
|
|
assert.Error(t, err)
|
|
assert.True(t, strings.Contains(err.Error(), "contains invalid UTF-8 data"))
|
|
}
|
|
|
|
// TestParseSparseFloatRowVector tests the parseSparseFloatRowVector function
|
|
func TestParseSparseFloatRowVector(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
input string
|
|
wantMaxIdx uint32
|
|
wantErrMsg string
|
|
}{
|
|
{
|
|
name: "empty sparse vector",
|
|
input: "{}",
|
|
wantMaxIdx: 0,
|
|
},
|
|
{
|
|
name: "key-value format",
|
|
input: "{\"275574541\":1.5383775}",
|
|
wantMaxIdx: 275574542, // max index 275574541 + 1
|
|
},
|
|
{
|
|
name: "multiple key-value pairs",
|
|
input: "{\"1\":0.5,\"10\":1.5,\"100\":2.5}",
|
|
wantMaxIdx: 101, // max index 100 + 1
|
|
},
|
|
{
|
|
name: "invalid format - missing braces",
|
|
input: "\"275574541\":1.5383775",
|
|
wantErrMsg: "Invalid JSON string for SparseFloatVector",
|
|
},
|
|
{
|
|
name: "invalid JSON format",
|
|
input: "{275574541:1.5383775}",
|
|
wantErrMsg: "Invalid JSON string for SparseFloatVector",
|
|
},
|
|
{
|
|
name: "malformed JSON",
|
|
input: "{\"key\": value}",
|
|
wantErrMsg: "Invalid JSON string for SparseFloatVector",
|
|
},
|
|
{
|
|
name: "non-numeric index",
|
|
input: "{\"abc\":1.5}",
|
|
wantErrMsg: "Invalid JSON string for SparseFloatVector",
|
|
},
|
|
{
|
|
name: "non-numeric value",
|
|
input: "{\"123\":\"abc\"}",
|
|
wantErrMsg: "Invalid JSON string for SparseFloatVector",
|
|
},
|
|
{
|
|
name: "negative index",
|
|
input: "{\"-1\":1.5}",
|
|
wantErrMsg: "Invalid JSON string for SparseFloatVector",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
rowVec, maxIdx, err := parseSparseFloatRowVector(tt.input)
|
|
|
|
if tt.wantErrMsg != "" {
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), tt.wantErrMsg)
|
|
return
|
|
}
|
|
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, tt.wantMaxIdx, maxIdx)
|
|
|
|
// Verify the rowVec is properly formatted
|
|
if maxIdx > 0 {
|
|
elemCount := len(rowVec) / 8
|
|
assert.Greater(t, elemCount, 0)
|
|
|
|
// Check the last index matches our expectation
|
|
lastIdx := typeutil.SparseFloatRowIndexAt(rowVec, elemCount-1)
|
|
assert.Equal(t, tt.wantMaxIdx-1, lastIdx)
|
|
} else {
|
|
assert.Empty(t, rowVec)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestParseSparseFloatVectorStructs(t *testing.T) {
|
|
mem := memory.NewGoAllocator()
|
|
|
|
checkFunc := func(indices arrow.Array, values arrow.Array, expectSucceed bool) ([][]byte, uint32) {
|
|
st := make(map[string]arrow.Array)
|
|
if indices != nil {
|
|
st[sparseVectorIndice] = indices
|
|
}
|
|
if values != nil {
|
|
st[sparseVectorValues] = values
|
|
}
|
|
|
|
structs := make([]map[string]arrow.Array, 0)
|
|
structs = append(structs, st)
|
|
|
|
byteArr, maxDim, err := parseSparseFloatVectorStructs(structs)
|
|
if expectSucceed {
|
|
assert.NoError(t, err)
|
|
} else {
|
|
assert.Error(t, err)
|
|
}
|
|
return byteArr, maxDim
|
|
}
|
|
|
|
genInt32Arr := func(len int) *array.Int32 {
|
|
builder := array.NewInt32Builder(mem)
|
|
data := make([]int32, 0)
|
|
validData := make([]bool, 0)
|
|
for i := 0; i < len; i++ {
|
|
data = append(data, (int32)(i))
|
|
validData = append(validData, i%2 == 0)
|
|
}
|
|
|
|
builder.AppendValues(data, validData)
|
|
return builder.NewInt32Array()
|
|
}
|
|
|
|
genFloat32Arr := func(len int) *array.Float32 {
|
|
builder := array.NewFloat32Builder(mem)
|
|
data := make([]float32, 0)
|
|
validData := make([]bool, 0)
|
|
for i := 0; i < len; i++ {
|
|
data = append(data, (float32)(i))
|
|
validData = append(validData, i%2 == 0)
|
|
}
|
|
builder.AppendValues(data, validData)
|
|
return builder.NewFloat32Array()
|
|
}
|
|
|
|
genInt32ArrList := func(arr []uint32) *array.List {
|
|
builder := array.NewListBuilder(mem, &arrow.Int32Type{})
|
|
builder.Append(true)
|
|
for _, v := range arr {
|
|
builder.ValueBuilder().(*array.Int32Builder).Append((int32)(v))
|
|
}
|
|
return builder.NewListArray()
|
|
}
|
|
|
|
genUint32ArrList := func(arr []uint32) *array.List {
|
|
builder := array.NewListBuilder(mem, &arrow.Uint32Type{})
|
|
if arr != nil {
|
|
builder.Append(true)
|
|
for _, v := range arr {
|
|
builder.ValueBuilder().(*array.Uint32Builder).Append(v)
|
|
}
|
|
}
|
|
return builder.NewListArray()
|
|
}
|
|
|
|
genInt64ArrList := func(arr []uint32) *array.List {
|
|
builder := array.NewListBuilder(mem, &arrow.Int64Type{})
|
|
if arr != nil {
|
|
builder.Append(true)
|
|
for _, v := range arr {
|
|
builder.ValueBuilder().(*array.Int64Builder).Append((int64)(v))
|
|
}
|
|
}
|
|
return builder.NewListArray()
|
|
}
|
|
|
|
genUint64ArrList := func(arr []uint32) *array.List {
|
|
builder := array.NewListBuilder(mem, &arrow.Uint64Type{})
|
|
if arr != nil {
|
|
builder.Append(true)
|
|
for _, v := range arr {
|
|
builder.ValueBuilder().(*array.Uint64Builder).Append((uint64)(v))
|
|
}
|
|
}
|
|
return builder.NewListArray()
|
|
}
|
|
|
|
genFloat32ArrList := func(arr []float32) *array.List {
|
|
builder := array.NewListBuilder(mem, &arrow.Float32Type{})
|
|
if arr != nil {
|
|
builder.Append(true)
|
|
for _, v := range arr {
|
|
builder.ValueBuilder().(*array.Float32Builder).Append(v)
|
|
}
|
|
}
|
|
return builder.NewListArray()
|
|
}
|
|
|
|
genFloat64ArrList := func(arr []float32) *array.List {
|
|
builder := array.NewListBuilder(mem, &arrow.Float64Type{})
|
|
if arr != nil {
|
|
builder.Append(true)
|
|
for _, v := range arr {
|
|
builder.ValueBuilder().(*array.Float64Builder).Append((float64)(v))
|
|
}
|
|
}
|
|
return builder.NewListArray()
|
|
}
|
|
|
|
// idices field missed
|
|
checkFunc(nil, genFloat32ArrList([]float32{0.1}), false)
|
|
|
|
// values field missed
|
|
checkFunc(genUint32ArrList([]uint32{1, 2}), nil, false)
|
|
|
|
// indices is not array.List
|
|
checkFunc(genInt32Arr(2), genFloat32ArrList([]float32{0.1, 0.2}), false)
|
|
|
|
// values is not array.List
|
|
checkFunc(genUint32ArrList([]uint32{1, 2}), genFloat32Arr(2), false)
|
|
|
|
// indices is not list of int32/uint32/int64/uint64 array
|
|
checkFunc(genFloat32ArrList([]float32{0.1, 0.2, 0.3}), genFloat32ArrList([]float32{0.1, 0.2, 0.3}), false)
|
|
|
|
// values is not list of float32/float64 array
|
|
checkFunc(genUint32ArrList([]uint32{1, 2, 3}), genUint32ArrList([]uint32{1, 2, 3}), false)
|
|
|
|
// row number of indices and values are different
|
|
checkFunc(genUint32ArrList([]uint32{1, 2}), genFloat32ArrList(nil), false)
|
|
|
|
// element number of indices and values are different
|
|
checkFunc(genUint32ArrList([]uint32{1, 2}), genFloat32ArrList([]float32{0.1}), false)
|
|
|
|
// duplicated indices
|
|
checkFunc(genUint32ArrList([]uint32{4, 5, 4}), genFloat32ArrList([]float32{0.11, 0.22, 0.23}), false)
|
|
|
|
// check result is correct
|
|
// can handle empty indices/values
|
|
byteArr, maxDim := checkFunc(genUint32ArrList([]uint32{}), genFloat32ArrList([]float32{}), true)
|
|
assert.Equal(t, uint32(0), maxDim)
|
|
assert.Equal(t, 1, len(byteArr))
|
|
assert.Equal(t, 0, len(byteArr[0]))
|
|
|
|
// note that the input indices is not sorted, the parseSparseFloatVectorStructs
|
|
// returns correct maxDim and byteArr
|
|
indices := []uint32{25, 78, 56}
|
|
values := []float32{0.11, 0.22, 0.23}
|
|
sortedIndices, sortedValues := typeutil.SortSparseFloatRow(indices, values)
|
|
rowBytes := typeutil.CreateSparseFloatRow(sortedIndices, sortedValues)
|
|
|
|
isValidFunc := func(indices arrow.Array, values arrow.Array) {
|
|
byteArr, maxDim := checkFunc(indices, values, true)
|
|
assert.Equal(t, uint32(78), maxDim)
|
|
assert.Equal(t, 1, len(byteArr))
|
|
assert.Equal(t, rowBytes, byteArr[0])
|
|
}
|
|
|
|
// ensure all supported types are correct
|
|
isValidFunc(genUint32ArrList(indices), genFloat32ArrList(values))
|
|
isValidFunc(genUint32ArrList(indices), genFloat64ArrList(values))
|
|
isValidFunc(genInt32ArrList(indices), genFloat32ArrList(values))
|
|
isValidFunc(genInt32ArrList(indices), genFloat64ArrList(values))
|
|
isValidFunc(genUint64ArrList(indices), genFloat32ArrList(values))
|
|
isValidFunc(genUint64ArrList(indices), genFloat64ArrList(values))
|
|
isValidFunc(genInt64ArrList(indices), genFloat32ArrList(values))
|
|
isValidFunc(genInt64ArrList(indices), genFloat64ArrList(values))
|
|
}
|
|
|
|
func TestReadFieldData(t *testing.T) {
|
|
checkFunc := func(t *testing.T, nullPercent int, readScehamIsNullable bool, dataType schemapb.DataType, elementType schemapb.DataType) {
|
|
fieldName := dataType.String()
|
|
if elementType != schemapb.DataType_None {
|
|
fieldName = fieldName + "_" + elementType.String()
|
|
}
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: fieldName,
|
|
DataType: dataType,
|
|
ElementType: elementType,
|
|
Nullable: nullPercent != 0,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "dim",
|
|
Value: "16",
|
|
},
|
|
{
|
|
Key: "max_length",
|
|
Value: "1000",
|
|
},
|
|
{
|
|
Key: "max_capacity",
|
|
Value: "50",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
arrDataType, err := convertToArrowDataType(schema.Fields[0], false)
|
|
assert.NoError(t, err)
|
|
arrFields := make([]arrow.Field, 0)
|
|
arrFields = append(arrFields, arrow.Field{
|
|
Name: schema.Fields[0].Name,
|
|
Type: arrDataType,
|
|
Nullable: true,
|
|
Metadata: arrow.Metadata{},
|
|
})
|
|
pqSchema := arrow.NewSchema(arrFields, nil)
|
|
|
|
filePath := fmt.Sprintf("/tmp/test_%d_reader.parquet", rand.Int())
|
|
defer os.Remove(filePath)
|
|
wf, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o666)
|
|
assert.NoError(t, err)
|
|
|
|
fw, err := pqarrow.NewFileWriter(pqSchema, wf,
|
|
parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(100)), pqarrow.DefaultWriterProps())
|
|
assert.NoError(t, err)
|
|
|
|
rowCount := 5
|
|
insertData, err := testutil.CreateInsertData(schema, rowCount, nullPercent)
|
|
assert.NoError(t, err)
|
|
columns, err := testutil.BuildArrayData(schema, insertData, false)
|
|
assert.NoError(t, err)
|
|
|
|
recordBatch := array.NewRecord(pqSchema, columns, int64(rowCount))
|
|
err = fw.Write(recordBatch)
|
|
assert.NoError(t, err)
|
|
fw.Close()
|
|
|
|
ctx := context.Background()
|
|
f := storage.NewChunkManagerFactory("local", objectstorage.RootPath(testOutputPath))
|
|
cm, err := f.NewPersistentStorageChunkManager(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
schema.Fields[0].Nullable = readScehamIsNullable
|
|
reader, err := NewReader(ctx, cm, schema, filePath, 64*1024*1024)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, reader)
|
|
defer reader.Close()
|
|
|
|
_, err = reader.Read()
|
|
if !readScehamIsNullable && nullPercent != 0 {
|
|
assert.Error(t, err)
|
|
} else {
|
|
assert.NoError(t, err)
|
|
}
|
|
}
|
|
|
|
type testCase struct {
|
|
name string
|
|
nullPercent int
|
|
readScehamIsNullable bool
|
|
dataType schemapb.DataType
|
|
elementType schemapb.DataType
|
|
}
|
|
buildCaseFunc := func(nullPercent int, readScehamIsNullable bool, dataType schemapb.DataType, elementType schemapb.DataType) *testCase {
|
|
name := fmt.Sprintf("nullPercent='%v' schemaNullable='%v' dataType='%s' elementType='%s'",
|
|
nullPercent, readScehamIsNullable, dataType, elementType)
|
|
return &testCase{
|
|
name: name,
|
|
nullPercent: nullPercent,
|
|
readScehamIsNullable: readScehamIsNullable,
|
|
dataType: dataType,
|
|
elementType: elementType,
|
|
}
|
|
}
|
|
cases := make([]*testCase, 0)
|
|
|
|
nullableDataTypes := []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,
|
|
schemapb.DataType_FloatVector,
|
|
schemapb.DataType_BinaryVector,
|
|
schemapb.DataType_SparseFloatVector,
|
|
schemapb.DataType_Float16Vector,
|
|
schemapb.DataType_BFloat16Vector,
|
|
schemapb.DataType_Int8Vector,
|
|
}
|
|
for _, dataType := range nullableDataTypes {
|
|
for _, nullPercent := range []int{0, 50} {
|
|
for _, readScehamIsNullable := range []bool{true, false} {
|
|
cases = append(cases, buildCaseFunc(nullPercent, readScehamIsNullable, dataType, schemapb.DataType_None))
|
|
}
|
|
}
|
|
}
|
|
|
|
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,
|
|
}
|
|
for _, elementType := range elementTypes {
|
|
for _, nullPercent := range []int{0, 50} {
|
|
for _, readScehamIsNullable := range []bool{true, false} {
|
|
cases = append(cases, buildCaseFunc(nullPercent, readScehamIsNullable, schemapb.DataType_Array, elementType))
|
|
}
|
|
}
|
|
}
|
|
|
|
notNullableTypes := []schemapb.DataType{
|
|
schemapb.DataType_JSON,
|
|
}
|
|
for _, dataType := range notNullableTypes {
|
|
cases = append(cases, buildCaseFunc(0, false, dataType, schemapb.DataType_None))
|
|
}
|
|
|
|
for _, tt := range cases {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
checkFunc(t, tt.nullPercent, tt.readScehamIsNullable, tt.dataType, tt.elementType)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestTypeMismatch(t *testing.T) {
|
|
checkFunc := func(srcDataType schemapb.DataType, srcElementType schemapb.DataType, dstDataType schemapb.DataType, dstElementType schemapb.DataType, nullalbe bool) {
|
|
fieldName := "test_field"
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: fieldName,
|
|
DataType: srcDataType,
|
|
ElementType: srcElementType,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "dim",
|
|
Value: "16",
|
|
},
|
|
{
|
|
Key: "max_length",
|
|
Value: "1000",
|
|
},
|
|
{
|
|
Key: "max_capacity",
|
|
Value: "50",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
arrDataType, err := convertToArrowDataType(schema.Fields[0], false)
|
|
assert.NoError(t, err)
|
|
arrFields := make([]arrow.Field, 0)
|
|
arrFields = append(arrFields, arrow.Field{
|
|
Name: schema.Fields[0].Name,
|
|
Type: arrDataType,
|
|
Nullable: true,
|
|
Metadata: arrow.Metadata{},
|
|
})
|
|
pqSchema := arrow.NewSchema(arrFields, nil)
|
|
|
|
filePath := fmt.Sprintf("/tmp/test_%d_reader.parquet", rand.Int())
|
|
defer os.Remove(filePath)
|
|
wf, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o666)
|
|
assert.NoError(t, err)
|
|
|
|
fw, err := pqarrow.NewFileWriter(pqSchema, wf,
|
|
parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(100)), pqarrow.DefaultWriterProps())
|
|
assert.NoError(t, err)
|
|
|
|
rowCount := 5
|
|
insertData, err := testutil.CreateInsertData(schema, rowCount, 0)
|
|
assert.NoError(t, err)
|
|
columns, err := testutil.BuildArrayData(schema, insertData, false)
|
|
assert.NoError(t, err)
|
|
|
|
recordBatch := array.NewRecord(pqSchema, columns, int64(rowCount))
|
|
err = fw.Write(recordBatch)
|
|
assert.NoError(t, err)
|
|
fw.Close()
|
|
|
|
ctx := context.Background()
|
|
f := storage.NewChunkManagerFactory("local", objectstorage.RootPath(testOutputPath))
|
|
cm, err := f.NewPersistentStorageChunkManager(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
schema.Fields[0].DataType = dstDataType
|
|
schema.Fields[0].ElementType = dstElementType
|
|
schema.Fields[0].Nullable = nullalbe
|
|
cmReader, err := cm.Reader(ctx, filePath)
|
|
assert.NoError(t, err)
|
|
reader, err := file.NewParquetReader(cmReader, file.WithReadProps(&parquet.ReaderProperties{
|
|
BufferSize: 65535,
|
|
BufferedStreamEnabled: true,
|
|
}))
|
|
assert.NoError(t, err)
|
|
|
|
readProps := pqarrow.ArrowReadProperties{
|
|
BatchSize: int64(rowCount),
|
|
}
|
|
fileReader, err := pqarrow.NewFileReader(reader, readProps, memory.DefaultAllocator)
|
|
assert.NoError(t, err)
|
|
columnReader, err := NewFieldReader(ctx, fileReader, 0, schema.Fields[0], common.DefaultTimezone)
|
|
assert.NoError(t, err)
|
|
|
|
_, _, err = columnReader.Next(int64(rowCount))
|
|
if srcDataType != dstDataType || srcElementType != dstElementType {
|
|
assert.Error(t, err)
|
|
} else {
|
|
assert.NoError(t, err)
|
|
}
|
|
}
|
|
|
|
type testCase struct {
|
|
name string
|
|
srcDataType schemapb.DataType
|
|
srcElementType schemapb.DataType
|
|
dstDataType schemapb.DataType
|
|
dstElementType schemapb.DataType
|
|
nullable bool
|
|
}
|
|
buildCaseFunc := func(srcDataType schemapb.DataType, srcElementType schemapb.DataType, dstDataType schemapb.DataType, dstElementType schemapb.DataType, nullable bool) *testCase {
|
|
name := fmt.Sprintf("srcDataType='%s' srcElementType='%s' dstDataType='%s' dstElementType='%s' nullable='%v'",
|
|
srcDataType, srcElementType, dstDataType, dstElementType, nullable)
|
|
return &testCase{
|
|
name: name,
|
|
srcDataType: srcDataType,
|
|
srcElementType: srcElementType,
|
|
dstDataType: dstDataType,
|
|
dstElementType: dstElementType,
|
|
nullable: nullable,
|
|
}
|
|
}
|
|
cases := make([]*testCase, 0)
|
|
|
|
scalarDataTypes := []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 scalarDataTypes {
|
|
srcDataType := schemapb.DataType_Bool
|
|
if dataType == schemapb.DataType_Bool {
|
|
srcDataType = schemapb.DataType_Int8
|
|
}
|
|
cases = append(cases, buildCaseFunc(srcDataType, schemapb.DataType_None, dataType, schemapb.DataType_None, true))
|
|
cases = append(cases, buildCaseFunc(srcDataType, schemapb.DataType_None, dataType, schemapb.DataType_None, false))
|
|
}
|
|
|
|
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,
|
|
}
|
|
for _, elementType := range elementTypes {
|
|
srcElementType := schemapb.DataType_Bool
|
|
if elementType == schemapb.DataType_Bool {
|
|
srcElementType = schemapb.DataType_Int8
|
|
}
|
|
// element type mismatch
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Array, srcElementType, schemapb.DataType_Array, elementType, true))
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Array, srcElementType, schemapb.DataType_Array, elementType, false))
|
|
// not a list
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Bool, schemapb.DataType_None, schemapb.DataType_Array, elementType, true))
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Bool, schemapb.DataType_None, schemapb.DataType_Array, elementType, false))
|
|
}
|
|
|
|
vectorTypes := []schemapb.DataType{
|
|
schemapb.DataType_FloatVector,
|
|
schemapb.DataType_BinaryVector,
|
|
schemapb.DataType_SparseFloatVector,
|
|
schemapb.DataType_Float16Vector,
|
|
schemapb.DataType_BFloat16Vector,
|
|
schemapb.DataType_Int8Vector,
|
|
}
|
|
for _, dataType := range vectorTypes {
|
|
srcDataType := schemapb.DataType_Bool
|
|
cases = append(cases, buildCaseFunc(srcDataType, schemapb.DataType_None, dataType, schemapb.DataType_None, true))
|
|
cases = append(cases, buildCaseFunc(srcDataType, schemapb.DataType_None, dataType, schemapb.DataType_None, false))
|
|
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Array, schemapb.DataType_Bool, dataType, schemapb.DataType_None, true))
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Array, schemapb.DataType_Bool, dataType, schemapb.DataType_None, false))
|
|
}
|
|
|
|
notNullableTypes := []schemapb.DataType{
|
|
schemapb.DataType_JSON,
|
|
}
|
|
for _, dataType := range notNullableTypes {
|
|
srcDataType := schemapb.DataType_Bool
|
|
if dataType == schemapb.DataType_Bool {
|
|
srcDataType = schemapb.DataType_Int8
|
|
}
|
|
// not a list
|
|
cases = append(cases, buildCaseFunc(srcDataType, schemapb.DataType_None, dataType, schemapb.DataType_None, false))
|
|
// element type mismatch
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Array, schemapb.DataType_Bool, dataType, schemapb.DataType_None, false))
|
|
}
|
|
|
|
for _, tt := range cases {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
checkFunc(tt.srcDataType, tt.srcElementType, tt.dstDataType, tt.dstElementType, tt.nullable)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestArrayNullElement(t *testing.T) {
|
|
checkFunc := func(dataType schemapb.DataType, elementType schemapb.DataType) {
|
|
fieldName := "test_field"
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: fieldName,
|
|
DataType: dataType,
|
|
ElementType: elementType,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "dim",
|
|
Value: "16",
|
|
},
|
|
{
|
|
Key: "max_length",
|
|
Value: "1000",
|
|
},
|
|
{
|
|
Key: "max_capacity",
|
|
Value: "50",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
arrDataType, err := convertToArrowDataType(schema.Fields[0], false)
|
|
assert.NoError(t, err)
|
|
arrFields := make([]arrow.Field, 0)
|
|
arrFields = append(arrFields, arrow.Field{
|
|
Name: schema.Fields[0].Name,
|
|
Type: arrDataType,
|
|
Nullable: true,
|
|
Metadata: arrow.Metadata{},
|
|
})
|
|
pqSchema := arrow.NewSchema(arrFields, nil)
|
|
|
|
filePath := fmt.Sprintf("/tmp/test_%d_reader.parquet", rand.Int())
|
|
defer os.Remove(filePath)
|
|
wf, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o666)
|
|
assert.NoError(t, err)
|
|
|
|
fw, err := pqarrow.NewFileWriter(pqSchema, wf,
|
|
parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(100)), pqarrow.DefaultWriterProps())
|
|
assert.NoError(t, err)
|
|
|
|
mem := memory.NewGoAllocator()
|
|
columns := make([]arrow.Array, 0, len(schema.Fields))
|
|
switch elementType {
|
|
case schemapb.DataType_Bool:
|
|
builder := array.NewListBuilder(mem, &arrow.BooleanType{})
|
|
valueBuilder := builder.ValueBuilder().(*array.BooleanBuilder)
|
|
valueBuilder.AppendValues([]bool{true, false}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
case schemapb.DataType_Int8:
|
|
builder := array.NewListBuilder(mem, &arrow.Int8Type{})
|
|
valueBuilder := builder.ValueBuilder().(*array.Int8Builder)
|
|
valueBuilder.AppendValues([]int8{1, 2}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
case schemapb.DataType_Int16:
|
|
builder := array.NewListBuilder(mem, &arrow.Int16Type{})
|
|
valueBuilder := builder.ValueBuilder().(*array.Int16Builder)
|
|
valueBuilder.AppendValues([]int16{1, 2}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
case schemapb.DataType_Int32:
|
|
builder := array.NewListBuilder(mem, &arrow.Int32Type{})
|
|
valueBuilder := builder.ValueBuilder().(*array.Int32Builder)
|
|
valueBuilder.AppendValues([]int32{1, 2}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
case schemapb.DataType_Int64:
|
|
builder := array.NewListBuilder(mem, &arrow.Int64Type{})
|
|
valueBuilder := builder.ValueBuilder().(*array.Int64Builder)
|
|
valueBuilder.AppendValues([]int64{1, 2}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
case schemapb.DataType_Float:
|
|
builder := array.NewListBuilder(mem, &arrow.Float32Type{})
|
|
valueBuilder := builder.ValueBuilder().(*array.Float32Builder)
|
|
valueBuilder.AppendValues([]float32{0.1, 0.2}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
case schemapb.DataType_Double:
|
|
builder := array.NewListBuilder(mem, &arrow.Float64Type{})
|
|
valueBuilder := builder.ValueBuilder().(*array.Float64Builder)
|
|
valueBuilder.AppendValues([]float64{0.1, 0.2}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
case schemapb.DataType_String, schemapb.DataType_VarChar:
|
|
builder := array.NewListBuilder(mem, &arrow.StringType{})
|
|
valueBuilder := builder.ValueBuilder().(*array.StringBuilder)
|
|
valueBuilder.AppendValues([]string{"a", "b"}, []bool{true, false})
|
|
builder.AppendValues([]int32{0}, []bool{true})
|
|
columns = append(columns, builder.NewListArray())
|
|
default:
|
|
break
|
|
}
|
|
|
|
recordBatch := array.NewRecord(pqSchema, columns, int64(1))
|
|
err = fw.Write(recordBatch)
|
|
assert.NoError(t, err)
|
|
fw.Close()
|
|
|
|
ctx := context.Background()
|
|
f := storage.NewChunkManagerFactory("local", objectstorage.RootPath(testOutputPath))
|
|
cm, err := f.NewPersistentStorageChunkManager(ctx)
|
|
assert.NoError(t, err)
|
|
|
|
reader, err := NewReader(ctx, cm, schema, filePath, 64*1024*1024)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, reader)
|
|
defer reader.Close()
|
|
|
|
_, err = reader.Read()
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
type testCase struct {
|
|
name string
|
|
dataType schemapb.DataType
|
|
elementType schemapb.DataType
|
|
}
|
|
buildCaseFunc := func(dataType schemapb.DataType, elementType schemapb.DataType) *testCase {
|
|
name := fmt.Sprintf("dataType='%s' elementType='%s'", dataType, elementType)
|
|
return &testCase{
|
|
name: name,
|
|
dataType: dataType,
|
|
elementType: elementType,
|
|
}
|
|
}
|
|
cases := make([]*testCase, 0)
|
|
|
|
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,
|
|
}
|
|
for _, elementType := range elementTypes {
|
|
cases = append(cases, buildCaseFunc(schemapb.DataType_Array, elementType))
|
|
}
|
|
|
|
for _, tt := range cases {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
checkFunc(tt.dataType, tt.elementType)
|
|
})
|
|
}
|
|
}
|