BigQuery のページネーションを Go でいい感じに実装する方法
はじめに
アプリケーションの実装を進める中で、BigQuery からデータを取得する際に、1ページあたりのサイズを適切に制御する必要があるユースケースに遭遇しました。
しかし、BigQuery のページネーションに関する情報があまり見つからなかったため、備忘録としてこの記事をまとめることにしました。
また、BigQuery のページネーションに関する情報が少ないのはそもそも、BigQuery ではデフォルトで maxResults が設定されていたり、Go の SDKでは自動的にページ分割が行われるため、パフォーマンス上の問題が発生しにくく大きな課題として取り上げられることが少ないのではないかと考えています。
BigQuery のページネーション実装例
前述した通り BigQuery は基本的によしなにページネーションを行ってくれるとのことですが、ページあたりのサイズをハンドリングしたい場合まず SQL で頑張る方法が挙げられ、具体的にはオフセット方式やカーソル方式が挙げられます。
ただ、これだと BigQuery というより SQL 側で解決することになるためスキーマ設計に影響を与える可能性があったり、そもそも実装が少々面倒だったりするのでもう少し楽な実現方法を後述します。
いい感じの実装方法
結論から述べると google.golang.org/api/iterator に含まれる Pager という型を使用することでいい感じに実装することができます。
サンプルコードは以下になりますが、実装自体は基本的に iterator が用意してくれているメソッドを呼び出しているだけなのでメソッド名からやっていることはなんとなく読み取れると思います。
また、処理の流れ自体は前述したカーソル方式に似たような形になっています。
package main
import (
"context"
"fmt"
"cloud.google.com/go/bigquery"
"github.com/Reimei1213/lab/paging-bigquery/pager"
)
func main() {
const (
projectID = ""
tableName = ""
)
ctx := context.Background()
cli, err := bigquery.NewClient(ctx, projectID)
if err != nil {
panic(err)
}
defer cli.Close()
p := pager.NewPager(cli)
if err := p.InitIterator(ctx, tableName); err != nil {
panic(err)
}
pageSize := 2
for {
rows, err := p.GetDataPage(ctx, pageSize)
if err != nil {
panic(err)
}
for _, r := range rows {
fmt.Println(r)
}
if p.GetPageToken() == "" {
break
}
}
}
package pager
import (
"context"
"fmt"
"cloud.google.com/go/bigquery"
"google.golang.org/api/iterator"
)
type Pager struct {
cli *bigquery.Client
iterator *bigquery.RowIterator
pageToken string
}
func NewPager(cli *bigquery.Client) *Pager {
return &Pager{
cli: cli,
}
}
func (p *Pager) InitIterator(ctx context.Context, tableName string) error {
q := p.cli.Query(fmt.Sprintf("SELECT * FROM `%s` LIMIT 10", tableName))
it, err := q.Read(ctx)
if err != nil {
return err
}
p.iterator = it
p.pageToken = ""
return nil
}
func (p *Pager) GetDataPage(_ context.Context, pageSize int) ([][]bigquery.Value, error) {
var rows [][]bigquery.Value
pager := iterator.NewPager(p.iterator, pageSize, p.pageToken)
nextToken, err := pager.NextPage(&rows)
if err != nil {
return nil, err
}
p.pageToken = nextToken
return rows, nil
}
func (p *Pager) GetPageToken() string {
return p.pageToken
}
※ 別の実現方法として iterator.PageInfo() 内の Token, MaxSize を変更して実装することもできますが、諸々 Pager が吸収してくれているので使用するメリットは特にないのかなと思ってます。
あと、動作確認後本当にリクエスト・レスポンスが分割されているのか不安だったので以下の環境変数を設定してログ出力をしてみましたが、 pager.NextPage() を実行する度にリクエストが飛んでいそうでした。
$ export GODEBUG=http2debug=2
最後に
今回はページごとのサイズをハンドリングしつついい感じに BigQuery のページネーションを Go で実装する方法についてご紹介しました。
この記事で紹介した方法を使用しないといけないユースケースに当たることは少ないとは思いますが誰かの役に立てたら幸いです。
また、記事の内容に間違い等ありましたら 優しく 教えてください。
P.S. ちなみに今回ご紹介した方法は最初 bigquery-emulator を使って検証していたましたが、謎のバグが発生して動かなかったので issue を作成しました。
サンプルコードはこちら
Discussion