github.com/fraugster/parquet-go@v0.12.0/compatibility/build.go (about)

     1  // +build ignore
     2  
     3  package main
     4  
     5  import (
     6  	"flag"
     7  	"os"
     8  	"strings"
     9  
    10  	"github.com/fraugster/parquet-go/parquetschema"
    11  
    12  	goparquet "github.com/fraugster/parquet-go"
    13  	"github.com/fraugster/parquet-go/parquet"
    14  )
    15  
    16  func main() {
    17  	var (
    18  		file    string
    19  		pq      string
    20  		version string
    21  		comp    string
    22  	)
    23  	flag.StringVar(&file, "json", "/data.json", "json file to load")
    24  	flag.StringVar(&pq, "pq", "/data.pq", "pq to save")
    25  	flag.StringVar(&version, "version", "v1", "Page v1 or Page v2 (v1 / v2)")
    26  	flag.StringVar(&comp, "compression", "snappy", "compression method, snappy, gzip, none")
    27  
    28  	flag.Parse()
    29  
    30  	data, err := loadDataFromJson(file)
    31  	if err != nil {
    32  		panic(err)
    33  	}
    34  
    35  	var opts []goparquet.FileWriterOption
    36  	switch strings.ToUpper(comp) {
    37  	case "SNAPPY":
    38  		opts = append(opts, goparquet.WithCompressionCodec(parquet.CompressionCodec_SNAPPY))
    39  	case "GZIP":
    40  		opts = append(opts, goparquet.WithCompressionCodec(parquet.CompressionCodec_GZIP))
    41  	case "NONE":
    42  		opts = append(opts, goparquet.WithCompressionCodec(parquet.CompressionCodec_UNCOMPRESSED))
    43  	default:
    44  		panic("invalid codec: " + comp)
    45  	}
    46  
    47  	switch strings.ToUpper(version) {
    48  	case "V1":
    49  		// it is default
    50  	case "V2":
    51  		opts = append(opts, goparquet.WithDataPageV2())
    52  	default:
    53  		panic("invalid version: " + version)
    54  	}
    55  
    56  	sc, err := parquetschema.ParseSchemaDefinition(schema)
    57  	if err != nil {
    58  		panic(err)
    59  	}
    60  
    61  	opts = append(opts, goparquet.WithSchemaDefinition(sc))
    62  
    63  	fl, err := os.Create(pq)
    64  	if err != nil {
    65  		panic(err)
    66  	}
    67  
    68  	writer := goparquet.NewFileWriter(fl, opts...)
    69  	for i := range data {
    70  		if err := writer.AddData(data[i].toMap()); err != nil {
    71  			panic(err)
    72  		}
    73  	}
    74  
    75  	if err := writer.Close(); err != nil {
    76  		panic(err)
    77  	}
    78  }