BigQueryのLanguage Serverであるbqlsを作っている話

AI・機械学習チームの北川(@kitagry)です。このブログはエムスリーアドベントカレンダー11日目の記事です。

このブログではタイトル通りOSSで作成しているbqlsというLanguage Serverについて紹介します。

github.com

bqlsについて
Language Serverとは
何故bqlsを作ったのか
どうやって補完をするのか
Language Serverならではのパースの難しさ
なんとかパースする
修正方法を考える
まとめ
We are hiring

bqlsについて

bqlsはGoogle CloudのデータウェアハウスであるBigQueryのLanguage Serverです。主な特徴としては次のとおりです。

Language Serverなので主要なエディタ(VS Code, Vim, Emacsなど)で利用できる。
BigQueryの編集支援(補完・解析・SQLの実行など)

Language Serverとは

Language Serverとはコードの補完や解析といった機能を提供するサーバーです。 goplsやrust-analyzerといった各Language ServerはLanguage Server Protocol（以降LSP）というプロトコルに沿った実装をしています。このプロトコルに従うことによって、エディタやIDEなどは同じ実装で補完や定義ジャンプ、リファクタリングなどを様々な言語で利用できます。 LSP以前では各エディタが各言語のサポートを実装していたため、エディタごとに言語サポートを実装する必要がありました。 LSPを使うことで、エディタやIDEはLanguage Serverとの通信に集中し、Language Serverは言語固有の機能を提供することに集中できます。

つまり、今回のBigQuery用のLanguage ServerはLSPに従うことで補完などの機能の恩恵を得ることができるのです。

何故bqlsを作ったのか

BigQueryはWeb UIを提供しており、その中で補完やフォーマットの機能を提供しています。しかし、この補完が個人的にはとても遅く感じており普段書くときにはほとんど利用出来ない状態でした。なにより僕は普段Neovimを利用しているのですが、これらのショートカットが利用出来ないことに毎回ストレスを感じていました。そこで、ローカルのエディタ(Neovim)でBigQueryのSQLを快適に書く方法を探すことにしました。

BigQueryはSQLによって実行出来ます。そして、もちろんですがSQLのLanguage Serverも複数存在します。僕が知っているLanguage Serverとしてはsql-language-serverやsqlsがあります。 sqlsはBigQueryをサポートしていませんが、sql-language-serverはadapterとしてBigQueryのサポートをしています。しかし、実際に僕が使ってみたところこのサポートは不十分に感じていました。

個人的な感覚ですが、各SQLの実行環境には方言がありBigQueryは他のRDBMSなどと比較して特にこの方言が強く感じています。そのため、BigQueryで快適にSQLを書くためにはBigQuery専用のLanguage Serverでないといけないのではないかと漠然と感じていました。そんなものは当然あるはずもなく、BigQuery用のSQLを書くときにはWeb UI上で書いたり脳内補完をしたりして書いていました。

そんなモヤモヤを感じているときに、goccyさんのBigQueryエミュレータの作り方というスライドに出会いました。これはGoogleのGitHub Orgで提供されているOSSであるzetasqlを利用しているものでした。さらに驚くべきことにzetasqlがC++で書かれているのをGoでラップしたgo-zetasqlを公開してくれていました。僕はGoでLanguage Serverを作った経験があった*1ので、BigQueryのLanguage Serverも作っちゃうかーというノリでbqlsを作成することにしました。もし、go-zetasqlが無ければ作ろうと思いもしなかったのでgoccyさんにはすごく感謝しています。

どうやって補完をするのか

補完をするためにはSQLのパースを行う必要があります。パースには前述した通り、go-zetasqlを利用してパースをします。

catalog := types.NewSimpleCatalog(catalogName)

langOpt := zetasql.NewLanguageOptions()
opts := zetasql.NewAnalyzerOptions()
opts.SetLanguage(langOpt)
output, _ := zetasql.AnalyzeStatement("SELECT created_at FROM `bigquery-public-data.samples.github_nested`", catalog, opts)
fmt.Println(output.Statement().DebugString())

結果は次のような形になります。

QueryStmt
+-output_column_list=
| +-bigquery-public-data.samples.github_nested.created_at#3 AS created_at [STRING]
+-query=
  +-ProjectScan
    +-parse_location=0-67
    +-column_list=[bigquery-public-data.samples.github_nested.created_at#3]
    +-input_scan=
      +-TableScan(parse_location=23-67, column_list=bigquery-public-data.samples.github_nested.[repository#1, actor_attributes#2, created_at#3, public#4, actor#5, payload#6, url#7, type#8], table=bigquery-public-data.samples.github_nested, column_index_list=[0, 1, 2, 3, 4, 5, 6, 7])

簡単に説明をすると、QueryStmtが１つのSQLクエリ文を表しています。そして、QueryStmtのqueryの要素であるProjectScanがSELECTから始まるクエリの内容について表しています。その要素であるinput_scanがFROM以降の内容を表しており、今回はbigquery-public-data.samples.github_nestedというテーブルを指しています。各テーブルが何のカラムを持っているかという情報をanalyzerが識別できるようにする必要があるのですが、どのように登録するかについては本ブログでは省略します。

とにかく、go-zetasqlを利用することによって簡単にパースを行うことが出来ました。この情報を利用して、SELECTの後ろで補完するときにはテーブル情報からカラムデータを引っ張り出して補完をすればよいのです。

Language Serverならではのパースの難しさ

前節ではパースさえ出来れば補完が動くというお話をしました。まあ、実際にパースが出来ても補完を出すのは簡単では無いのですが、簡単ということにしてこのまま進めます。ここではLanguage Serverならではのパースの難しさについて考えてみましょう。

みなさんが補完を欲しい場面とはどのような状況でしょうか？次のようなクエリを書いたときのパースについて考えてみましょう（|はカーソルの位置を表しています）。

SELECT | FROM `project.dataset.table`

このパターンは頻出します。 project.dataset.tableにどのようなカラムがあるかが分からないときに補完によって選択したい場面です。さて、このクエリをzetasqlでパースするとどうなるでしょうか？

INVALID_ARGUMENT: Syntax error: SELECT list must not be empty

SELECTの中のリストが空ではいけないという趣旨のエラーが返ってきます。もし、BigQueryのクエリの結果を返すサーバーを作るだけであればSQLがパースに失敗した場合はそのままユーザーにエラーを返すだけで良いでしょう。しかし、編集する上ではSQLが文法的に正しくない状況がほとんどなのです。そんな中でも補完などの機能を使うためになんとかパースした情報を取得しないといけないのがLanguage Serverならではのパースの難しさです。

なんとかパースする

このようなパースの難しさはLanguage Server特有のものでしょう。そのため、各Language Serverがどのようにこれらの問題を解決しているかを調べてみることにしました。

まず初めに前述したsqlsがどのようにパースしているのかを調べてみます。 sqlsについては、作者のlighttiger2505さんがブログにどのような方針で実装したかを書いてくださっていたのでそれを参考にさせていただきました。 sqlsでは構文をグルーピング処理し、途中でエラーになってもスルーするようなパーサーを自作しているようです（まちがっていたらすいません）。他にもRustのLanguage Serverであるrust-analyzerも同様にRust本体とは異なる独自のパーサーを実装しているようです*2。 Rustが公式ライブラリとして提供しているパーサーについては定義取得のときなどに利用しているようです。公式パーサーと違い独自パーサーではエラーがあっても処理を止めること無く結果を返すような作りになっているようです。

sqlsやrust-analyzerのような独自パーサーを実装する以外には方法は無いのでしょうか？ GoのLanguage Serverであるgoplsは独自パーサーを実装せずに公式ライブラリであるgo/parserを用いて実装されています。 goplsではなんと、パースができるまで機械的に修正を試みるという方法が取られています。 bqlsではgoplsと同様にパースが成功するまでファイルを修正するという方法を採用しました。この方法のメリットはなんといってもパーサーを実装するコストを払うことなく実装ができることです。デメリットとしてはファイルの修正が不完全だとパースが出来ないことです*3。

修正方法を考える

先程の例である次のSQL文について考えてみます。

SELECT | FROM `project.dataset.table`  -- INVALID_ARGUMENT: Syntax error: SELECT list must not be empty

この例ではSELECTとFROMの間に何も無いことが問題になっています。そのため、この例では適当に*などをいれることによって文法エラーは無くなりそうです。

SELECT * FROM `project.dataset.table`  -- ok

では続いて次のSQL文について考えてみましょう。

SELECT i FROM `project.dataset.table`  -- Unrecognized name: i

このようにテーブルに存在しない識別子を書くことは編集する上では多くあります。例えば、テーブルにidentityという名前のカラムがあった場合にはiから順番に入力していきます。ユーザーがidentityを入力しきるまでこのエラーが出続けますが、正直補完でぱっと選択したくなるところです。このエラーについても*を入れるだけで良いでしょうか？

SELECT * FROM `project.dataset.table`  -- ok

実は次のようなケースでは*だとエラーが解消出来ません。

SELECT
    id,
    i|  -- Unrecognized name: i
    name
FROM
    `project.dataset.table`

これを*で置き換えると次のようなエラーになってしまいます。

SELECT
    id,
    *
    name -- Syntax error: Expected end of input but got identifier "name"
FROM
    `project.dataset.table`

このようにSQLを機械的に直すのはなかなか考えることが多く正直大変でした。これならパーサーを自作したほうが良かったのではと自問自答をよくしています。

ちなみに上の例についてbqlsでは1というINT64リテラルで置き換えることによってパースをしています。

SELECT
    id,
    1
    name --ok
FROM
    `project.dataset.table`

1 nameのように書くと1 AS nameと同等の文法になるため、パースは可能です。これでもいくつか対応できないパターンがあるのでより良い方法を模索中です。

まとめ

本ブログではBigQuery用のLanguage Serverを実装した話について書きました。 bqlsは自分が満足する形になればいいかなと思って、コソコソ作っていたのですが最近Issueがちらちら届くようになり、それなら公開までしちゃうかと思ってブログ化するに至りました。

正直、zetasqlというC++ライブラリをcgoで利用するようになった結果クロスコンパイルが出来なくなり、それと格闘した話などもあるのですがこの話はまた後日どこかで書けたら良いなと思います。自分としては得意ではない言語を触るときや編集が難しいときにLanguage Serverを作ることで自分の（そして他の人の！）開発体験を向上させることができるという経験が出来たのは良かったです。