Go言語でnull値を含むParquetファイルを生成する

環境

• macOS Catalina
• Go 1.14
• parquet-go 1.5.1
• pyarrow 0.16.0（確認用）
• pandas 1.0.3（確認用）

Parquetファイルの生成

このソースコードをmain.goとして保存する。

package main

import (
	"log"

	"github.com/xitongsys/parquet-go-source/local"
	"github.com/xitongsys/parquet-go/writer"
)

type Example struct {
	Col *string `parquet:"name=col, type=UTF8"`
}

func main() {
	fw, err := local.NewLocalFileWriter("example.parquet")
	if err != nil {
		log.Fatal(err)
	}

	pw, err := writer.NewParquetWriter(fw, new(Example), 4)
	if err != nil {
		log.Fatal(err)
	}

	row0 := "foo"
	row2 := "bar"
	rows := []*string{&row0, nil, &row2}
	for i := 0; i < 3; i++ {
		ex := Example{Col: rows[i]}
		if err := pw.Write(ex); err != nil {
			log.Fatal(err)
		}
	}
	if err = pw.WriteStop(); err != nil {
		log.Fatal(err)
	}
	fw.Close()
}

null値を含むParquetファイルを生成する上で重要なのは、下記行である。

	Col *string `parquet:"name=col, type=UTF8"`

つまり、stringではなく*stringを指定することで、parquet-goのREADME.mdにある通りRepetition TypeをOPTIONALにする。

作成したファイルに対してgo run main.goを実行すると、example.parquetが生成される。

Parquetファイルの確認

下記のPythonスクリプトを使えば、確かにnull値を含むParquetファイルが生成されたことを確認できる。

import pyarrow.parquet as pq

table = pq.read_table('example.parquet')
print(table.to_pandas())

このファイルをread.pyとして保存し、python3 read.pyを実行すると

    col
0   foo
1  None
2   bar

と出力される。

もちろんGo言語でも、このexample.parquetを読み込める。

package main

import (
	"fmt"
	"log"

	"github.com/xitongsys/parquet-go-source/local"
	"github.com/xitongsys/parquet-go/reader"
)

type Example struct {
	Col *string `parquet:"name=col", type=UTF8`
}

func main() {
	fr, err := local.NewLocalFileReader("example.parquet")
	if err != nil {
		log.Fatal(err)
	}
	pr, err := reader.NewParquetReader(fr, new(Example), 4)
	if err != nil {
		log.Fatal(err)
	}
	num := int(pr.GetNumRows())
	ex := make([]Example, num)
	if err := pr.Read(&ex); err != nil {
		log.Fatal(err)
	}
	for i := 0; i < num; i++ {
		col := ex[i].Col
		if col != nil {
			fmt.Println(*col)
		} else {
			fmt.Println("<nil>")
		}
	}
	pr.ReadStop()
	fr.Close()
}

このファイルをmain.goとして保存し、go run main.goを実行すると

foo
<nil>
bar

と出力される。