あけましておめでとうございます。今年の目標を「Qiita に一本くらい記事を書く」としたため、さっさと目標達成をしようと箱根駅伝をチラ見しつつ本文を書いております。
動機
自分が携わるプロジェクトでパフォーマンスを上げる必要性が出てきた。それは高速化かつ省メモリ化の両方を満たす必要があった。
元ネタ
タイトルに関して、参考にさせていただいた元ネタがあります。
Rails の ActiveRecord で膨大なレコード数を高速かつ省メモリで処理する
大変参考になりました。ありがとうございます。
ちょっと掘り下げ
紹介記事にあるように「膨大なレコードを高速かつ省メモリで処理したい」というのが本来の目的です。では、「なんで pluck_each を使ってないの？」と問われると、「パフォーマンスチューニングをするにあたり、少し欲が出てきたから」が回答になります。
pluck_each は find_each を踏襲しているので、「結果を一件ずつ返す」というものになります。これをもうちょっと効率よくしたくなってきました。なので「一件ずつの処理ではなく、もうちょっとまとめて処理をしたい」という機能を満たすために find_in_batches を踏襲した pluck_in_batches があるといいなぁ、と思った次第です。
無いなら作れ
で、上手く解決できないかなぁ、とぐぐっていたのですが、以前の勤め先の先輩の格言「無いなら作れ」という言葉を思い出し自作してみることにしました。
さっきぐぐったら既に作っている人がいたっぽいけど、気にしない。
で、できたのがこんな感じ。
# 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 ブランチの pluck と find_in_bathces が 4.2 系とは別物くさかったので、これにも対応できてからやるかどうか考えます（多分、やらない）。