Browse Source

adjust test

mq-subscribe
chrislu 8 months ago
parent
commit
f93c4cf48c
  1. 16
      weed/mq/schema/write_parquet_test.go

16
weed/mq/schema/write_parquet_test.go

@ -32,17 +32,19 @@ func TestWriteParquet(t *testing.T) {
filename := "example.parquet"
testWritingParquetFile(t, filename, parquetSchema, recordType)
count := 3
testWritingParquetFile(t, count, filename, parquetSchema, recordType)
total := testReadingParquetFile(t, filename, parquetSchema, recordType)
if total != 128*1024 {
if total != count {
t.Fatalf("total != 128*1024: %v", total)
}
}
func testWritingParquetFile(t *testing.T, filename string, parquetSchema *parquet.Schema, recordType *schema_pb.RecordType) {
func testWritingParquetFile(t *testing.T, count int, filename string, parquetSchema *parquet.Schema, recordType *schema_pb.RecordType) {
// create a parquet file
file, err := os.OpenFile(filename, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0664)
if err != nil {
@ -51,7 +53,7 @@ func testWritingParquetFile(t *testing.T, filename string, parquetSchema *parque
defer file.Close()
writer := parquet.NewWriter(file, parquetSchema, parquet.Compression(&zstd.Codec{Level: zstd.SpeedDefault}))
rowBuilder := parquet.NewRowBuilder(parquetSchema)
for i := 0; i < 128*1024; i++ {
for i := 0; i < count; i++ {
rowBuilder.Reset()
// generate random data
recordValue := NewRecordValueBuilder().
@ -107,11 +109,13 @@ func testReadingParquetFile(t *testing.T, filename string, parquetSchema *parque
for i := 0; i < rowCount; i++ {
row := rows[i]
// convert parquet row to schema_pb.RecordValue
_, err := ToRecordValue(recordType, row)
recordValue, err := ToRecordValue(recordType, row)
if err != nil {
t.Fatalf("ToRecordValue failed: %v", err)
}
// fmt.Printf("RecordValue: %v\n", recordValue)
if rowCount < 10 {
fmt.Printf("RecordValue: %v\n", recordValue)
}
}
total += rowCount
}

Loading…
Cancel
Save