find_each、find_in_batchesとprimary

RailsのActiveRecordのメソッドであるfind_each、find_in_batchesとprimary_keyの備忘録。両メソッドの詳しい説明に関しては参考URLを参照してくださいー。

まずはコードリーディングから。find_each、find_in_batchesともにin_batchesを呼び出してクエリ発行＆ブロック実行をしています。

def in_batches(of: 1000, start: nil, finish: nil, load: false, error_on_ignore: nil)
  relation = self
  unless block_given?
    return BatchEnumerator.new(of: of, start: start, finish: finish, relation: self)
  end

  if arel.orders.present? || arel.taken.present?
    act_on_order_or_limit_ignored(error_on_ignore)
  end

  relation = relation.reorder(batch_order).limit(of)
  relation = apply_limits(relation, start, finish)
  batch_relation = relation

  loop do
    if load
      records = batch_relation.records
      ids = records.map(&:id)
      yielded_relation = self.where(primary_key => ids)
      yielded_relation.load_records(records)
    else
      ids = batch_relation.pluck(primary_key)
      yielded_relation = self.where(primary_key => ids)
    end

    break if ids.empty?

    primary_key_offset = ids.last
    raise ArgumentError.new("Primary key not included in the custom select clause") unless primary_key_offset

    yield yielded_relation

    break if ids.length < of
    batch_relation = relation.where(arel_attribute(primary_key).gt(primary_key_offset))
  end
end

挙動としては以下の３つのコードに着目すればOKです

relation = relation.reorder(batch_order).limit(of)
primary_key_offset = ids.last
batch_relation = relation.where(arel_attribute(primary_key).gt(primary_key_offset))

※yielded_relationはload_records(records)でrecordsを設定しているのでto_aでレコード抽出時にSQLは発行されない。

ざっくり説明すると

relationでlimit（処理するチャンクのレコード数）を指定してクエリを発行
次のチャンクは前のチャンクの最後のprimary_keyの値より大きいprimary_keyのレコードを取得

という流れになっています。

例えばbatch_sizeパラメータを10にすると以下のようなSQLが逐次発行され、SQLの結果をActiveRecordのArrayに格納してブロック内の処理が実行されます。

Hoge Load (1.1ms)  SELECT  `hoges`.* FROM `hoges` ORDER BY `hoges`.`int` ASC LIMIT 10
Hoge Load (1.0ms)  SELECT  `hoges`.* FROM `hoges` WHERE (`hoges`.`int` > 10) ORDER BY `hoges`.`int` ASC LIMIT 10
Hoge Load (1.0ms)  SELECT  `hoges`.* FROM `hoges` WHERE (`hoges`.`int` > 20) ORDER BY `hoges`.`int` ASC LIMIT 10
...

primary_keyをモデルで設定していないケース（例えばMySQLのパーティショニングで複合PRIMARY KEYを使わざるを得ないなど）でfind_each、find_in_batchesを使いたい場合は、self.primary_key=によって主キーを明示的に設定する必要があります。ただし、文字通りprimary_keyなのでユニークでないと、primary_key_offsetの部分でうまくレコードが抽出できません。