【Rails】大量のデータを DB から CSV へ書き出す時のメモリ消費を抑える

こんにちは。 iys5 です。
今回は Ruby on Rails で DB から大量のデータを CSV に出力する際にメモリを無駄に消費しないような実装について考えたいと思います。
なお、以下のバージョンを前提としています。

• Ruby 3.2.2
• Ruby on Rails 7.0

まずは普通に書いてみる

今回は users テーブルからユーザの id, name, created_at を取得し CSV に出力することを目指します。
普通に実装してみると以下のようになります。

CSV.open('/path/to/file.csv', 'w') do |csv|
  User.all.each do |user|
    csv << user.values_at(:id, :name, :created_at)
  end
end

ユーザ数が少なければこの実装で問題ありません。
しかし、ユーザ数が膨大だと User.all.each が実行されたタイミングでそれらすべてのデータがメモリに一時的に乗ることになり、メモリを大幅に消費してしまいます。

find_each を使用する

上記の問題に関しては ActiveRecord ではおなじみの find_each を使用します。

 CSV.open('/path/to/file.csv', 'w') do |csv|
-  User.all.each do |user|
+  User.all.find_each do |user|
     csv << user.values_at(:id, :name, :created_at)
   end
 end

これにより DB からは一定のレコード数（デフォルトでは1000）ずつ取得するため、メモリの消費は一定量に抑えられます。

CSV に出力されるタイミング

上記までの内容から、DB からサーバまでデータを持ってくるところまでは良さそうです。
では持ってきたデータをサーバから CSV へ出力するタイミングはどうなるでしょうか。

試しに以下の処理を実行し、binding.pry のタイミングで CSV ファイルがどうなっているか確認してみます。

CSV.open('/path/to/file.csv', 'w') do |csv|
  csv << [1, 'foo', '1970-01-01 09:00:00 +0900']
  binding.pry
end

binding.pry のタイミングでのCSV

処理終了後のCSV

1,foo,1970-01-01 09:00:00 +0900

binding.pry のタイミングでは CSV ファイルにはまだデータが出力されず、処理終了後は出力されていました。

つまり、csv << [...] が実行された後でも end でブロックを抜けるまでメモリに配列の情報が残っていることになります。
これでは find_each で DB からのデータ取得を分割したとしても、結局すべてのデータを一時的にメモリに保持してしまうことになります。
これを解決する方法はないでしょうか。

実は csv << [...] で CSV への出力が行われないわけではなく、一定量のデータがバッファに積まれた段階で CSV へ出力されるようです。

試しに以下の処理を実行します。

CSV.open('/path/to/file.csv', 'w') do |csv|
  10000.times do |i|
    csv << [i, "name_#{i}"]
    if i % 100 == 0
      puts i
      sleep 1
    end
  end
end

この処理では100行/秒ずつバッファに積んでおり、このときのコンソール画面と、CSV の状態を less +F でリアルタイムに確認したものが以下になります。

コンソール画面（左）と CSV （右）

バッファに積まれたデータが一定の量を超えたタイミングで CSV に出力されていることがわかります。

これにより CSV への出力に関してのメモリ効率については Ruby が上手くやってくれていることがわかりました。

CSV#<< の実装では指定したオプションに関する処理が行われたあと、最終的に IO クラスの IO#<< を使用していそうだったので、CSVに限らず他のファイル形式でも同じことが言えそうです。

CSV#flush を使用する

Ruby には CSV#flush というメソッドも用意されており、これを明示的に呼ぶことでも CSV に書き込むことができます。
先程の1行出力する処理の binding.pry の前に csv.flush を追加して実行してみます。

CSV.open('/path/to/file.csv', 'w') do |csv|
  csv << [1, 'foo', '1970-01-01 09:00:00 +0900']
  csv.flush
  binding.pry
end

binding.pry のタイミングでのCSV

1,foo,1970-01-01 09:00:00 +0900

処理終了後のCSV

1,foo,1970-01-01 09:00:00 +0900

両方のタイミングで CSV に追加されていることが確認できました。

CSV#<< と同様に、 CSV#flush は IO クラスのメソッド IO#flush を継承して使用しているため、他のファイル形式でも同じことが言えそうです。

まとめ

CSV への出力タイミングに関する話が主になりましたが、結論としては DB から取得する際に find_each を使用すれば良いということになりました。

なお、こちらの記事によるとメンテナンス性は落ちますが、pluck を使用した実装をしたり、直接 SQL から CSV を出力することで、さらに低負荷な処理にすることができそうです。