Ruby
ActiveRecord
Rails4.2

pluck + find_in_batches = pluck_in_batches

あけましておめでとうございます。今年の目標を「Qiita に一本くらい記事を書く」としたため、さっさと目標達成をしようと箱根駅伝をチラ見しつつ本文を書いております。

動機

自分が携わるプロジェクトでパフォーマンスを上げる必要性が出てきた。それは高速化かつ省メモリ化の両方を満たす必要があった。

元ネタ

タイトルに関して、参考にさせていただいた元ネタがあります。

Rails の ActiveRecord で膨大なレコード数を高速かつ省メモリで処理する

大変参考になりました。ありがとうございます。

ちょっと掘り下げ

紹介記事にあるように「膨大なレコードを高速かつ省メモリで処理したい」というのが本来の目的です。では、「なんで pluck_each を使ってないの?」と問われると、「パフォーマンスチューニングをするにあたり、少し欲が出てきたから」が回答になります。
pluck_each は find_each を踏襲しているので、「結果を一件ずつ返す」というものになります。これをもうちょっと効率よくしたくなってきました。なので「一件ずつの処理ではなく、もうちょっとまとめて処理をしたい」という機能を満たすために find_in_batches を踏襲した pluck_in_batches があるといいなぁ、と思った次第です。

無いなら作れ

で、上手く解決できないかなぁ、とぐぐっていたのですが、以前の勤め先の先輩の格言「無いなら作れ」という言葉を思い出し自作してみることにしました。
さっきぐぐったら既に作っている人がいたっぽいけど、気にしない。
で、できたのがこんな感じ。

calculations.rb
# coding: utf-8
module ActiveRecord
  module Calculations

    def pluck_in_batches(columns = [], option = {})
      return self if columns.blank?

      relation   = self
      start      = option[:start]
      batch_size = option[:batch_size] || 1000

      # ブロックで使われることが大前提。
      unless block_given?
        raise "Do not allow without block."
      end

      # primary key の位置特定(なければ columns の最後に追加)
      columns_count = columns.count
      pk_idx = columns.map(&:to_s).index(self.primary_key) || columns_count
      columns << self.primary_key if pk_idx == columns_count

      relation = self.reorder(self.primary_key).limit(batch_size)
      if relation.present?
        records = relation.where(table[primary_key].gteq(relation[0].id)).pluck(*columns)
      else
        return self
      end

      while records.any?
        records_size = records.size
        last_record = records.last
        primary_key_offset = (last_record.is_a?(Array) : last_record[pk_idx] ? last_record)
        raise "Primary key not included in the custom select clause" unless primary_key_offset

        yield records

        break if records_size < batch_size

        records = relation.where(table[primary_key].gt(primary_key_offset)).pluck(*columns)
      end
    end

  end # Calculations
end # ActiveRecord

ちなみに私の環境は Rails 4.2 系ですので、Active Record もそれに準拠しています。よって、pluck と find_in_batches も 4.2 系を踏襲しています。
以下は参考までにオリジナルの pluck と find_in_batches。

4-2-stable ブランチの pluck
4-2-stable ブランチの find_in_batches

制限事項

pluck_in_batches は以下の制限があります。

  • 複合キーには未対応
  • なんだかんだでプライマリキー名が id に固定されている
  • pluck_in_batches を使用するブロックが、プライマリキーでレコードが一意に絞られない場合、batch_size の設定値次第で取りこぼしが発生する

最後に

もうちょっと融通が効くものにしようかと思いましたが、それは in_batches ではない別物になりそうな気がしたので作ってません。また必要にせまられたら作るかもしれません。既に master ブランチの pluckfind_in_bathces が 4.2 系とは別物くさかったので、これにも対応できてからやるかどうか考えます(多分、やらない)。