Skip to content

Commit

Permalink
fix: Fix 0 read count during import (#38694)
Browse files Browse the repository at this point in the history
issue: #38693

Signed-off-by: bigsheeper <[email protected]>
  • Loading branch information
bigsheeper authored Dec 24, 2024
1 parent 5304707 commit 2b53b09
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 1 deletion.
9 changes: 8 additions & 1 deletion internal/util/importutilv2/common/util.go
Original file line number Diff line number Diff line change
Expand Up @@ -84,8 +84,15 @@ func EstimateReadCountPerBatch(bufferSize int, schema *schemapb.CollectionSchema
if err != nil {
return 0, err
}
if sizePerRecord <= 0 || bufferSize <= 0 {
return 0, fmt.Errorf("invalid size, sizePerRecord=%d, bufferSize=%d", sizePerRecord, bufferSize)
}
if 1000*sizePerRecord <= bufferSize {
return 1000, nil
}
return int64(bufferSize) / int64(sizePerRecord), nil
ret := int64(bufferSize) / int64(sizePerRecord)
if ret <= 0 {
return 1, nil
}
return ret, nil
}
40 changes: 40 additions & 0 deletions internal/util/importutilv2/common/util_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,3 +66,43 @@ func TestUtil_EstimateReadCountPerBatch(t *testing.T) {
_, err = EstimateReadCountPerBatch(16*1024*1024, schema)
assert.Error(t, err)
}

func TestUtil_EstimateReadCountPerBatch_InvalidBufferSize(t *testing.T) {
schema := &schemapb.CollectionSchema{}
count, err := EstimateReadCountPerBatch(16*1024*1024, schema)
assert.Error(t, err)
assert.Equal(t, int64(0), count)
t.Logf("err=%v", err)

schema = &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
DataType: schemapb.DataType_Int64,
},
},
}
count, err = EstimateReadCountPerBatch(0, schema)
assert.Error(t, err)
assert.Equal(t, int64(0), count)
t.Logf("err=%v", err)
}

func TestUtil_EstimateReadCountPerBatch_LargeSchema(t *testing.T) {
schema := &schemapb.CollectionSchema{}
for i := 0; i < 100; i++ {
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
FieldID: int64(i),
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: "10000000",
},
},
})
}
count, err := EstimateReadCountPerBatch(16*1024*1024, schema)
assert.NoError(t, err)
assert.Equal(t, int64(1), count)
}

0 comments on commit 2b53b09

Please sign in to comment.