forked from pingcap/br
-
Notifications
You must be signed in to change notification settings - Fork 0
/
parquet.go
65 lines (56 loc) · 1.38 KB
/
parquet.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
package main
import (
"flag"
"fmt"
"log"
"path/filepath"
"strconv"
"github.com/xitongsys/parquet-go-source/local"
"github.com/xitongsys/parquet-go/writer"
)
var (
schema = flag.String("schema", "test", "Test schema name")
table = flag.String("table", "parquet", "Test table name")
chunks = flag.Int("chunk", 10, "Chunk files count")
rowNumbers = flag.Int("rows", 1000, "Row number for each test file")
sourceDir = flag.String("dir", "", "test directory path")
)
func genParquetFile(dir, name string, count int) error {
type Test struct {
I int32 `parquet:"name=iVal, type=INT32"`
S string `parquet:"name=s, type=UTF8, encoding=PLAIN_DICTIONARY"`
}
w, err := local.NewLocalFileWriter(filepath.Join(dir, name))
if err != nil {
return err
}
test := &Test{}
dataWriter, err := writer.NewParquetWriter(w, test, 2)
if err != nil {
return err
}
for i := 0; i < count; i++ {
test.I = int32(i)
test.S = strconv.Itoa(i)
err := dataWriter.Write(test)
if err != nil {
return err
}
}
err = dataWriter.WriteStop()
if err != nil {
return err
}
w.Close()
return nil
}
func main() {
flag.Parse()
for i := 0; i < *chunks; i++ {
name := fmt.Sprintf("%s.%s.%04d.parquet", *schema, *table, i)
err := genParquetFile(*sourceDir, name, *rowNumbers)
if err != nil {
log.Fatalf("generate test source failed, name: %s, err: %+v", name, err)
}
}
}