CSVファイルをフィルタリングする - エンジニアのソフトウェア的愛情

「元のCSVファイルから特定の条件の行だけを抜き出した新しいCSVファイルを作成する」という仕事中の話題から。

Ruby

Rubyでしたら、不都合がない限り標準添付のcsvライブラリを利用すると思います。

onliner

「ファイルから読み込んで、行を選択して、ファイルに書き出す」という手続きをそのままコードにしたもの。

require 'csv'

File.write('dst.csv', CSV::Table.new(CSV.read('src.csv', headers: true).select {|row| row['max'].to_i >= 100 }))

手続きごとに分解すると。

require 'csv'

src_data = CSV.read('src.csv', headers: true)              # (1) ヘッダ情報付きでファイルから読み込む
dst_data = src_data.select {|row| row['max'].to_i >= 100 } # (2) 行を選択する
dst_csv = CSV::Table.new(dst_data)                         # (3) 行の集まりをCSVにする
File.write('dst.csv', dst_csv)                             # (4) ファイルに書き出す

このうち(4)の「行の集まりをCSVにする」部分は、CSV::Table.newが「CSV::Rowのインスタンスの配列」を引数として受け取ることを利用しています。

docs.ruby-lang.org

CSV::Table.new([CSV::Row.new(%w(a b c), [1, 2, 3]), CSV::Row.new(%w(a b c), [4, 5, 6])])
# => #<CSV::Table mode:col_or_row row_count:3>

またCSV::Tableのインスタンスは文字列に変換するとCSV形式の文字列になります。ここではIOへの出力されたときに自動的に文字列に変換されることを利用しています。

csv = CSV::Table.new([CSV::Row.new(%w(a b c), [1, 2, 3]), CSV::Row.new(%w(a b c), [4, 5, 6])])

csv.to_s
# => "a,b,c\n1,2,3\n4,5,6\n"

puts csv
# => a,b,c
#    1,2,3
#    4,5,6

Stream

入力となるCSVと出力となるCSVを同時に開き、読み込む1行ごとに選択と書き込みをする、というアイディアです。

書き込み時にもヘッダをつける場合には:headersオプションの他に:write_headersオプションを指定する必要があります。

また足をすくわれた点として。 CSV#headersは、ヘッダの利用を指定していても、読み込み前の状態ではtrueしか返してくれません。

[https://docs.ruby-lang.org/ja/latest/method/CSV/i/headers.html:docs.ruby-lang.org

ヘッダの配列を取得するために、ここでは最初に1回CSV#getsを呼び出しています。

require 'csv'

CSV.open('src.csv', headers: true) do |src|
  CSV.open('dst.csv', 'w', headers: true, write_headers: true) do |dst|
    # 1行目を読み込む
    row = src.gets

    # ヘッダ行を書き込む
    dst.puts src.headers

    while row # nil だったら、ファイルの終端に達した
    # 条件に合う場合、その行を書き込む
    dst.puts row if row['max'].to_i >= 100

    # 次の行を読み込む
    row = src.gets
    end
  end
end

SQLite3

あるデータから条件にあったデータを取得するというのならば。そもそもSQLなどを使うのがよいのでは？　と思い浮かんだので書いてみました。

CSV SQLで検索すると色々出てきますが、今回はSQLite3を使うことにしました。

import

www.sqlite.org

まず、SQLite3 のコンソールを起動して動きを確認します。

$ sqlite3 temp.db
sqlite>

.modeコマンドでモードをCSVに変更し、.importコマンドでCSVファイルをテーブルにインポートします。

sqlite> .mode csv
sqlite> .import src.csv data

インポートされたデータはSQLite3のデータですので、当然SQLで検索することができます。

sqlite> select * from data;

型

ただし、型を指定せずにインポートしたばあい、textとして扱われてしまうため、数値として扱いたいばあい具合がよくありません。

sqlite> select typeof(max) from data limit 1;
text

テキストを数値に変換する関数が必要ですが、明確な変換関数は用意されていないようです。調べたところround関数が目的のはたらきをしてくれることがわかりました。

www.sqlite.org

sqlite> select round('123');
123.0
sqlite> select typeof(round('123')); -- 型を確認
real
sqlite> select round('123.456');
123.0
sqlite> select round('123.456', 2); -- 桁数を指定
123.46

次のように書くことで目的を達成できました。

sqlite> select * from data where round(max) >= 100;

ヘッダ

SQLite3は、デフォルトでは結果にヘッダを出力しません。

sqlite> select 123 as value;
123

結果をCSVとして取得したいので、.headersコマンドでヘッダの出力を有効にします。

sqlite> .headers on
sqlite> select 123 as value;
value
123

SQLにまとめる

SQLite3にはファイルに出力する方法も用意されています。

www.sqlite.org

ですが今回はリダイレクトを使ってファイルに出力する方法を取ることにします。

まず、ここまでの作業をまとめたSQLファイルを作成します。

-- filtering.sql
.mode csv
.import src.csv data
.headers on
select * from data where round(max) >= 100;

これをリダイレクトでsqlite3コマンドに流し込み、結果もリダイレクトでファイルに出力します。

$ sqlite3 tmp.db < filtering.sql > dst.csv

Rubyでsqlite3コマンドを実行する

RubyからSQLite3を操作するなら、gem sqlite3 を利用するのがマットウな手段と理解していますが。

rubygems.org

手っ取り早く、sqlite3コマンドを実行することで解決してしまいます。

require 'tempfile'

src_filename = 'src.csv'
dst_filename = 'dst.csv'

sql = <<~SQL
.mode csv
.import #{src_filename} data
.headers on
select * from data where round(max) >= 100;
SQL

Tempfile.open do |tempfile|
  IO.popen("sqlite3 #{tempfile.path}", 'r+') do |io|
    io.write sql
    io.close_write

    File.open(dst_filename, 'w') do |dst_file|
      IO.copy_stream(io, dst_file)
    end
  end
end