Use MapReduce to repartition parquet files in S3

mapreduce_repartition.go

package main

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/in4it/gomap/pkg/context"
	"github.com/in4it/gomap/pkg/types"
	"github.com/in4it/gomap/pkg/utils"
	"github.com/xitongsys/parquet-go-source/local"
	"github.com/xitongsys/parquet-go/writer"
)

type DataRecord struct {
	Id        *int32  `parquet:"name=id, type=INT32, repetitiontype=OPTIONAL"`
	Firstname *string `parquet:"name=firstname, type=UTF8, repetitiontype=OPTIONAL"`
	Lastname  *string `parquet:"name=lastname, type=UTF8, repetitiontype=OPTIONAL"`
	Street    *string `parquet:"name=street, type=UTF8, repetitiontype=OPTIONAL"`
	Zipcode   *string `parquet:"name=zipcode, type=UTF8, repetitiontype=OPTIONAL"`
}
type DataRecordPartitioned struct {
	Id        *int32  `parquet:"name=id, type=INT32, repetitiontype=OPTIONAL"`
	Firstname *string `parquet:"name=firstname, type=UTF8, repetitiontype=OPTIONAL"`
	Lastname  *string `parquet:"name=lastname, type=UTF8, repetitiontype=OPTIONAL"`
	Street    *string `parquet:"name=street, type=UTF8, repetitiontype=OPTIONAL"`
}

func repartitionParquet(bucket string) {
	c := context.New()
	var t DataRecordPartitioned
	c.ReadParquet(bucket, new(DataRecord)).MapToKV(func(input types.RawInput) (types.RawOutput, types.RawOutput) {
		var line DataRecord
		err := utils.RawDecode(input, &line)
		if err != nil {
			panic(err)
		}
		return utils.StringToRawOutput(*line.Zipcode), utils.RawEncode([]DataRecord{line})
	}).ReduceByKey(func(a, b types.RawInput) types.RawOutput {
		var line1 []DataRecord
		var line2 []DataRecord
		err := utils.RawDecode(a, &line1)
		if err != nil {
			panic(err)
		}
		err = utils.RawDecode(b, &line2)
		if err != nil {
			panic(err)
		}
		return utils.RawEncode(append(line1, line2...))
	}).Run().Foreach(func(key, value types.RawOutput) {
		var lines []DataRecord

		x := time.Now()
		filename := "data-partitioned-" + fmt.Sprintf("%v", (x.Format("20060102150405"))) + ".parquet"

		err := utils.RawDecode(value, &lines)
		if err != nil {
			panic(err)
		}
		fmt.Println("key: " + string(key))

		fw, err := local.NewLocalFileWriter("/tmp/" + filename)
		if err != nil {
			log.Println("Can't create local file", err)
			return
		}

		pw, err := writer.NewParquetWriter(fw, new(DataRecordPartitioned), 4)
		if err != nil {
			log.Println("Can't create parquet writer", err)
			return
		}

		for _, l := range lines {
			t = DataRecordPartitioned{
				Id:        l.Id,
				Firstname: l.Firstname,
				Lastname:  l.Lastname,
				Street:    l.Street,
			}
			if err = pw.Write(t); err != nil {
				log.Println("Write error", err)
			}
		}

		if err = pw.WriteStop(); err != nil {
			log.Println("WriteStop error", err)
			return
		}
		log.Println("Write Finished")
		fw.Close()
		os.Remove(filename)
	})
}