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

以前、あるシステムのGo言語のバージョンをアップするとともに古くなったParquetの読み込みライブラリも変更しようと話になりました。
Go言語のParquetライブラリについて調べると「Apache Arrow」、「xitongsys/parquet-go」の2つが候補になると思いますが、今回は現在でもバージョンアップがされていて、信頼性と将来性から「Apache Arrow」を選択しました。
こちらはまだ未完成な部分がありParquetファイルの書き込みには対応していないようですが、Parquet自体Apacheが開発したものですので今後に期待しても良いでしょう。

追記：Parquetファイルでは整数型・浮動小数点型のカラムにnullのデータを持つことが許されています。このライブラリにおいてRecord Batchの時点ではnullは保持されていますが、下記の例による「カラムデータの各要素の取得」でnullを取得した場合、0か1の数字が勝手に挿入されてしまうようです。この場合は整数型・浮動小数点型のカラムを文字列型に変換して値を取り出すような対策が必要となります。また、Parquetファイルの書き込みに対応していたので別の記事で公開予定です。

Apache Arrowを扱うための基礎知識

今回の記事では軽く触れる程度ですが、Apache Arrowには以下のようなデータ形式があります。

• Array : 配列。
• Chunked Array : いくつかの配列をひとまとめにして1つの配列として扱えるようにしたもの。Chunked Arrayに含まれる配列の長さは異なっていてもいいが、すべて同じ型でなければならない。
• Record Batch : 一連の配列で構成されるテーブルのようなデータ構造。配列はそれぞれ異なる型にすることはできるがすべて同じ長さでなければならない。なお、Apache Arrowのソースコードでは”Record”と略されている。
• Table : Record BatchのArrayがChunked Arrayになったもの。Chunked Arrayはそれぞれ異なる型にできるがすべて同じ長さでなければならない。

なお、「Apache Arrow」について詳しく知りたい方は他の解説サイトを御覧ください。

Parquetファイルを読み込む

このライブラリは解説サイトやサンプルがほとんどないため、今回は簡単な使い方をまとめてみました。

go.modファイルを作ったあと、ライブラリを追加します。

go get github.com/apache/arrow/go/v11@latest

今回は以下のようなParquetファイルを用意しました。

test.parquet
{
    "id": [1, 2, 3, 4],
    "name": ["Apple", "Orange", "Grape", "Lemon"],
    "quantity": [50, 100, 200, 350],
}

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

import (
    "context"

    "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"
    "github.com/apache/arrow/go/v11/parquet/file"
    "github.com/apache/arrow/go/v11/parquet/pqarrow"
)

これで準備はできました。

ここからはParquetファイルを読み込みTableReaderを取得するところまでの処理となります。

filename := "test.parquet"

mem := memory.NewCheckedAllocator(memory.DefaultAllocator)

// ローカルファイルシステム上の指定されたParquetファイルのReaderを返します。
// ここでは読み込みを高速化するためにファイルをメモリマップしています。
reader, err := file.OpenParquetFile(filename, true, file.WithReadProps(parquet.NewReaderProperties(mem)))
if err != nil {
    panic(err)
}
defer reader.Close()

// ParquetファイルReaderオブジェクトからArrowオブジェクトへ変換して読むための
// Readerを作成します
arrowRdr, err := pqarrow.NewFileReader(reader, pqarrow.ArrowReadProperties{}, mem)
if err != nil {
    panic(err)
}

// ReadTableはファイル全体をTable形式で読み込みます
tbl, err := arrowRdr.ReadTable(context.Background())
if err != nil {
    panic(err)
}
defer tbl.Release()

// NewTableReaderはテーブルを反復処理するための新しいTableReaderを返します。
tr := array.NewTableReader(tbl, tbl.NumRows())
defer tr.Release()

tr.Next()

以上の処理によってParquetファイルのデータを詳細な情報を得ることができるようになりました。Table形式でも情報は取得できますが、今回はより簡単に情報が取得できるRecordBatchに変換しています。

// RecordBatch形式のデータを取得します。
	rec := tr.Record()
	defer rec.Release()

// Recordを出力した内容
// record:
//  schema:
// fields: 3
//    - id: type=int64, nullable
//    metadata: ["PARQUET:field_id": "-1"]
//    - name: type=utf8, nullable
//    metadata: ["PARQUET:field_id": "-1"]
//    - quantity: type=int64, nullable
//    metadata: ["PARQUET:field_id": "-1"]
//  rows: 4
//  col[0][id]: [1 2 3 4]
//  col[1][name]: ["Apple" "Orange" "Grape" "Lemon"]
//  col[2][quantity]: [50 100 200 350]

// スキーマ情報の取得
	schema := tr.Schema()

// schemaを出力した内容
//  schema:
//    fields: 3
//      - id: type=int64, nullable
//      metadata: ["PARQUET:field_id": "-1"]
//      - name: type=utf8, nullable
//        metadata: ["PARQUET:field_id": "-1"]
//      - quantity: type=int64, nullable
//        metadata: ["PARQUET:field_id": "-1"]

// RecordBatch形式のデータを取得します。
	rec := tr.Record()
	defer rec.Release()


// 指定したカラムのデータ型を取得
	rec.Column(0).DataType()
// 結果
// int64


// 列数を取得
	rec.NumCols()
// 結果
// 3


// 行数を取得
	rec.NumRows()
// 結果
// 4


// 全カラムデータを取得
	rec.Columns()
// 結果
// [[1 2 3 4] ["Apple" "Orange" "Grape" "Lemon"] [50 100 200 350]]


// 指定したのカラムデータを取得
	rec.Column(2)
// 結果
// [50 100 200 350]


// カラムデータの各要素の取得
	name := rec.Column(1).(*array.String)
	name.Value(1)
// 結果
// Orange


// 指定したカラム名を取得
	rec.ColumnName(2)
// 結果
// quantity


// レコードの一部を取り出す
	rec.NewSlice(1, 2)
// 結果
// record:
//  schema:
//  fields: 3
//    - id: type=int64, nullable
//    metadata: ["PARQUET:field_id": "-1"]
//    - name: type=utf8, nullable
//    metadata: ["PARQUET:field_id": "-1"]
//    - quantity: type=int64, nullable
//    metadata: ["PARQUET:field_id": "-1"]
//  rows: 1
//  col[0][id]: [2]
//  col[1][name]: ["Orange"]
//  col[2][quantity]: [100]

以上、Apache Arrowを使ったParquetファイルの読み込み方法でした。

最後に

Apache Arrowのライブラリとはいえ、機能が未完成でドキュメントやサンプルが少ないことが原因でなかなか普及していない現状です。メモのような記事ではありますが、少しでもお役に立てればと思います。