Enumerator::Lazyで遅延評価する #Ruby

Enumerator::Lazy を使うとメモリ効率の良いプログラムが書ける。
けれどあまり実際のコードで見かけないので、実際に使えそうな例を上げたい。

動作確認環境

Ruby 2.4.1
activerecord 5.1.4
sqlite3 1.3.13

まずはよく見るEnumerator::Lazyの例

(1..Float::INFINITY).lazy.select { |n| n % 3 == 0 }.first(5)  # => [3, 6, 9, 12, 15]

無限リストから3の倍数の最初の５つを取得するコード。
これだけだと実用性は感じないけれど、 next を使えば何かのgeneratorとしては使えそう。

def generator
  (1..Float::INFINITY).lazy.select { |n| n % 3 == 0 }
end

gen = generator()
5.times do
  p gen.next
end

ActiveRecordのfind_eachと組み合わせる例

find_each はブロックを渡すと each のように1件ずつループ処理するが、裏では batch_size 件ずつDBからデータを取得してくれる。

（↓の例のDBにデータを用意する部分は http://qiita.com/xkumiyu/items/2ecee242b7e6e6c6d9a1 のコードを使わせて頂きました。）

Gemfile

source 'https://rubygems.org'

gem 'activerecord'
gem 'sqlite3'

require "active_record"

# データベースへの接続
ActiveRecord::Base.establish_connection(
  adapter:   'sqlite3',
  database:  ':memory:'
)

# スキーマの設定
class InitialSchema < ActiveRecord::Migration[5.1]
  def self.up
    create_table :users do |t|
      t.string :name
    end
  end

  def self.down
    drop_table :users
  end
end

# マイグレーション
InitialSchema.migrate(:up)

class User < ActiveRecord::Base
end

# モデルを生成
20.times do |i|
  User.create(name: "user-#{i}")
end

ActiveRecord::Base.logger = Logger.new(STDOUT)

# ↓ここから本題

User.find_each(batch_size: 10) do |user|
  p user.name
end

実行すると↓のように10件ごとにDBからデータ取得していることがわかる。

また、

User.find_each(batch_size: 10) do |user|
  p user.name
end

の部分を

def users
  User.find_each(batch_size: 10)
end

users.each do |user|
  p user.name
end

にしても結果は同じになる。

しかし、

def users
  User.find_each(batch_size: 10)
end

users.map(&:name).each do |name|
  p name
end

にすると↓のように（クエリは10件ずつ取得するものだが）最初にすべて取得してしまう。

これは map(&:name) の部分で User インスタンスをその名前に変換する部分では10件ごとのデータ取得を行いなから変換されているが、すべての変換処理が終わってから .each do |name| でループに入るため。
これではせっかくの find_each が台無しである。

そこで lazy を使う。

def users
  User.find_each(batch_size: 10).lazy
end

最後に .lazy を付けただけだが、結果は↓のように最初と同じようになる。

これで users メソッドを利用する側では特に何も気にせずに　map や select 等の操作が行える。

例えば、

users.map(&:name).each_slice(2) do |names|
  p names
end

とすれば、データ取得は10件ごと、処理は2件ごと、という挙動になる。

自前でEnumerator::Lazyを生成する例

.lazy がいつでも使えれば良いが、そのままでは使えない場合もある。
↓に find_in_batches を使って find_each(batch_size: 10).lazy を再実装した例を挙げる。

class Source
  def each
    User.find_in_batches(batch_size: 10) do |users|
      yield users
    end
  end
end

def users
  src = Source.new

  Enumerator::Lazy.new(src) do |yielder, users|
    users.each { |user| yielder << user }
  end
end

users.map(&:name).each do |name|
  p name
end

実行結果は↓のように find_each(batch_size: 10).lazy の場合と同じである。

ここで、 Source#each の中の User.find_in_batches を、pagingでデータ取得する外部APIだと思ってほしい。
すると、外部データをlazyにmapやeachできるインターフェースが実装できたことになる。

また、ファイルやネットワークからのデータ読み込みで .lazy が使えない場合も上記の例を少し変更すれば対応できると思う。
（もちろん .lazy が使えるなら使った方が良い。）

まとめ

このように Enumerator::Lazy を使用すると「必要になったら取得する」という動作が簡単に実装できる。