Rails の task で超巨大なテーブルに対して集計を行いたくて、
いろいろ調べてみたところ、以下の 2 つの方法をみつけました。
- pluck
- find_each
pluck は任意のカラムのデータを配列で取得するメソッドです。
find メソッドのようにインスタンスを取得することはできませんが、
その代わりにインスタンス生成にかかるオーバーヘッドがかからないため非常に高速です。
メモリ消費の点でも、インスタンスを生成しない分有利ですが、
取得するレコード数に比例してメモリ使用量は大きくなるため、
巨大なテーブルに対して使用するのは危険が伴います。
find_each は少しずつデータをメモリに展開して処理を行うメソッドです。
一度にメモリに展開するレコード数を制限できるため、
使用メモリ量をコントロールすることができます。
ただし、インスタンス生成のオーバーヘッドがかかるため、
データ量が膨大だと pluck に比べて速度が出ません。
そこで、これらの考えを組み合わせることで
高速かつ省メモリで大量レコードを処理するメソッドを書いてみました。
それが以下の pluck_each です。
# coding: utf-8
module ActiveRecord
module Calculations
def pluck_each(columns = [], batch_size = 100000)
# primary key の位置特定 (なければ columns の最後に追加)
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)
part = relation
last_pk_val = nil
loop do
part = relation.where('? > ?', self.primary_key, last_pk_val) unless last_pk_val.nil?
# 全レコードを yield したら終了
c = part.count
break if c == 0
# LIMIT を残り件数よりも多く設定すると遅くなる問題への対処
part = part.limit([batch_size, c].min)
# 部分的に pluck を実行
part.pluck(*columns).each do |ary|
yield ary
last_pk_val = ary[pk_idx]
end
end
end
end
end
使い方
Rails で使用する場合は lib/core_ext/active_record/calculations.rb に上記のコードを記述し、
config/initializers/extensions_loader.rb 等に以下のようなコードを書いておけば読み込めます。
Dir.glob("#{Rails.root}/lib/core_ext/**/*.rb").each {|ext| require ext }
pluck_each を使用する際は、取得したいカラム名 (シンボル) の配列を渡します。第 2 引数にはメモリに展開するレコード数も渡せます (デフォルトは 10 万)。
以下のような感じです。おそらく pluck と find_each を使用したことがあれば、(汎用性には乏しいですが) それなりに違和感なく使えると思います。
# 1000 レコードずつメモリに展開
User.where(last_login: s..e).pluck_each([:id, :username, :email, :last_login], 1000) do |ary|
# レコード 1 件毎の処理内容
# (ary = [1, 'user1', 'user1@mail.com', '2015-04-01 13:00:00'], ...)
end
注意点
- Primary Key が複合キーのテーブルには使用できません
- 第 1 引数に Primary Key が入っていない場合、ブロックパラメータの配列の最後に Primary Key の値が追加されます
最後に
MySQL だけなのかはよくわからないのですが、LIMIT を残り件数よりも多く設定するとかなり遅くなるため、
pluck_each にはその対策として毎ループで count も取得してます。
無駄な気もしますが、これしか思いつかなかったので、もっと良い方法があれば教えてください。
(それ以前に、LIMIT が残り件数よりも多いときになぜ遅くなるかもよくわかってないので、そっちも教えて欲しいです。)
ちなみに、find_each にはこの対策がされていなくて、最終ループだけ遅くなってしまうという問題があり、
それを pluck_each と似たような方法で対策すると find_each も速くなったりします。