More than 1 year has passed since last update.

ActiveRecord で order しつつ find_each する

Last updated at 2022-02-27Posted at 2022-02-27

動機

大量のレコードを読み込んで処理をする場合、each を使ってイテレートすると、レコードを一括で取得するためメモリを大量に消費してしまう場合があります。そこで ActiveRecord::Batches#find_each を使ってバッチサイズ (デフォルトは 1,000 件) ごとにレコードを分割して読み込み、イテレートするのが常套手段ではないでしょうか。しかし find_each を使う場合、その仕組み上、リレーションの order 情報を無視して id の昇順¹に並べてイテレートすることになります。そこで、今回独自に order を保ったまま find_each する方法を用意しました。

方法

最初は ActiveRecord::Relation をオープンクラスやモジュールの prepend で拡張しようと考えましたが、方法ではそのサブクラス (Pokemon モデルに対する Pokemon::ActiveRecord_Relation など) にメソッドを定義することができませんでした (なぜだろう 🤔) 。そこで ActiveRecord::QueryMethods#extending を使用し、 ActiveRecord::Relation オブジェクトを必要に応じて拡張する方法を採用しました。

lib/find_each_in_order.rb

# ActiveRecord::Relation#find_each_in_order を提供するためのモジュール。
# ActiveRecord::QueryMethods#extending を使ってリレーションにメソッドを追加する。
module FindEachInOrder
  class DoNotUseLimitError < StandardError; end

  # ソート順を保ちつつ find_each のようにバッチサイズずつ読み込み処理を行う。
  # 一貫性を保証するために DB 読み取りのみの場合でもトランザクションで囲むことを推奨する。
  # @param [Integer] batch_size バッチサイズ
  # @return [Enumerator]
  def find_each_in_order(batch_size: 1_000, &block)
    # batch_size は 1 より大きくないと find_each にする意味がない。
    raise(ArgumentError.new('batch_size must be greater than 1')) unless batch_size > 1

    # バッチの取得に LIMIT を使う関係で、
    # すでに LIMIT を使用しているリレーションではこのメソッドは呼べないようにする。
    # find_each は LIMIT も考慮しているが、このメソッドはなるべく実装をシンプルにしたい。
    raise(DoNotUseLimitError.new('Do not use limit in relation')) if limit_value

    # find_each と同様、ブロックを渡さない場合は Enumerator オブジェクトを返すようにする。
    unless block
      return enum_for(__method__, batch_size: batch_size) { size }
    end

    page = 1

    loop do
      offset = (page - 1) * batch_size
      batch = limit(batch_size).offset(offset)
      batch.each(&block)

      # 現在の件数がバッチサイズより少なければ、もう次のバッチは存在しない。
      break if batch.size < batch_size

      page += 1
    end
  end
end

そして以下のように使用します。

ActiveRecord::Base.transaction do
  pokemons = Pokemon.order(:number).extending(FindEachInOrder)
  pokemons.find_each_in_order(batch_size: 300) do |pokemon|
    puts(pokemon.name)
  end
end
# SELECT `pokemons`.* FROM `pokemons` ORDER BY `pokemons`.`number` ASC LIMIT 300 OFFSET 0
# SELECT `pokemons`.* FROM `pokemons` ORDER BY `pokemons`.`number` ASC LIMIT 300 OFFSET 300
# SELECT `pokemons`.* FROM `pokemons` ORDER BY `pokemons`.`number` ASC LIMIT 300 OFFSET 600

注意点

※ 以下は MySQL 5.7 について調べた話です。他の RDBMS については検証の上使用してください。

OFFSET について

OFFSET を使って件数を絞り込む場合、index が効かず全件走査になる場合があるので注意が必要です。

トランザクションについて

この方法ではバッチごとに SELECT 文を発行しますが、その結果の一貫性を保証するため、取得 (SELECT) のみでもトランザクションで囲むことを推奨します。

MySQL の InnoDB のトランザクション分離レベルはデフォルトで REPEATABLE READ です。そして MySQL の REPEATABLE READ ではトランザクション内で発行する SELECT は一貫性読み取りです。

MySQL :: MySQL 5.7 Reference Manual :: 14.7.2.1 Transaction Isolation Levels より

REPEATABLE READ
This is the default isolation level for InnoDB. Consistent reads within the same transaction read the snapshot established by the first read. This means that if you issue several plain (nonlocking) SELECT statements within the same transaction, these SELECT statements are consistent also with respect to each other. See Section 14.7.2.3, “Consistent Nonlocking Reads”.

これは「トランザクションの途中で別のトランザクションがレコードを追加したり削除したりしても、実行中のトランザクションでの SELECT には影響しない。そのため id 順に並べるのではなく LIMIT と OFFSET を使用する方法でも一貫性を保証できる」という想定のもと実装しました。

参考

Rails 6.1 から find_each, find_in_batches, in_batches で order: :desc というオプションを渡すことで、id の降順でもソートできるようになりました。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up