fix ut

Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Jan 10, 2025 · e3459ef · e3459ef
1 parent c850086
commit e3459ef
Showing 1 changed file with 58 additions and 58 deletions.
diff --git 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