parquet_example.go

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"strconv"

	"github.com/apache/arrow/go/v7/arrow"
	"github.com/apache/arrow/go/v7/arrow/array"
	"github.com/apache/arrow/go/v7/arrow/memory"
	"github.com/apache/arrow/go/v7/parquet"
	"github.com/apache/arrow/go/v7/parquet/pqarrow"
)

func main() {

	// profile
	f, _ := os.Create("profile.pb.gz")
	defer f.Close()
	runtime.GC()
	pprof.StartCPUProfile(f)
	defer pprof.StopCPUProfile()

	// open a file
	file, err := ioutil.TempFile(".", "out.*.parquet")
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()
	fmt.Println(file.Name())

	// example schema
	schema := arrow.NewSchema(
		[]arrow.Field{
			{Name: "stringer", Type: &arrow.StringType{}},
		},
		nil,
	)

	// Setup a FileWriter
	// https://pkg.go.dev/github.com/apache/arrow/go/[email protected]/parquet/pqarrow#FileWriter
	fw, err := pqarrow.NewFileWriter(
		schema,
		file,
		parquet.NewWriterProperties(),
		pqarrow.DefaultWriterProps(),
	)
	defer fw.Close()

	// Setup a builder
	// https://pkg.go.dev/github.com/apache/arrow/go/[email protected]/arrow/array#RecordBuilder
	builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
	defer builder.Release()

	// example ETL
	for i := 0; i < 10; i++ {

		// Append values via builder (see docs)
		// https://pkg.go.dev/github.com/apache/arrow/go/arrow#example-package-Table
		str := strconv.Itoa(i)
		builder.Field(0).(*array.StringBuilder).AppendValues([]string{str}, nil)

		// get records from builder
		r := builder.NewRecord()
		defer r.Release()

		// write
		if err := fw.WriteBuffered(r); err != nil {
			log.Fatalf("failed to write: %s", err)
		}

		// WARN: Don't do this, every write creates a new RowGroup
		// if err := fw.Write(r); err != nil {
		// 	log.Fatalf("failed to write: %s", err)
		// }

	}

}