github.com/fraugster/parquet-go@v0.12.0/examples/high-level-custom-marshalling/main.go (about)

     1  package main
     2  
     3  import (
     4  	"log"
     5  	"reflect"
     6  
     7  	goparquet "github.com/fraugster/parquet-go"
     8  	"github.com/fraugster/parquet-go/floor"
     9  	"github.com/fraugster/parquet-go/floor/interfaces"
    10  	"github.com/fraugster/parquet-go/parquet"
    11  	"github.com/fraugster/parquet-go/parquetschema"
    12  )
    13  
    14  func main() {
    15  	schemaDef, err := parquetschema.ParseSchemaDefinition(
    16  		`message test {
    17  			required binary name (STRING);
    18  			required binary data;
    19  			required double score;
    20  		}`)
    21  	if err != nil {
    22  		log.Fatalf("Parsing schema definition failed: %v", err)
    23  	}
    24  
    25  	parquetFilename := "output.parquet"
    26  
    27  	fw, err := floor.NewFileWriter(parquetFilename,
    28  		goparquet.WithSchemaDefinition(schemaDef),
    29  		goparquet.WithCompressionCodec(parquet.CompressionCodec_SNAPPY),
    30  	)
    31  	if err != nil {
    32  		log.Fatalf("Opening parquet file for writing failed: %v", err)
    33  	}
    34  
    35  	input := []*record{
    36  		{
    37  			n: "Test",
    38  			d: []byte{0xFF, 0x0A, 0x8E, 0x00, 0x12},
    39  			s: 23.5,
    40  		},
    41  	}
    42  
    43  	for _, rec := range input {
    44  		if err := fw.Write(rec); err != nil {
    45  			log.Fatalf("Writing record failed: %v", err)
    46  		}
    47  	}
    48  
    49  	if err := fw.Close(); err != nil {
    50  		log.Fatalf("Closing parquet writer failed: %v", err)
    51  	}
    52  
    53  	fr, err := floor.NewFileReader(parquetFilename)
    54  	if err != nil {
    55  		log.Fatalf("Opening parquet file failed: %v", err)
    56  	}
    57  
    58  	var fileContent []*record
    59  
    60  	for fr.Next() {
    61  		rec := &record{}
    62  		if err := fr.Scan(rec); err != nil {
    63  			log.Fatalf("Scanning record failed: %v", err)
    64  		}
    65  		fileContent = append(fileContent, rec)
    66  	}
    67  
    68  	equal := reflect.DeepEqual(input, fileContent)
    69  	if equal {
    70  		log.Printf("Congratulations! The input and the data read back are identical!")
    71  	} else {
    72  		log.Printf("This is strange... the data read back does not back what has been written to the file.")
    73  	}
    74  }
    75  
    76  type record struct {
    77  	n string
    78  	d []byte
    79  	s float64
    80  }
    81  
    82  func (r *record) MarshalParquet(obj interfaces.MarshalObject) error {
    83  	obj.AddField("name").SetByteArray([]byte(r.n))
    84  	obj.AddField("data").SetByteArray(r.d)
    85  	obj.AddField("score").SetFloat64(r.s)
    86  	return nil
    87  }
    88  
    89  func (r *record) UnmarshalParquet(obj interfaces.UnmarshalObject) error {
    90  	name, err := obj.GetField("name").ByteArray()
    91  	if err != nil {
    92  		return nil
    93  	}
    94  	r.n = string(name)
    95  
    96  	data, err := obj.GetField("data").ByteArray()
    97  	if err != nil {
    98  		return nil
    99  	}
   100  	r.d = data
   101  
   102  	score, err := obj.GetField("score").Float64()
   103  	if err != nil {
   104  		return nil
   105  	}
   106  	r.s = score
   107  
   108  	return nil
   109  }