github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/tests/lightning_checkpoint_parquet/parquet.go (about)

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