seaweedfs/weed/mq/schema/write_parquet_test.go

154 lines
4.0 KiB
Go
Raw Normal View History

2024-04-19 13:41:20 +08:00
package schema
import (
"fmt"
2024-04-22 15:42:18 +08:00
"github.com/parquet-go/parquet-go"
"github.com/parquet-go/parquet-go/compress/zstd"
"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
"io"
"os"
2024-04-19 13:41:20 +08:00
"testing"
)
2024-04-25 23:54:57 +08:00
func TestWriteReadParquet(t *testing.T) {
2024-04-19 13:41:20 +08:00
// create a schema_pb.RecordType
2024-04-29 04:00:52 +08:00
recordType := RecordTypeBegin().
2024-05-02 23:32:15 +08:00
WithField("ID", TypeLong).
WithField("CreatedAt", TypeLong).
WithRecordField("Person",
2024-04-29 04:00:52 +08:00
RecordTypeBegin().
2024-05-02 23:32:15 +08:00
WithField("zName", TypeString).
WithField("emails", ListOf(TypeString)).
2024-04-29 04:00:52 +08:00
RecordTypeEnd()).
2024-05-02 23:32:15 +08:00
WithField("Company", TypeString).
WithRecordField("Address",
2024-04-29 04:00:52 +08:00
RecordTypeBegin().
2024-05-02 23:32:15 +08:00
WithField("Street", TypeString).
WithField("City", TypeString).
2024-04-29 04:00:52 +08:00
RecordTypeEnd()).
RecordTypeEnd()
2024-04-19 13:41:20 +08:00
fmt.Printf("RecordType: %v\n", recordType)
// create a parquet schema
parquetSchema, err := ToParquetSchema("example", recordType)
if err != nil {
t.Fatalf("ToParquetSchema failed: %v", err)
}
fmt.Printf("ParquetSchema: %v\n", parquetSchema)
2024-04-25 14:04:47 +08:00
fmt.Printf("Go Type: %+v\n", parquetSchema.GoType())
2024-04-19 13:41:20 +08:00
2024-04-22 15:42:18 +08:00
filename := "example.parquet"
2024-04-25 14:24:53 +08:00
count := 3
testWritingParquetFile(t, count, filename, parquetSchema, recordType)
2024-04-22 15:42:18 +08:00
total := testReadingParquetFile(t, filename, parquetSchema, recordType)
2024-04-25 14:24:53 +08:00
if total != count {
2024-04-22 15:42:18 +08:00
t.Fatalf("total != 128*1024: %v", total)
}
2024-04-25 14:29:55 +08:00
if err = os.Remove(filename); err != nil {
t.Fatalf("os.Remove failed: %v", err)
}
2024-04-22 15:42:18 +08:00
}
2024-04-25 14:24:53 +08:00
func testWritingParquetFile(t *testing.T, count int, filename string, parquetSchema *parquet.Schema, recordType *schema_pb.RecordType) {
2024-04-26 14:59:30 +08:00
parquetLevels, err := ToParquetLevels(recordType)
if err != nil {
t.Fatalf("ToParquetLevels failed: %v", err)
}
2024-04-22 15:42:18 +08:00
// create a parquet file
file, err := os.OpenFile(filename, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0664)
if err != nil {
t.Fatalf("os.Open failed: %v", err)
}
defer file.Close()
2024-04-26 13:02:03 +08:00
writer := parquet.NewWriter(file, parquetSchema, parquet.Compression(&zstd.Codec{Level: zstd.DefaultLevel}))
2024-04-22 15:42:18 +08:00
rowBuilder := parquet.NewRowBuilder(parquetSchema)
2024-04-25 14:24:53 +08:00
for i := 0; i < count; i++ {
2024-04-22 15:42:18 +08:00
rowBuilder.Reset()
// generate random data
2024-04-29 04:42:45 +08:00
recordValue := RecordBegin().
SetLong("ID", 1+int64(i)).
SetLong("CreatedAt", 2+2*int64(i)).
SetRecord("Person",
RecordBegin().
SetString("zName", fmt.Sprintf("john_%d", i)).
SetStringList("emails",
fmt.Sprintf("john_%d@a.com", i),
fmt.Sprintf("john_%d@b.com", i),
fmt.Sprintf("john_%d@c.com", i),
fmt.Sprintf("john_%d@d.com", i),
fmt.Sprintf("john_%d@e.com", i)).
RecordEnd()).
SetString("Company", fmt.Sprintf("company_%d", i)).
2024-04-29 04:56:55 +08:00
RecordEnd()
2024-04-26 14:59:30 +08:00
AddRecordValue(rowBuilder, recordType, parquetLevels, recordValue)
2024-04-25 14:04:47 +08:00
2024-04-25 23:54:57 +08:00
if count < 10 {
fmt.Printf("RecordValue: %v\n", recordValue)
}
2024-04-22 15:42:18 +08:00
row := rowBuilder.Row()
2024-04-25 14:04:47 +08:00
2024-04-25 23:54:57 +08:00
if count < 10 {
fmt.Printf("Row: %+v\n", row)
}
2024-04-25 14:04:47 +08:00
2024-04-22 15:42:18 +08:00
if err != nil {
t.Fatalf("rowBuilder.Build failed: %v", err)
}
if _, err = writer.WriteRows([]parquet.Row{row}); err != nil {
t.Fatalf("writer.Write failed: %v", err)
}
}
if err = writer.Close(); err != nil {
t.Fatalf("writer.WriteStop failed: %v", err)
}
}
func testReadingParquetFile(t *testing.T, filename string, parquetSchema *parquet.Schema, recordType *schema_pb.RecordType) (total int) {
2024-04-26 14:59:30 +08:00
parquetLevels, err := ToParquetLevels(recordType)
if err != nil {
t.Fatalf("ToParquetLevels failed: %v", err)
}
2024-04-22 15:42:18 +08:00
// read the parquet file
file, err := os.Open(filename)
if err != nil {
t.Fatalf("os.Open failed: %v", err)
}
defer file.Close()
reader := parquet.NewReader(file, parquetSchema)
rows := make([]parquet.Row, 128)
for {
rowCount, err := reader.ReadRows(rows)
if err != nil {
if err == io.EOF {
break
}
t.Fatalf("reader.Read failed: %v", err)
}
for i := 0; i < rowCount; i++ {
row := rows[i]
// convert parquet row to schema_pb.RecordValue
2024-04-26 14:59:30 +08:00
recordValue, err := ToRecordValue(recordType, parquetLevels, row)
2024-04-22 15:42:18 +08:00
if err != nil {
t.Fatalf("ToRecordValue failed: %v", err)
}
2024-04-25 14:24:53 +08:00
if rowCount < 10 {
fmt.Printf("RecordValue: %v\n", recordValue)
}
2024-04-22 15:42:18 +08:00
}
total += rowCount
}
fmt.Printf("total: %v\n", total)
return
2024-04-19 13:41:20 +08:00
}