Batching について、 GraphQL の文脈で

この記事は何か

Rails で GraphQL をやって得た知見から、 graphql-batch という gem と Batching という仕組みについて書いてみます。

記事中に出てくるコードは https://github.com/tanaka51/sample-rails-graphql-batch こちらのリポジトリにコミットしているので、実際に動かしてみたい方は参考にしてみてください。

書いていないこと

• GraphQL について
• N+1 について

アドベントカレンダー

この記事は、CBcloud アドベントカレンダー( https://qiita.com/advent-calendar/2020/cbcloud )の18日目の記事として投稿しています。

Batching で GraphQL の N+1 を解決する

GraphQL で素朴に実装して N+1 を確認する

ここからは、下記に示すような、ブログシステム的なものを考えます。

ユーザー(User)がいて、ユーザーは記事(Article)をかけます。また、お気に入りの記事情報を保存(Favorite)できる、というシステムです。

さて、 GraphQL で以下のようなクエリを叩く事を考えましょう。

{
  users {
    name
    favorites {
      article {
        title
      }
    }
  }
}

これは、ユーザー全員のお気に入りの記事のタイトルをまとめて取得する、というクエリになります。

このクエリに対応する実装を graphql-ruby を用いて書くと、こうなります。

module Types
  class QueryType < Types::BaseObject
    field :users, [UserType], null: false
    def users
      User.all
    end
  end

  class UserType < Types::BaseObject
    field :name, String, null: false
    field :articles, [Types::ArticleType], null: false
    field :favorites, [Types::FavoriteType], null: false
  end

  class FavoriteType < Types::BaseObject
    field :user, Types::UserType, null: false
    field :article, Types::ArticleType, null: false
  end

  class ArticleType < Types::BaseObject
    field :title, String, null: false
    field :author, Types::UserType, null: false
  end
end

詳しいことは省きますが、 graphql-ruby では必要な type を定義していき、それらを宣言的に配置していくことで query を定義できます。

そして、これをそのまま実行すると、以下のように N+1 が出ます。

Started POST "/graphql" for 127.0.0.1 at 2020-12-14 21:58:27 +0900
Processing by GraphqlController#execute as HTML
  Parameters: {"query"=>"{\n  users {\n    favorites {\n      article {\n        title\n      }\n    }\n  }\n}", "variables"=>nil, "graphql"=>{"query"=>"{\n  users {\n    favorites {\n      article {\n        title\n      }\n    }\n  }\n}", "variables"=>nil}}
Can't verify CSRF token authenticity.
   (0.1ms)  SELECT sqlite_version(*)
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  User Load (0.5ms)  SELECT "users".* FROM "users"
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  Favorite Load (0.8ms)  SELECT "favorites".* FROM "favorites" WHERE "favorites"."user_id" = ?  [["user_id", 1]]
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  Article Load (0.5ms)  SELECT "articles".* FROM "articles" WHERE "articles"."id" = ? LIMIT ?  [["id", 4], ["LIMIT", 1]]
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  Article Load (0.2ms)  SELECT "articles".* FROM "articles" WHERE "articles"."id" = ? LIMIT ?  [["id", 5], ["LIMIT", 1]]
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  Favorite Load (0.2ms)  SELECT "favorites".* FROM "favorites" WHERE "favorites"."user_id" = ?  [["user_id", 2]]
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  Article Load (0.2ms)  SELECT "articles".* FROM "articles" WHERE "articles"."id" = ? LIMIT ?  [["id", 1], ["LIMIT", 1]]
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  Article Load (0.1ms)  SELECT "articles".* FROM "articles" WHERE "articles"."id" = ? LIMIT ?  [["id", 2], ["LIMIT", 1]]
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
Completed 200 OK in 28ms (Views: 0.3ms | ActiveRecord: 2.8ms | Allocations: 9994)

通常の N+1 を回避する書き方

ここで比較のため、通常の Rails のコードを書いてみます。

def index
  users = User.all.includes(favorites: :article).map do |user|
    {
      name: user.name,
      favorites: user.favorites.map do |favorite|
        article: {
          title: favorite.article.title
        }
      end
    }
  end

  render :json, { data: users }
end

おそらくこんな形になるでしょうか。 典型的な N+1 を回避するためのコードです。User というリソースに対して「何が下にひっつくのか」を考え、それをコード中に記述して N+1 を回避します。

この考え方は GraphQL ではうまくいきません。なぜならユーザーはクエリを好きなように組み替える事ができるからです。例えば、ユーザーが書いた記事一覧も、同じエンドポイントで提供をしないといけません。

{
  users {
    name
    articles {
      title
    }
  }
}

つまり、事前に「何が下にひっつくのか」を定義できないのです。クエリのすべての組み合わせを列挙してそれぞれに関して定義しておく、という方法も考えられますが、現実的ではありませんよね。

これを解決する手段として Batching という仕組みを導入するのが graphql-batch という gem です。この gem を使いながら、 batching がどういう仕組みか、紹介してみたいと思います。

graphql-batch を用いて N+1 を回避する

まずは graphql-batch で、どのようなコードになるか、素朴に書いてみます。

最初に Loader を用意します。

# app/loaders/user_favorites_loader.rb
class UserFavoritesLoader < GraphQL::Batch::Loader
  def perform(user_ids)
    favorites = Favorite.where(user_id: user_ids)
    favorites.group_by(&:user_id).each do |user_id, favorites|
      fulfill(user_id, favorites)
    end
  end
end

# app/loaders/article_loader.rb
class ArticleLoader < GraphQL::Batch::Loader
  def perform(ids)
    Article.where(id: ids).each{|a| fulfill(a.id, a)}
  end
end

そして型定義の所で、 Loader を使うようにします。

module Types
  class UserType < Types::BaseObject
    field :id, ID, null: false
    field :name, String, null: false
    field :created_at, GraphQL::Types::ISO8601DateTime, null: false
    field :updated_at, GraphQL::Types::ISO8601DateTime, null: false
    field :articles, [Types::ArticleType], null: true
    field :favorites, [Types::FavoriteType], null: true

    def favorites
      UserFavoritesLoader.for.load(object.id)
    end
  end
end

module Types
  class FavoriteType < Types::BaseObject
    field :id, ID, null: false
    field :created_at, GraphQL::Types::ISO8601DateTime, null: false
    field :updated_at, GraphQL::Types::ISO8601DateTime, null: false
    field :user, Types::UserType, null: true
    field :article, Types::ArticleType, null: true

    def article
      ArticleLoader.for.load(object.article_id)
    end
  end
end

これで実行するとサーバーのログが以下のようになり、 N+1 が解決した事がわかります

Started POST "/graphql" for 127.0.0.1 at 2020-12-18 13:51:59 +0900
Processing by GraphqlController#execute as HTML
  Parameters: {"query"=>"{\n  users {\n    favorites {\n      article {\n        title\n      }\n    }\n  }\n}", "variables"=>nil, "graphql"=>{"query"=>"{\n  users {\n    favorites {\n      article {\n        title\n      }\n    }\n  }\n}", "variables"=>nil}}
Can't verify CSRF token authenticity.
  User Load (0.1ms)  SELECT "users".* FROM "users"
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  Favorite Load (0.2ms)  SELECT "favorites".* FROM "favorites" WHERE "favorites"."user_id" IN (?, ?)  [[nil, 1], [nil, 2]]
  ↳ app/loaders/favorites_loader.rb:4:in `group_by'
  Article Load (0.3ms)  SELECT "articles".* FROM "articles" WHERE "articles"."id" IN (?, ?, ?, ?)  [[nil, 4], [nil, 5], [nil, 1], [nil, 2]]
  ↳ app/loaders/article_loader.rb:3:in `perform'
Completed 200 OK in 10ms (Views: 0.4ms | ActiveRecord: 0.7ms | Allocations: 3675)

以上が Loader の一番素朴な使い方です。

graphql-batch の使い方

grapql-batch について、少し掘り下げて説明します。

まずは Loader を使った最小限のコードを見ていきましょう。

class ArticleLoader < GraphQL::Batch::Loader
  def perform(ids)
    Article.where(id: ids).each{|a| fulfill(a.id, a)}
  end
end

module Types
  class FavoriteType < Types::BaseObject
    field :article, Types::ArticleType, null: true

    def article
      ArticleLoader.for.load(object.article_id)
    end
  end
end

基本的な Loader の使い方としては、
1. for メソッドで初期化する
2. load メソッドにキーとなる値を渡す

となります。
そして Loader の実装は、
1. GraphQL::Batch::Loader を継承したクラスをつくる
2. perform メソッドを実装する

となります。

次に perform メソッドをどうやって実装するか説明します。

perform は、最終的にレスポンスを生成する際に呼ばれます。
そして実際に呼ばれる際には、 load で渡された値が配列で渡ってきます。
このタイミングで、リソースをまとめて取得し、 fulfill メソッドに詰め込みます。

fulfill には キー, 値 という順番で値を渡します。
ここでキーとは、 load に渡した値と同じ値を指定します。
こうすることで、 load は戻り値として値を取り出すことができます。

…と文字で説明しても分かりづらいので、絵も載せてみます。

これでうまくイメージできると嬉しいのですが、、、
要点を整理すると、

• load で値を集める
• 集めた値は perform で処理される
• perform 内でリソースをまとめて取得し、 fulfill で詰めていく
• fulfill で詰めた値を load が取り出して戻り値となる

となります。

上記の要点が、つまり Batching というアイディアだと私は思っています。
目の前の処理をすぐに実行するのではなく、一旦別の場所に貯めておいて、最後にまとめて処理します。

実装例

さて、 graphql-batch については、初期化のメソッドが .for になっている事から分かるように、実際には汎用的な作りを意識したものになっています。
以下は README からの抜粋です。

field :product, Types::Product, null: true do
  argument :id, ID, required: true
end

def product(id:)
  RecordLoader.for(Product).load(id)
end

また、 Loader の例は examples 配下にあるので、基本的にはこれ使えば良いと思います。
https://github.com/Shopify/graphql-batch/tree/master/examples

Loader の実装をいきなり見ても分かりづらいと思いますが、この記事が助けに読み進められると嬉しいなぁと思います。

リソースカウンティングにおける N+1 の解決例

examples に無い実装例として、リソースカウンティングの Loader を書いてみたいと思います。
例えばこういうクエリーを処理するのに使います。

{
  users {
    favorites_count
  }
}

これを素朴に実装するとこうなります

module Types
  class UserType < Types::BaseObject
    field :favorites_count, Integer, null: false

    def favorites_count
      object.favorites.count
    end
  end
end

Started POST "/graphql" for 127.0.0.1 at 2020-12-18 17:39:42 +0900
Processing by GraphqlController#execute as HTML
  Parameters: {"query"=>"{\n  users {\n    favoritesCount\n  }\n}", "variables"=>nil, "graphql"=>{"query"=>"{\n  users {\n    favoritesCount\n  }\n}", "variables"=>nil}}
Can't verify CSRF token authenticity.
  User Load (0.6ms)  SELECT "users".* FROM "users"
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
   (0.5ms)  SELECT COUNT(*) FROM "favorites" WHERE "favorites"."user_id" = ?  [["user_id", 1]]
  ↳ app/graphql/types/user_type.rb:16:in `favorites_count'
   (0.2ms)  SELECT COUNT(*) FROM "favorites" WHERE "favorites"."user_id" = ?  [["user_id", 2]]
  ↳ app/graphql/types/user_type.rb:16:in `favorites_count'
Completed 200 OK in 34ms (Views: 0.3ms | ActiveRecord: 2.0ms | Allocations: 10884)

こんな形で、ユーザー数分だけ COUNT 文が発行されます。
Rails をやっていた場合に counter_cache を検討するような場面でも、 Loader を使えばすっきり書く事ができます。

class UserFavoritesCountLoader < GraphQL::Batch::Loader
  def perform(user_ids)
    Favorite.where(user_id: user_ids).group(:user_id).count.each do |user_id, count|
      fulfill(user_id, count)
    end
  end
end

module Types
  class UserType < Types::BaseObject
    field :favorites_count, Integer, null: true

    def favorites_count
      UserFavoritesCountLoader.for.load(object.id)
    end
  end
end

上記で書いた UserFavoritesLoader ほぼそのままですが、内容としては、集めた user_id で group by + count で、対象の user_id をキーにしてカウントを保存しています。

Started POST "/graphql" for 127.0.0.1 at 2020-12-18 17:43:16 +0900
Processing by GraphqlController#execute as HTML
  Parameters: {"query"=>"{\n  users {\n    favoritesCount\n  }\n}", "variables"=>nil, "graphql"=>{"query"=>"{\n  users {\n    favoritesCount\n  }\n}", "variables"=>nil}}
Can't verify CSRF token authenticity.
   (0.1ms)  SELECT sqlite_version(*)
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
  User Load (0.1ms)  SELECT "users".* FROM "users"
  ↳ app/controllers/graphql_controller.rb:15:in `execute'
   (0.2ms)  SELECT COUNT(*) AS count_all, "favorites"."user_id" AS favorites_user_id FROM "favorites" WHERE "favorites"."user_id" IN (?, ?) GROUP BY "favorites"."user_id"  [[nil, 1], [nil, 2]]
  ↳ app/loaders/user_favorites_count_loader.rb:3:in `perform'
Completed 200 OK in 25ms (Views: 0.3ms | ActiveRecord: 1.8ms | Allocations: 9009)

これで N+1 を避ける事ができました。

まとめ

graphql-batch gem を用いて、素朴な Loader を実装して使い方を説明しました。
今回の記事では限定的な Loader を作成しましたが、汎用的な Loader を作れば、「リソースの取得にはとりあえず Loader 使っておけばよい」とすればよく、これだけで N+1 の事を考える必要がほぼなくなるので、より本質的な所に集中できるようになると思います。