【Python】行数・列数が大きいCSVファイルから行番号指定で特定の行を効率よく読み出す方法についての検討

はじめに

CSVは扱いやすいデータ保存形式ですが、非常に大きなデータを扱う場合には全てのデータをメモリに載せることができない場合があります。

このような場合、ファイルから1行ずつ読み込んで処理していくことになりますが、不特定の順でデータを取得したり、データを繰り返し取得する場合は非常に効率が悪くなります。

そこで、行数の大きいCSVファイルと列数の大きいCSVファイルを対象に、行番号指定で特定の行を読み出すにはどのような方法が効率的なのかを、簡単な実験を行なって検討してみました。

テスト環境

Google Colaboratory 上で実行しました。
ハードウェア アクセラレータは None を選択しています。
2019年2月時点では以下の環境のようです。

• OS: Ubuntu 18.04.1 LTS
• Python: 3.6.7
• CPU: Intel(R) Xeon(R) CPU @ 2.30GHz

テスト用CSVファイルの作成

以下の4つのCSVファイルを使います。

• MEDIUM ROWS : 1,000,000行 × 10列 のCSV
• LARGE ROWS : 10,000,000行 × 10列 のCSV
• MEDIUM COLUMNS : 1,000行 × 10,000列 のCSV
• LARGE COLUMNS : 1,000行 × 100,000列 のCSV

今回はデータの中身は使いませんが、どの行と列なのかが確認できるように (行インデックス)-(列インデックス)という形になっているデータを用いました。
ヘッダ行はありません。また、行番号は0開始です。

medium_rows.csv

0-0,0-1,0-2, ... ,0-10
1-0,1-1,1-2, ... ,1-10
...
999998-0,999998-1,999998-3, ... ,999998-10
999999-0,999999-1,999999-3, ... ,999999-10

それぞれのCSVファイルは以下のコードで作成しました。

import csv

MEDIUM_ROWS = 1000000 # 100万行
LARGE_ROWS = 10000000 # 1000万行
MEDIUM_COLUMNS = 10000 # 1万列
LARGE_COLUMNS = 100000 # 10万列

# MEDIUM_ROWS
with open('medium_rows.csv', 'w') as f:
    writer = csv.writer(f, lineterminator='\n')
    for i in range(MEDIUM_ROWS):
        row = [str(i) + '-' + str(j) for j in range(10)]
        writer.writerow(row)

# LARGE_ROWS
with open('large_rows.csv', 'w') as f:
    writer = csv.writer(f, lineterminator='\n')
    for i in range(LARGE_ROWS):
        row = [str(i) + '-' + str(j) for j in range(10)]
        writer.writerow(row)

# MEDIUM COLUMNS
with open('medium_columns.csv', 'w') as f:
    writer = csv.writer(f, lineterminator='\n')
    for i in range(1000):
        row = [str(i) + '-' + str(j) for j in range(MEDIUM_COLUMNS)]
        writer.writerow(row)

# LARGE COLUMNS
with open('large_columns.csv', 'w') as f:
    writer = csv.writer(f, lineterminator='\n')
    for i in range(1000):
        row = [str(i) + '-' + str(j) for j in range(LARGE_COLUMNS)]
        writer.writerow(row)

ファイルの大きさは

• MEDIUM ROWS : 85Mb
• LARGE ROWS : 944Mb
• MEDIUM COLUMNS : 84Mb
• LARGE COLUMNS : 933Mb

となりました。
このくらいの大きさであればメモリに十分載りますが、今回は実験ですのでこれらを用います。

実験

いくつかの方法で、各ファイルの最初の行の読み出しと最後の行の読み出しにかかる時間を測定しました。
測定にはJupyter Notebook の %%timeit コマンドを用いました。

%%timeit -r3 -n5

セルを5回繰り返し実行して平均を取るという処理を3回行い、そのうち一番成績のいいものが表示されます。

1. Pandasで特定の行を取得

まずは普通にCSVファイルを読み込む場合によく使われるPandasを試してみます。
以下の関数を使って読み出します。
データフレームをリストに変換する処理は行なっていません。

import pandas as pd

# path: CSVファイルのパス
# idx: 行インデックス
def pd_read_row(path, idx):
    return pd.read_csv(path, header=None, skiprows=lambda x: x not in [idx])

結果

MEDIUM ROWS の最初の行を取得

%%timeit -r3 -n5
result = pd_read_row('medium_rows.csv', 0)

5 loops, best of 3: 657 ms per loop

MEDIUM ROWS の最後の行を取得

%%timeit -r3 -n5
result = pd_read_row('medium_rows.csv', MEDIUM_ROWS - 1)

5 loops, best of 3: 665 ms per loop

LARGE ROWS の最初の行を取得

%%timeit -r3 -n5
result = pd_read_row('large_rows.csv', 0)

5 loops, best of 3: 6.76 s per loop

LARGE ROWS の最後の行を取得

%%timeit -r3 -n5
result = pd_read_row('large_rows.csv', LARGE_ROWS - 1)

5 loops, best of 3: 6.82 s per loop

MEDIUM COLUMNS の最初の行を取得

%%timeit -r3 -n5
result = pd_read_row('medium_columns.csv', 0)

5 loops, best of 3: 963 ms per loop

MEDIUM COLUMNS の最後の行を取得

%%timeit -r3 -n5
result = pd_read_row('medium_columns.csv', 999)

5 loops, best of 3: 903 ms per loop

LARGE COLUMNS の最初の行を取得

%%timeit -r3 -n5
result = pd_read_row('large_columns.csv', 0)

5 loops, best of 3: 9.27 s per loop

LARGE COLUMNS の最後の行を取得

%%timeit -r3 -n5
result = pd_read_row('large_columns.csv', 999)

5 loops, best of 3: 9.28 s per loop

以上を表にすると以下のようになります。

	M ROWS	L ROWS	M COLUMNS	L COLUMNS
最初の行	657 ms	6.76 s	963 ms	9.27 s
最後の行	665 ms	6.82 s	903 ms	9.28 s

• 特定の行を指定するための引数は用意されていないので、skiprowsで必要な行以外を除外しています。そのため、最初の行でも最後の行でも時間はあまり変わりません。
• 行数に比例して取得にかかる時間も増大します。
• おそらくデータフレームを作成しているためでしょうが、行数の増加よりも列数の増加の方が負荷がかかります。

単一の行を取得するには向いていませんが、複数の行をまとめて取得して、そのままデータフレームとして扱いたい場合はPandasが一番便利だと思います。

2. csvモジュールで特定の行を取得

次に、csvモジュールを使って特定の行を取得してみます。
以下の関数を使って読み出します。

def csv_mod_read_row(path, idx):
    with open(path, 'r') as f:
        reader = csv.reader(f)
        for row in reader:
            if reader.line_num - 1 == idx:
                return row
    return None

結果

以降の結果は表だけを示します。

	M ROWS	L ROWS	M COLUMNS	L COLUMNS
最初の行	28.8 µs	26.6 µs	1.03 ms	13.5 ms
最後の行	1.39 s	15.2 s	1.25 s	16.4 s

• 行を取得した時点で処理を打ち切っているため、最初の行を取得する場合はPandasを使う場合より速くなっています。
• 一方で、最後の行を取得する場合はPandasより時間がかかっています。

ここで、行送りの

for row in reader:

のところで、行を毎回パースしているのが非常に無駄な気がします。
次はこの点を改善します。

3. 行送りをreadline()に変更してcsvモジュールで特定の行を取得

行送りをファイルオブジェクトのreadline()に変更してみます。
本来は「read」しなくてもいいのですが、純粋に行を進めるだけの方法はあるのでしょうか？

def csv_mod_read_row2(path, idx):
    with open(path, 'r') as f:
        reader = csv.reader(f)
        line_num = 0
        while True:
            if line_num == idx:
                return next(reader)
            f.readline()
            line_num += 1
    return None

結果

	M ROWS	L ROWS	M COLUMNS	L COLUMNS
最初の行	21.8 µs	29.4 µs	1 ms	13.7 ms
最後の行	269 ms	2.69 s	61.3 ms	824 ms

行送りの際のパース処理がなくなったおかげでだいぶ速くなりました。

4. split()でパース

普通はCSVをパースする場合はモジュールを使うべきですが、データが数値だけであることが保証されているような「綺麗な」データではモジュールでパースするよりも、単純にコンマでパースした方が効率がいいかもしれません。
そこで、split()でのパースを試してみます。

def simple_split_read_row(path, idx):
    with open(path, 'r') as f:
        reader = csv.reader(f)
        line_num = 0
        while True:
            if line_num == idx:
                return f.readline().strip().split(',')
            f.readline()
            line_num +=1
    return None

結果

	M ROWS	L ROWS	M COLUMNS	L COLUMNS
最初の行	24.7 µs	28.8 µs	688 µs	7.44 ms
最後の行	277 ms	2.8 s	61.1 ms	983 ms

これでは差がわかりにくいので、LARGE COLUMNSの最初の行を100回取得する時間で比較してみます。

CSVモジュールでパース

%%timeit -r3 -n5
for i in range(100):
    result = csv_mod_read_row3('large_columns.csv', 0)

5 loops, best of 3: 1.19 s per loop

split()でパース

%%timeit -r3 -n5
for i in range(100):
    result = simple_split_read_row('large_columns.csv', 0)

5 loops, best of 3: 528 ms per loop

split()でパースした方が速いという結果になりました。
扱うデータによって使い分けると良さそうです。

5. 各行の開始位置を記憶して行取得の繰り返しに対応する

ここまでの方法は一度だけデータを取得する場合はそれほど速度が気になりませんが、繰り返して行を取得する場合にはどうしてもある程度の時間がかかってしまいます。
そこで、まず事前に行の開始位置を記憶しておいて、行を取得する時はその位置からデータを読み出すようにします。

まず、以下のようなクラスを定義します。
各行の開始位置を offset_list に格納しています。offset_list のインデックスが行番号に対応しているので、seek()でその位置に移動してから行を読み出します。

import csv

class CsvRowReader:
    def __init__(self, path):
        f = open(path, 'r')
        self.file = f
        self.reader = csv.reader(f)
        self.offset_list = []
        while True:
            self.offset_list.append(f.tell())
            line = f.readline()
            if line == '':
                break
        self.offset_list.pop() # remove offset at end of file
        self.num_rows = len(self.offset_list)

    def __del__(self):
        self.file.close()

    def read_row(self, idx):
        self.file.seek(self.offset_list[idx])
        return next(self.reader)

以下のようにして使います。

# インスタンスを作成
reader = CsvRowReader('medium_rows.csv')

# 行を取得
result = reader.read_row(0)

結果

次のような結果になりました。
インスタンスの作成にかかる時間は「初期化」の行で示しています。

	M ROWS	L ROWS	M COLUMNS	L COLUMNS
初期化	5.69 s	56.3 s	64.3 ms	879 ms
最初の行	7.61 µs	10.4 µs	1 ms	16.5 ms
最後の行	7 µs	9.38 µs	1.12 ms	19.3 ms

特に行が多いと初期化に相応の時間がかかりますが、一度作成してしまえば読み出しは他の方法に比べるとかなり速いので、機械学習などでメモリに載り切らないデータを何度もロードしなければならない場合は有効な方法だと思います。

おわりに

比較的単純な方法ですが、意外とこのような情報が見つからなかったので記事にしてみました。
列数がそれほど大きくなければSQLite、列演算を行いたい場合はParquetなども選択肢に入ってくると思います。

もっといい方法があれば、ぜひ教えていただけると幸いです。