mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-01 03:18:29 +08:00
633a749880
Signed-off-by: xiaofan-luan <xiaofan.luan@zilliz.com> Signed-off-by: xiaofan-luan <xiaofan.luan@zilliz.com>
79 lines
1.7 KiB
Go
79 lines
1.7 KiB
Go
package storage
|
|
|
|
import (
|
|
"fmt"
|
|
"testing"
|
|
|
|
"github.com/apache/arrow/go/v8/parquet/file"
|
|
"github.com/milvus-io/milvus-proto/go-api/schemapb"
|
|
"github.com/stretchr/testify/suite"
|
|
)
|
|
|
|
type ReadDataFromAllRowGroupsSuite struct {
|
|
suite.Suite
|
|
size int
|
|
|
|
logData []byte
|
|
|
|
reader *PayloadReader
|
|
}
|
|
|
|
func (s *ReadDataFromAllRowGroupsSuite) SetupSuite() {
|
|
w := NewIndexFileBinlogWriter(0, 0, 1, 2, 3, 100, "", 0, "test")
|
|
defer w.Close()
|
|
// make sure it's still written int8 data
|
|
w.PayloadDataType = schemapb.DataType_Int8
|
|
ew, err := w.NextIndexFileEventWriter()
|
|
s.Require().NoError(err)
|
|
defer ew.Close()
|
|
|
|
s.size = 1 << 10
|
|
|
|
data := make([]byte, s.size)
|
|
err = ew.AddByteToPayload(data)
|
|
s.Require().NoError(err)
|
|
|
|
ew.SetEventTimestamp(1, 1)
|
|
w.SetEventTimeStamp(1, 1)
|
|
|
|
w.AddExtra(originalSizeKey, fmt.Sprintf("%v", len(data)))
|
|
|
|
err = w.Finish()
|
|
s.Require().NoError(err)
|
|
|
|
buffer, err := w.GetBuffer()
|
|
s.Require().NoError(err)
|
|
|
|
s.logData = buffer
|
|
}
|
|
|
|
func (s *ReadDataFromAllRowGroupsSuite) TearDownSuite() {}
|
|
|
|
func (s *ReadDataFromAllRowGroupsSuite) SetupTest() {
|
|
br, err := NewBinlogReader(s.logData)
|
|
s.Require().NoError(err)
|
|
er, err := br.NextEventReader()
|
|
s.Require().NoError(err)
|
|
|
|
reader, ok := er.PayloadReaderInterface.(*PayloadReader)
|
|
s.Require().True(ok)
|
|
|
|
s.reader = reader
|
|
}
|
|
|
|
func (s *ReadDataFromAllRowGroupsSuite) TearDownTest() {
|
|
s.reader.Close()
|
|
s.reader = nil
|
|
}
|
|
|
|
func (s *ReadDataFromAllRowGroupsSuite) TestNormalRun() {
|
|
values := make([]int32, s.size)
|
|
valuesRead, err := ReadDataFromAllRowGroups[int32, *file.Int32ColumnChunkReader](s.reader.reader, values, 0, int64(s.size))
|
|
s.Assert().NoError(err)
|
|
s.Assert().EqualValues(s.size, valuesRead)
|
|
}
|
|
|
|
func TestReadDataFromAllRowGroupsSuite(t *testing.T) {
|
|
suite.Run(t, new(ReadDataFromAllRowGroupsSuite))
|
|
}
|