Skip to content

Commit

Permalink
add ut
Browse files Browse the repository at this point in the history
Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
  • Loading branch information
shaoting-huang committed Jul 3, 2024
1 parent 76cf96a commit ec2e083
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 1 deletion.
3 changes: 2 additions & 1 deletion internal/storage/serde.go
Original file line number Diff line number Diff line change
Expand Up @@ -676,9 +676,10 @@ func (sfw *singleFieldRecordWriter) Close() {

func newSingleFieldRecordWriter(fieldId FieldID, field arrow.Field, writer io.Writer) (*singleFieldRecordWriter, error) {
schema := arrow.NewSchema([]arrow.Field{field}, nil)

// use writer properties as same as payload writer's for now
fw, err := pqarrow.NewFileWriter(schema, writer,
parquet.NewWriterProperties(
parquet.WithMaxRowGroupLength(math.MaxInt64), // No additional grouping for now.
parquet.WithCompression(compress.Codecs.Zstd),
parquet.WithCompressionLevel(3)),
pqarrow.DefaultWriterProps())
Expand Down
18 changes: 18 additions & 0 deletions internal/storage/serde_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (
"reflect"
"testing"

"github.com/apache/arrow/go/v12/arrow"
"github.com/apache/arrow/go/v12/arrow/array"
"github.com/apache/arrow/go/v12/arrow/memory"
"github.com/stretchr/testify/assert"
Expand Down Expand Up @@ -100,6 +101,23 @@ func TestSerDe(t *testing.T) {
}
}

func TestArrowSchema(t *testing.T) {
t.Run("test composite record", func(t *testing.T) {
fields := []arrow.Field{{Name: "1", Type: arrow.BinaryTypes.String, Nullable: true}}
builder := array.NewBuilder(memory.DefaultAllocator, arrow.BinaryTypes.String)
builder.AppendValueFromString("1")
arrays := []arrow.Array{builder.NewArray()}
cr := &compositeRecord{
recs: make(map[FieldID]arrow.Record, 1),
schema: make(map[FieldID]schemapb.DataType, 1),
}
cr.recs[0] = array.NewRecord(arrow.NewSchema(fields, nil), arrays, 1)
cr.schema[0] = schemapb.DataType_String
expected := arrow.NewSchema(fields, nil)
assert.Equal(t, expected, cr.ArrowSchema())
})
}

func BenchmarkDeserializeReader(b *testing.B) {
len := 1000000
blobs, err := generateTestData(len)
Expand Down

0 comments on commit ec2e083

Please sign in to comment.