diff --git a/go/packed/packed_test.go b/go/packed/packed_test.go index 731b2e4..a1612ad 100644 --- a/go/packed/packed_test.go +++ b/go/packed/packed_test.go @@ -24,67 +24,67 @@ import ( "golang.org/x/exp/rand" ) -// func TestPackedOneFile(t *testing.T) { -// batches := 100 -// schema := arrow.NewSchema([]arrow.Field{ -// {Name: "a", Type: arrow.PrimitiveTypes.Int32}, -// {Name: "b", Type: arrow.PrimitiveTypes.Int64}, -// {Name: "c", Type: arrow.BinaryTypes.String}, -// }, nil) +func TestPackedOneFile(t *testing.T) { + batches := 100 + schema := arrow.NewSchema([]arrow.Field{ + {Name: "a", Type: arrow.PrimitiveTypes.Int32}, + {Name: "b", Type: arrow.PrimitiveTypes.Int64}, + {Name: "c", Type: arrow.BinaryTypes.String}, + }, nil) -// b := array.NewRecordBuilder(memory.DefaultAllocator, schema) -// defer b.Release() -// for idx := range schema.Fields() { -// switch idx { -// case 0: -// b.Field(idx).(*array.Int32Builder).AppendValues( -// []int32{int32(1), int32(2), int32(3)}, nil, -// ) -// case 1: -// b.Field(idx).(*array.Int64Builder).AppendValues( -// []int64{int64(4), int64(5), int64(6)}, nil, -// ) -// case 2: -// b.Field(idx).(*array.StringBuilder).AppendValues( -// []string{"a", "b", "c"}, nil, -// ) -// } -// } -// rec := b.NewRecord() -// defer rec.Release() -// path := "/tmp" -// bufferSize := 10 * 1024 * 1024 // 10MB -// pw, err := newPackedWriter(path, schema, bufferSize) -// assert.NoError(t, err) -// for i := 0; i < batches; i++ { -// err = pw.writeRecordBatch(rec) -// assert.NoError(t, err) -// } -// err = pw.close() -// assert.NoError(t, err) + b := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer b.Release() + for idx := range schema.Fields() { + switch idx { + case 0: + b.Field(idx).(*array.Int32Builder).AppendValues( + []int32{int32(1), int32(2), int32(3)}, nil, + ) + case 1: + b.Field(idx).(*array.Int64Builder).AppendValues( + []int64{int64(4), int64(5), int64(6)}, nil, + ) + case 2: + b.Field(idx).(*array.StringBuilder).AppendValues( + []string{"a", "b", "c"}, nil, + ) + } + } + rec := b.NewRecord() + defer rec.Release() + path := "/tmp" + bufferSize := 10 * 1024 * 1024 // 10MB + pw, err := newPackedWriter(path, schema, bufferSize) + assert.NoError(t, err) + for i := 0; i < batches; i++ { + err = pw.writeRecordBatch(rec) + assert.NoError(t, err) + } + err = pw.close() + assert.NoError(t, err) -// readerOption := NewPackedReaderOption() -// reader, err := newPackedReader(path, schema, bufferSize, readerOption) -// assert.NoError(t, err) -// rr, err := reader.readNext() -// assert.NoError(t, err) -// defer rr.Release() -// assert.Equal(t, int64(3*batches), rr.NumRows()) + readerOption := NewPackedReaderOption() + reader, err := newPackedReader(path, schema, bufferSize, readerOption) + assert.NoError(t, err) + rr, err := reader.readNext() + assert.NoError(t, err) + defer rr.Release() + assert.Equal(t, int64(3*batches), rr.NumRows()) -// // test packed partial read -// readerOption = NewPackedReaderOption() -// readerOption.WithNeededColumns([]int{1}) -// schema = arrow.NewSchema([]arrow.Field{ -// {Name: "b", Type: arrow.PrimitiveTypes.Int64}, -// }, nil) -// reader, err = newPackedReader(path, schema, bufferSize, readerOption) -// assert.NoError(t, err) -// rr, err = reader.readNext() -// assert.Equal(t, int64(1), rr.NumCols()) -// assert.Equal(t, int64(3*batches), rr.NumRows()) -// assert.NoError(t, err) -// defer rr.Release() -// } + // test packed partial read + readerOption = NewPackedReaderOption() + readerOption.WithNeededColumns([]int{1}) + schema = arrow.NewSchema([]arrow.Field{ + {Name: "b", Type: arrow.PrimitiveTypes.Int64}, + }, nil) + reader, err = newPackedReader(path, schema, bufferSize, readerOption) + assert.NoError(t, err) + rr, err = reader.readNext() + assert.Equal(t, int64(1), rr.NumCols()) + assert.Equal(t, int64(3*batches), rr.NumRows()) + assert.NoError(t, err) + defer rr.Release() +} func TestPackedMultiFiles(t *testing.T) { batches := 1000