文字列を変換しながらNumpyで読み込もうとした話

仕事で必要になったので、思考実験も兼ねてテストしてまとめてみました。
データ解析とがで同じようなことをやってる人は多いと思うので、他の人の参考になれば幸いです。

もっと早そうな手法があればぜひ教えて下さい。

目的

0,1,0,0,1,10,42,3.9,0.001のような文字が入ったcsvファイルをbool型のnumpyの配列に変換し、できるだけ効率よくアクセスしたい。
ただし、0はFalse、1はTrueに変換し、それ以外の値はどちらでもよいものとする。

実行環境

OS: Windows 10
Python 3.9.1
CPU: Core i7-9700(3GHz)
メモリ: 32GB(とりあえず十分量)
ディスク：HDD(想定する使用環境もHDDなため問題なし)

テストケース作成

あまり複雑なことを考えず、半分は0か1、残り半分は適当な桁数の正の実数から構成されているcsvファイルを作成した。

import os
import random

os.chdir("TestFiles")

ROW = 10 ** 5
COL = 1000


def create_test_case(case_num: int):
    print("start")
    for i in range(case_num):
        with open(f"case{i:02}.csv", "w") as file:
            for r in range(ROW):
                line = []
                for c in range(COL):
                    if c & 1:
                        line.append(str(random.randint(0, 1)))
                    else:
                        line.append(f"{(random.randint(0,9999) / 100):.2f}")

                file.write("{}\n".format(",".join(line)))
                if r % 1000 == 0:
                    print(f"\r{i} case, {r} row written", end="")

    print("\nAll compleated.")


create_test_case(10)

これを実行して、とりあえず10個のcsvファイルを作った。各ファイルの容量はおおよそ385MB。

テスト方法と手法その1

あまりちゃんと作っているとは言えないが、下記ソースのように複数のテストケースを縦につなげて時間の測定を行った。
1件のテストにかかる時間が30秒を超えることもあったので、複数回テストせずに一発勝負とした。

import datetime
import numpy as np

ROW = 10 ** 5
COL = 1000


def read_test_0():
    dest = np.zeros((ROW, COL), dtype=bool)
    with open("TestFiles\\case00.csv", "r") as f:
        for row_index, l in enumerate(f):
            dest[row_index, row_index % COL] = True

    print(dest.sum())


def read_test_1():
    dest = np.zeros((ROW, COL), dtype=bool)
    with open("TestFiles\\case00.csv", "r") as f:
        for row_index, l in enumerate(f):
            row = list(map(float, l.split(",")))
            for col_index, val in enumerate(row):
                dest[row_index, col_index] = val

    print(dest.sum())


def test_str_benchmark():
    start = datetime.datetime.now()
    read_test_0()
    end0 = datetime.datetime.now()
    read_test_1()
    end1 = datetime.datetime.now()
    print(f"読み込み時間\n{(end0-start).total_seconds():.3}秒")
    print(f"float→bool変換\n{(end1-end0).total_seconds():.4}秒")


if __name__ == "__main__":
    test_str_benchmark()

手法その1はpythonはbool型に数値を代入したとき、0をfalse、それ以外をtrueとおくことを利用した。
全部のcsvのセルをfloatに直し、それをそのままnumpyのbool型の配列に代入することで、
bool型の列は仕様どおりに代入される。

出力は次の通り

100000
74989369
読み込み時間
1.64秒
float→bool変換
28.65秒

全体の半分のセルは実数なのでほぼ全部True、残りの半分は0か1かランダムなので、
この出力結果は正しいように見える。

個人的にはこれだけ処理がかかっていると無限ループになってないか心配になるので、
もう少し早く処理を行いたい。
欲を言えば10秒以内。

手法その2

最初に掲げた仕様なら、別に全部少数に変えなくても、最初の1文字目が1かどうかで確認したほうが早そうに見える。
その点を突いたプログラムが下記のもの。

def read_test_2():
    dest = np.zeros((ROW, COL), dtype=bool)
    with open("TestFiles\\case00.csv", "r") as f:
        for row_index, l in enumerate(f):
            row = l.split(",")
            for col_index, val in enumerate(row):
                dest[row_index, col_index] = (val[0] == "1")

    print(dest.sum())

出力と時間は次の通り

30493448
最初の文字だけの読み取り1
16.28秒

大体25%+5%がTrueになるはずなので、出力は正しそうだ。
おおよそ処理時間が半分になっている。

手法その3

では、読み込みをバイナリとしての読み込みにして、文字の比較を行わなければどうなる？

def read_test_3():
    dest = np.zeros((ROW, COL), dtype=bool)
    one = ord("1")
    comma = ord(",")
    nl = ord("\n")
    with open("TestFiles\\case00.csv", "rb") as f:
        result_bytes = f.read(1)
        index = 0
        dest[index // ROW, index % COL] = (result_bytes[0] == one)
        index += 1
        result_bytes = f.read(1)
        use_next = False
        while len(result_bytes) != 0: # EOFになるまで
            if use_next:
                dest[index // ROW, index % COL] = (result_bytes[0] == one)
                index += 1
                use_next = False
            else: # ラインフィードかコンマの次が読み取り対象
                if result_bytes[0] == nl or result_bytes[0] == comma:
                    use_next = True
            result_bytes = f.read(1)

    print(dest.sum())

出力。想定外に遅い上に何か出力がおかしいな。
後で最終的なindexの値をデバッグで確認してもROW*COLに見えたので、何か見落としがありそうだ。

304822
最初の文字だけの読み取り2
63.94秒

出力が遅いのはファイルを1文字ごと読み込むのが原因のような気がするので、
バグは置いておいて文字列として読み込む方法の確認を継続する。

手法その4

ファイルを1行単位で読み込んだ上で、手法3でやった処理方法を合わせてみる。

def read_test_4():
    dest = np.zeros((ROW, COL), dtype=bool)
    with open("TestFiles\\case00.csv", "r") as f:
        for row_index, l in enumerate(f):
            use_next = True
            col_index = 0
            for c in l:
                if use_next:
                    dest[row_index, col_index] = c == "1"
                    col_index += 1
                    use_next = False
                if c == ",":
                    use_next = True

    print(dest.sum())

出力結果は次の通り。

30493448
最初の文字だけの読み取り3
23.55秒

手法2より遅くなっているのが意外だった。pythonぐらいの高級言語になってくると
コンパイラ側に文字列の処理は任せたほうが早いのかもしれない。

手法その5

手法2をベースとして、頭文字を文字列じゃなく数値として比較させたらちょっとぐらい早くなったりしないのか？

def read_test_5():
    dest = np.zeros((ROW, COL), dtype=bool)
    with open("TestFiles\\case00.csv", "r") as f:
        for row_index, l in enumerate(f):
            row = l.split(",")
            for col_index, val in enumerate(row):
                dest[row_index, col_index] = (ord(val[0]) & 1)

    print(dest.sum())

出力結果は次の通り。

52491005
最初の文字だけの読み取り1-2
21.31秒

これでだめとなると、コードの短さがそのまま処理の速さにつながっていそうだ。

手法その6(現時点での最速)

自分のコードミスから生まれた下記コードが結局最速だった。

def read_test_6():
    dest = np.zeros((ROW, COL), dtype=bool)
    with open("TestFiles\\case00.csv", "r") as f:
        for row_index, l in enumerate(f):
            row = l.split(",")
            for col_index, val in enumerate(row):
                dest[row_index, col_index] = (val == "1") # 間違えて文字列全体との比較をしている

    print(dest.sum())

24944449
最初の文字だけの読み取り1
14.78秒

個人的な教訓

C++みたいな言語ならともかく、pythonでは無駄に策を巡らせるよりコード量を減らしたほうが高速化につながることがあるようだ。
今後は保守性と速度の面から、文字列関連の処理はできるだけコードを簡潔にすることを心がけることにする。