wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
717 lines
18 KiB
Go
717 lines
18 KiB
Go
// 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 compactor
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
"go.uber.org/atomic"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/internal/allocator"
|
|
"github.com/milvus-io/milvus/internal/compaction"
|
|
"github.com/milvus-io/milvus/internal/mocks/flushcommon/mock_util"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
|
|
)
|
|
|
|
func TestMultiSegmentWriterSuite(t *testing.T) {
|
|
suite.Run(t, new(MultiSegmentWriterSuite))
|
|
}
|
|
|
|
type MultiSegmentWriterSuite struct {
|
|
suite.Suite
|
|
|
|
mockBinlogIO *mock_util.MockBinlogIO
|
|
mockAlloc *allocator.MockAllocator
|
|
mockID atomic.Int64
|
|
|
|
collectionID int64
|
|
partitionID int64
|
|
channel string
|
|
batchSize int
|
|
params compaction.Params
|
|
}
|
|
|
|
type testCompactionAllocator struct {
|
|
next int64
|
|
allocOneCalls int
|
|
failAllocOneAt int
|
|
failAllocOneErr error
|
|
}
|
|
|
|
type closeErrSerializeWriter struct {
|
|
err error
|
|
}
|
|
|
|
type closeErrBinlogRecordWriter struct {
|
|
writtenUncompressed uint64
|
|
}
|
|
|
|
func (w closeErrSerializeWriter) WriteValue(*storage.Value) error {
|
|
return nil
|
|
}
|
|
|
|
func (w closeErrSerializeWriter) Flush() error {
|
|
return nil
|
|
}
|
|
|
|
func (w closeErrSerializeWriter) Close() error {
|
|
return w.err
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) Write(storage.Record) error {
|
|
return nil
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) GetWrittenUncompressed() uint64 {
|
|
return w.writtenUncompressed
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) Close() error {
|
|
return nil
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) GetLogs() (
|
|
map[storage.FieldID]*datapb.FieldBinlog,
|
|
*datapb.FieldBinlog,
|
|
map[storage.FieldID]*datapb.FieldBinlog,
|
|
string,
|
|
[]int64,
|
|
) {
|
|
return nil, nil, nil, "", nil
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) GetRowNum() int64 {
|
|
return 0
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) FlushChunk() error {
|
|
return nil
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) GetBufferUncompressed() uint64 {
|
|
return w.writtenUncompressed
|
|
}
|
|
|
|
func (w closeErrBinlogRecordWriter) Schema() *schemapb.CollectionSchema {
|
|
return nil
|
|
}
|
|
|
|
func (a *testCompactionAllocator) Alloc(count uint32) (int64, int64, error) {
|
|
start := a.next
|
|
a.next += int64(count)
|
|
return start, a.next, nil
|
|
}
|
|
|
|
func (a *testCompactionAllocator) AllocOne() (int64, error) {
|
|
a.allocOneCalls++
|
|
if a.failAllocOneAt > 0 && a.allocOneCalls == a.failAllocOneAt {
|
|
if a.failAllocOneErr != nil {
|
|
return 0, a.failAllocOneErr
|
|
}
|
|
return 0, allocator.NewIDExhaustedError(0, 0, 1)
|
|
}
|
|
a.next++
|
|
return a.next, nil
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) SetupSuite() {
|
|
paramtable.Get().Init(paramtable.NewBaseTable())
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) SetupTest() {
|
|
paramtable.Get().Save(paramtable.Get().CommonCfg.StorageType.Key, "local")
|
|
|
|
s.mockBinlogIO = mock_util.NewMockBinlogIO(s.T())
|
|
s.mockBinlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(nil).Maybe()
|
|
|
|
s.mockAlloc = allocator.NewMockAllocator(s.T())
|
|
s.mockID.Store(time.Now().UnixMilli())
|
|
s.mockAlloc.EXPECT().Alloc(mock.Anything).RunAndReturn(func(x uint32) (int64, int64, error) {
|
|
start := s.mockID.Load()
|
|
end := s.mockID.Add(int64(x))
|
|
return start, end, nil
|
|
}).Maybe()
|
|
s.mockAlloc.EXPECT().AllocOne().RunAndReturn(func() (int64, error) {
|
|
end := s.mockID.Add(1)
|
|
return end, nil
|
|
}).Maybe()
|
|
|
|
s.collectionID = 100
|
|
s.partitionID = 101
|
|
s.channel = "test_channel"
|
|
s.batchSize = 100
|
|
s.params = compaction.GenParams()
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TearDownTest() {
|
|
paramtable.Get().Reset(paramtable.Get().CommonCfg.StorageType.Key)
|
|
}
|
|
|
|
// genSimpleSchema generates a simple collection schema for testing
|
|
func (s *MultiSegmentWriterSuite) genSimpleSchema() *schemapb.CollectionSchema {
|
|
return &schemapb.CollectionSchema{
|
|
Name: "test_schema",
|
|
Description: "test schema for multi segment writer",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.RowIDField,
|
|
Name: "row_id",
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: common.TimeStampField,
|
|
Name: "timestamp",
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "varchar_field",
|
|
DataType: schemapb.DataType_VarChar,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MaxLengthKey,
|
|
Value: "128",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "float_vector",
|
|
Description: "float vector field",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.DimKey,
|
|
Value: "4",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
// genTestValue generates a test storage.Value for the given ID
|
|
func (s *MultiSegmentWriterSuite) genTestValue(id int64) *storage.Value {
|
|
ts := tsoutil.ComposeTSByTime(time.Now())
|
|
return &storage.Value{
|
|
PK: storage.NewInt64PrimaryKey(id),
|
|
Timestamp: int64(ts),
|
|
Value: map[int64]interface{}{
|
|
common.RowIDField: id,
|
|
common.TimeStampField: int64(ts),
|
|
100: id, // pk
|
|
101: "test_varchar_value", // varchar
|
|
102: []float32{1.0, 2.0, 3.0, 4.0}, // float vector
|
|
},
|
|
}
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestNewMultiSegmentWriter() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
1024*1024, // 1MB segment size
|
|
schema,
|
|
s.params,
|
|
1000, // maxRows
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
|
|
s.NoError(err)
|
|
s.NotNil(writer)
|
|
s.Equal(s.collectionID, writer.collectionID)
|
|
s.Equal(s.partitionID, writer.partitionID)
|
|
s.Equal(s.channel, writer.channel)
|
|
s.Equal(int64(1024*1024), writer.segmentSize)
|
|
s.Equal(int64(1000), writer.maxRows)
|
|
s.Equal(s.batchSize, writer.batchSize)
|
|
s.Empty(writer.res) // No segments created yet
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestWriteSingleSegment() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
// Use a large segment size to ensure all data fits in one segment
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
10*1024*1024, // 10MB segment size
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
// Write some test data
|
|
numRows := 100
|
|
for i := 0; i < numRows; i++ {
|
|
value := s.genTestValue(int64(i))
|
|
err := writer.WriteValue(value)
|
|
s.NoError(err)
|
|
}
|
|
|
|
// Close the writer to finalize the segment
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
|
|
// Check that only one segment was created
|
|
segments := writer.GetCompactionSegments()
|
|
s.Equal(1, len(segments))
|
|
|
|
segment := segments[0]
|
|
s.Equal(int64(numRows), segment.NumOfRows)
|
|
s.Equal(s.channel, segment.Channel)
|
|
s.NotEmpty(segment.InsertLogs)
|
|
s.NotEmpty(segment.Field2StatslogPaths)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestWriteMultipleSegments() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
// Use a small segment size to force multiple segments
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
1024, // 1KB segment size - very small to force rotation
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
// Write enough data to exceed segment size multiple times
|
|
numRows := 1000
|
|
expectedSegments := 3 // Expect at least 3 segments given the small size
|
|
|
|
for i := 0; i < numRows; i++ {
|
|
value := s.genTestValue(int64(i))
|
|
err := writer.WriteValue(value)
|
|
s.NoError(err)
|
|
|
|
// Flush periodically to trigger size checks
|
|
if i%100 == 0 {
|
|
err = writer.Flush()
|
|
s.NoError(err)
|
|
}
|
|
}
|
|
|
|
// Close the writer to finalize all segments
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
|
|
// Check that multiple segments were created
|
|
segments := writer.GetCompactionSegments()
|
|
s.GreaterOrEqual(len(segments), expectedSegments)
|
|
|
|
// Verify all segments are valid
|
|
totalRows := int64(0)
|
|
for i, segment := range segments {
|
|
s.Greater(segment.NumOfRows, int64(0), "Segment %d should have rows", i)
|
|
s.Equal(s.channel, segment.Channel)
|
|
s.NotEmpty(segment.InsertLogs)
|
|
s.NotEmpty(segment.Field2StatslogPaths)
|
|
s.NotEqual(int64(0), segment.SegmentID)
|
|
totalRows += segment.NumOfRows
|
|
}
|
|
|
|
// Total rows should match what we wrote
|
|
s.Equal(int64(numRows), totalRows)
|
|
|
|
// Verify segment IDs are unique
|
|
segmentIDs := make(map[int64]bool)
|
|
for _, segment := range segments {
|
|
s.False(segmentIDs[segment.SegmentID], "Segment ID %d should be unique", segment.SegmentID)
|
|
segmentIDs[segment.SegmentID] = true
|
|
}
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestSegmentRotation() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
2048, // 2KB segment size
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
// Track segments as they are created
|
|
initialSegments := len(writer.GetCompactionSegments())
|
|
|
|
// Write data until we trigger rotation
|
|
for i := 0; i < 1000; i++ {
|
|
value := s.genTestValue(int64(i))
|
|
err := writer.WriteValue(value)
|
|
s.NoError(err)
|
|
|
|
// Check if rotation happened
|
|
if len(writer.GetCompactionSegments()) > initialSegments {
|
|
// First rotation detected
|
|
break
|
|
}
|
|
}
|
|
|
|
// Verify at least one segment was created during rotation
|
|
s.Greater(len(writer.GetCompactionSegments()), initialSegments)
|
|
|
|
// Continue writing to trigger more rotations
|
|
for i := 500; i < 1000; i++ {
|
|
value := s.genTestValue(int64(i))
|
|
err := writer.WriteValue(value)
|
|
s.NoError(err)
|
|
}
|
|
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
|
|
// Should have multiple segments
|
|
finalSegments := writer.GetCompactionSegments()
|
|
s.GreaterOrEqual(len(finalSegments), 2)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestSegmentIDExhaustionGrowsCurrentSegment() {
|
|
schema := s.genSimpleSchema()
|
|
segmentAlloc := allocator.NewLocalAllocator(1000, 1001)
|
|
logAlloc := allocator.NewLocalAllocator(2000, 3000)
|
|
allocator := NewCompactionAllocator(segmentAlloc, logAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
1,
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
1,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
s.Require().NoError(writer.WriteValue(s.genTestValue(1)))
|
|
s.Require().NoError(writer.WriteValue(s.genTestValue(2)))
|
|
s.Require().NoError(writer.WriteValue(s.genTestValue(3)))
|
|
s.Require().NoError(writer.Close())
|
|
|
|
segments := writer.GetCompactionSegments()
|
|
s.Require().Len(segments, 1)
|
|
s.EqualValues(1000, segments[0].SegmentID)
|
|
s.EqualValues(3, segments[0].NumOfRows)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestRotateAllocFailureKeepsCurrentWriterOpen() {
|
|
schema := s.genSimpleSchema()
|
|
allocErr := errors.New("rootcoord unavailable")
|
|
segmentAlloc := &testCompactionAllocator{next: 1000, failAllocOneAt: 2, failAllocOneErr: allocErr}
|
|
logAlloc := &testCompactionAllocator{next: 2000}
|
|
allocator := NewCompactionAllocator(segmentAlloc, logAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
1,
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
1,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
err = writer.WriteValue(s.genTestValue(1))
|
|
s.Require().NoError(err)
|
|
err = writer.WriteValue(s.genTestValue(2))
|
|
s.ErrorIs(err, allocErr)
|
|
s.NotNil(writer.writer)
|
|
s.Empty(writer.GetCompactionSegments())
|
|
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
s.Len(writer.GetCompactionSegments(), 1)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestLogIDExhaustionDuringRotationReturnsError() {
|
|
segmentAlloc := &testCompactionAllocator{next: 1000}
|
|
logAlloc := &testCompactionAllocator{next: 2000}
|
|
writer := &MultiSegmentWriter{
|
|
allocator: NewCompactionAllocator(segmentAlloc, logAlloc),
|
|
writer: &storage.BinlogValueWriter{
|
|
BinlogRecordWriter: closeErrBinlogRecordWriter{writtenUncompressed: 2},
|
|
SerializeWriter: closeErrSerializeWriter{err: allocator.NewIDExhaustedError(0, 0, 1)},
|
|
},
|
|
segmentSize: 1,
|
|
currentSegmentID: 1000,
|
|
collectionID: s.collectionID,
|
|
partitionID: s.partitionID,
|
|
channel: s.channel,
|
|
}
|
|
|
|
var err error
|
|
s.NotPanics(func() {
|
|
err = writer.rotateWriterOrGrowCurrent()
|
|
})
|
|
s.True(allocator.IsIDExhausted(err))
|
|
s.Nil(writer.writer)
|
|
s.Empty(writer.GetCompactionSegments())
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestCloseFailureClearsWriter() {
|
|
closeErr := errors.New("close failed")
|
|
writer := &MultiSegmentWriter{
|
|
writer: &storage.BinlogValueWriter{
|
|
SerializeWriter: closeErrSerializeWriter{err: closeErr},
|
|
},
|
|
}
|
|
|
|
err := writer.closeWriter()
|
|
s.ErrorIs(err, closeErr)
|
|
s.Nil(writer.writer)
|
|
s.Empty(writer.GetCompactionSegments())
|
|
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestWriterMethods() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
1024*1024, // 1MB
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
// Test initial state
|
|
s.Equal(uint64(0), writer.GetWrittenUncompressed())
|
|
s.Equal(uint64(0), writer.GetBufferUncompressed())
|
|
s.Empty(writer.GetCompactionSegments())
|
|
|
|
// Write some data
|
|
for i := 0; i < 10; i++ {
|
|
value := s.genTestValue(int64(i))
|
|
err := writer.WriteValue(value)
|
|
s.NoError(err)
|
|
}
|
|
|
|
// Test flush operations
|
|
err = writer.Flush()
|
|
s.NoError(err)
|
|
|
|
// Test after writing
|
|
s.Greater(writer.GetWrittenUncompressed(), uint64(0))
|
|
|
|
err = writer.FlushChunk()
|
|
s.NoError(err)
|
|
|
|
// Close and verify
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
|
|
segments := writer.GetCompactionSegments()
|
|
s.Equal(1, len(segments))
|
|
s.Equal(int64(10), segments[0].NumOfRows)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestWriteWithRecord() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
1024, // Small size to trigger rotation
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
// Create a test record (this would normally come from storage layer)
|
|
// For simplicity, we'll use WriteValue in this test since creating
|
|
// proper Arrow records is complex
|
|
numRows := 200
|
|
for i := 0; i < numRows; i++ {
|
|
value := s.genTestValue(int64(i))
|
|
err := writer.WriteValue(value)
|
|
s.NoError(err)
|
|
}
|
|
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
|
|
segments := writer.GetCompactionSegments()
|
|
s.GreaterOrEqual(len(segments), 2) // Should have multiple segments
|
|
|
|
totalRows := int64(0)
|
|
for _, segment := range segments {
|
|
totalRows += segment.NumOfRows
|
|
}
|
|
s.Equal(int64(numRows), totalRows)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestEmptyWriter() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
1024*1024,
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
// Close without writing any data
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
|
|
// Should have no segments
|
|
segments := writer.GetCompactionSegments()
|
|
s.Empty(segments)
|
|
}
|
|
|
|
func (s *MultiSegmentWriterSuite) TestLargeDataWrite() {
|
|
schema := s.genSimpleSchema()
|
|
allocator := NewCompactionAllocator(s.mockAlloc, s.mockAlloc)
|
|
|
|
writer, err := NewMultiSegmentWriter(
|
|
context.Background(),
|
|
s.mockBinlogIO,
|
|
allocator,
|
|
64*1024, // 64KB segments
|
|
schema,
|
|
s.params,
|
|
1000,
|
|
s.partitionID,
|
|
s.collectionID,
|
|
s.channel,
|
|
s.batchSize,
|
|
storage.WithStorageConfig(s.params.StorageConfig),
|
|
)
|
|
s.Require().NoError(err)
|
|
|
|
// Write a large amount of data
|
|
numRows := 5000
|
|
for i := 0; i < numRows; i++ {
|
|
value := s.genTestValue(int64(i))
|
|
err := writer.WriteValue(value)
|
|
s.NoError(err)
|
|
|
|
// Flush every 1000 rows to trigger size checks
|
|
if i%1000 == 0 {
|
|
err = writer.Flush()
|
|
s.NoError(err)
|
|
}
|
|
}
|
|
|
|
err = writer.Close()
|
|
s.NoError(err)
|
|
|
|
segments := writer.GetCompactionSegments()
|
|
s.GreaterOrEqual(len(segments), 1) // Should have many segments
|
|
|
|
// Verify data integrity
|
|
totalRows := int64(0)
|
|
for _, segment := range segments {
|
|
s.Greater(segment.NumOfRows, int64(0))
|
|
s.NotEmpty(segment.InsertLogs)
|
|
totalRows += segment.NumOfRows
|
|
}
|
|
s.Equal(int64(numRows), totalRows)
|
|
}
|