【Go言語】Apache Arrowを使ってParquetファイルを書き込む

先日、【Go言語】Apache Arrowを使ってParquetファイルを読み込む という記事を書いたのですが、こちらではParquetファイルの書き込みについて簡単に説明しています。

このParquetファイルはなぜかVisual Studio Codeのプラグインでは開けませんが、Parquet ViewerやAmazon Athenaから正常に読み込めたので実用に問題がないと判断し、公開することにしました。

しかし、現在このライブラリではRecordBuilderへ追加した整数型、浮動小数点型データをAppendValuesの第2引数で無効とした場合（値はnull値になる）、Parquetファイルに書き込まれるときにこの無効としたことが無視されすべて記録されてしまいます。つまりnull値として書き込まれません。この点はご注意ください。

インポートするライブラリは以下のとおりです。

import (
	"os"

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

NewRecordBuilderでRecord Batchを作成し、それをparquetファイル形式で書き込んでいます。

pool := memory.NewGoAllocator()

// スキーマの作成
schema := arrow.NewSchema(
	[]arrow.Field{
		{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
		{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
	},
	nil,
)

// Record Builderの作成
b := array.NewRecordBuilder(pool, schema)
defer b.Release()

// データの追加
b.Field(0).(*array.Int32Builder).AppendValues([]int32{10, 9, 8, 7, 6, 5, 4, 3, 2, 1}, nil)
b.Field(1).(*array.Float64Builder).AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)

// Record Batchの作成
rec := b.NewRecord()
defer rec.Release()

// 書き込みファイルの作成
f, err := os.Create("output.parquet")
if err != nil {
    panic(err)
}
defer f.Close()

// NewFileWriterの準備
w, err := pqarrow.NewFileWriter(rec.Schema(), f, nil, pqarrow.DefaultWriterProps())
if err != nil {
    panic(err)
}
defer w.Close()

// Record Batchの書き込み
err = w.Write(rec)
if err != nil {
    panic(err)
}

こちらは書き込んだ Record Batch の内容を出力したものになります。

record:
  schema:
  fields: 2
    - f1-i32: type=int32
    - f2-f64: type=float64
  rows: 10
  col[0][f1-i32]: [10 9 8 7 6 5 4 3 2 1]
  col[1][f2-f64]: [1 2 3 4 5 6 7 8 9 10]

作成された output.parquet ファイルをParquet Viewerで表示すると以下のようになります。

Record Batchができてしまえば非常に簡単に書き込むことができます。

最後に

前回の記事とあわせ、Apache ArrowでParquetファイルの読み込み、書き込みに対応することができました。しかし、Parquetファイルでは整数型、浮動小数点型データをnull値にすることがサポートされてますが、このライブラリではnullの取り扱いに注意が必要となります。（元からnullを含むRecord Batchを書き込む場合は問題ありません）

Apache Arrowライブラリは頻繁に更新されているため、今後の対応に期待したいところです。