大きめのCSVファイルの集計速度について比較してみた

Ruby, Go, Rust の3つの言語で大きめの CSV ファイルの集計速度について比較してみました。
ここでは、ビルドにかかる時間については触れません。

使ったコードは https://github.com/okkez/csv-aggregation-example にあります。

使用する CSV ファイルについて

クラウドサービスの使用料をサービスごとに集計することをイメージしてみました。

• サイズ: 235M

• 行数: 3000001 (ヘッダーを除くとちょうど300万行)

• カラム

  • id: integer
  • name: string
  • description: string
  • cost: float

生成には generator/main.go を使いました。
テストデータの生成も最初は Ruby で faker を使ってやろうと思ったんですが、時間がかかりすぎたので Go でサクっと書きました。

集計方法

前項で示したようなテーブル costs が RDBMS 上にあるとしたら、以下のようなイメージです。

select
  name,
  sum(cost) as cost
from
  costs
group by name
order by name
;

name ごとに cost の合計値を計算し name の順に出力する感じです。

Ruby

特にひねらずに、素直に書きました。空行を除くと8行で圧倒的に短いです。

$ ruby -v
ruby 3.0.0p0 (2020-12-25 revision 95aff21468) [x86_64-linux]

require "csv"

name_to_cost = Hash.new(0)

CSV.foreach(ARGV[0], headers: true) do |row|
  name_to_cost[row["name"]] += row["cost"].to_f
end

name_to_cost.sort_by {|k, _| k }.each do |name, cost|
  printf "%s\t%.3f\n", name, cost
end

Go

Ruby 版と同じ結果を出力するように書きました。型をわかりやすくするために一度 struct に変換しています。main だけで 46 行あるのとエラー処理に関するコードが多い印象です。あんまり調べないで書いたので、もう少し簡潔に書けそうな気はしました。

$ go version  
go version go1.15.6 linux/amd64

package main

import(
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"sort"
	"strconv"
)

type Record struct {
	ID          uint64  `csv:"id"`
	Name        string  `csv:"name"`
	Description string  `csv:"description"`
	Cost        float64 `csv:"cost"`
}

func main() {
	file, err := os.Open(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	nameToCost := map[string]float64{}

	csvReader := csv.NewReader(file)
	csvReader.LazyQuotes = true
	_, err = csvReader.Read()

	for {
		row, err := csvReader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		id, _ := strconv.ParseUint(row[0], 10, 64)
		cost, _ := strconv.ParseFloat(row[3], 64)
		record := &Record{
			ID: id,
			Name: row[1],
			Description: row[2],
			Cost: cost,
		}
		if _, ok := nameToCost[record.Name]; ok {
			nameToCost[record.Name] += record.Cost
		} else {
			nameToCost[record.Name] = record.Cost
		}
	}

	keys := make([]string, 0, len(nameToCost))
	for key := range nameToCost {
		keys = append(keys, key)
	}

	sort.Strings(keys)

	for _, key := range keys {
		fmt.Printf("%s\t%.3f\n", key, nameToCost[key])
	}
}

Rust

run が 21 行なので割と短く書けたと思います。あと集計後に name でソートしないようにするために BTreeMap を使いました。これも Ruby 版と同じ結果を出力するようにしました。

$ rustc --version
rustc 1.48.0 (7eac88abb 2020-11-16)

use std::collections::BTreeMap;
use std::env;
use std::error::Error;
use std::fs::File;
use std::io::BufReader;

use serde::{Deserialize};

#[derive(Deserialize, Debug)]
struct Record {
    id: u64,
    name: String,
    description: String,
    cost: f64,
}

fn run() -> Result<(), Box<dyn Error>> {
    let args: Vec<String> = env::args().collect();
    let f = File::open(&args[1])?;
    let b = BufReader::new(f);
    let mut csv_reader = csv::ReaderBuilder::new().has_headers(true).from_reader(b);
    let mut name_to_cost = BTreeMap::new();

    for result in csv_reader.deserialize() {
        let record: Record = result?;
        let cost = match name_to_cost.get(&record.name) {
            Some(value) => *value,
            None => 0.0,
        };
        name_to_cost.insert(record.name.clone(), record.cost + cost);
    }

    for (name, cost) in name_to_cost {
        println!("{}\t{:.3}", name, cost);
    }

    Ok(())
}

fn main() {
    run().unwrap()
}

結果

言語	処理時間(sec)	メモリ使用量 (kilo byte)
Go	2.89	8860
Rust	1.42	2068
Ruby	27.35	17344

順当に Rust が最速で省メモリでした。

番外: Ruby と Arrow

Arrow はあんまり詳しくないんですが CSV の読み込みやら集計が速いと聞いたのでやってみました。
読み込みと全ての行の集計はとても速いんですが group by 的なことをするとものすごく遅くなってしまいました。どう書けば速くできるのかはわかりませんでした。

なお、手元での環境構築が大変だったので Docker を使って環境を構築しました。

require "arrow"

schema = Arrow::Schema.new([
  Arrow::Field.new("id", :string),
  Arrow::Field.new("name", :string),
  Arrow::Field.new("description", :string),
  Arrow::Field.new("string", :float),
])

name_to_cost = Hash.new(0)
table = Arrow::Table.load(ARGV[0], schema: schema)
table = table.select_columns do |column|
  %w(name cost).include?(column.name)
end
table.each_record_batch do |batch|
  batch.each do |row|
    name_to_cost[row["name"]] += row["cost"]
  end
end

name_to_cost.sort_by {|k, _| k }.each do |name, cost|
  printf("%s\t%.3f\n", name, cost)
end

real    9m14.356s
user    9m14.916s
sys     0m1.228s

コード読んでたら group by できる API を見つけたので使ってみました。
自分が調べた時点ではこの API は experimental 扱いで Ruby で実装されていました。

require "arrow"

schema = Arrow::Schema.new([
  Arrow::Field.new("id", :string),
  Arrow::Field.new("name", :string),
  Arrow::Field.new("description", :string),
  Arrow::Field.new("string", :float),
])

name_to_cost = Hash.new(0)
table = Arrow::Table.load(ARGV[0], schema: schema)
table = table.select_columns do |column|
  %w(name cost).include?(column.name)
end
g = table.group("name")
result = g.sum
result.each_record_batch do |batch|
  batch.each do |row|
    name_to_cost[row["name"]] += row["cost"]
  end
end

name_to_cost.sort_by {|k, _| k }.each do |name, cost|
  printf("%s\t%.3f\n", name, cost)
end

real    164m23.954s
user    164m12.100s
sys     0m9.232s

実行時間は2時間半を越えました。