バイナリサーチで処理時間を 44時間→3分 に短縮した話

問題

＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝

2つの CSV ファイルを内部結合する形で、新たな CSV ファイルを生成するバッチ処理がありました。あなたならどう実装しますか？

＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝

もう少し要件を詳しく書きます。

• language.csv と book.csv で LANGUAGE が一致するレコードを result.csv に書き出す
• 画像のように結合キー LANGUAGE には重複がある
• language.csv は 800MB、50,000 件程度のデータ
• book.csv は 100MB、5,000件 程度のデータ
• Java（SE8）アプリケーションが EC2 上で稼働する

※ ここで登場するファイルやカラム名はあくまでイメージです。

最初の実装【 44時間 かかる処理】

2つの CSV ファイルを一度オンメモリに乗っけて、for で愚直になめるということをやりました。

実装はこんな感じでした。
なお1点補足すると、CSV は行×列という特性上 List<List<String>> のような2重リスト構造となっています。

List<LanguageLine> languageLines = CsvStream.lines(reader::read)
        .map(LanguageLine::new) // ファーストクラスコレクション化
        .collect(Collectors.toList());

List<BookLine> bookLines = CsvStream.lines(reader::read)
        .map(BookLine::new) // ファーストクラスコレクション化
        .collect(Collectors.toList());

for(LanguageLine language: languageLines) {
    for(BookLine book: bookLines) {
        if(!language.getLanguage().equals(book.getLanguage())) {
            break;
        }

        writer.write(language + book);
    }
}

結論から言うと、44 時間かかる ことになります（無論 44 時間走らせたわけではなく計算上の話です）。

ちょっと改善した実装【 4 時間かかる処理】

最初の実装には2つの問題点があったと考えられます。

１．CSV を無駄に走査していること
２．2つの比較的大きな CSV を同時にオンメモリに乗せていること

2つの問題点に対する改善案として、（既存の実装を最大限生かしつつ）最初に思いついたのが次のようなアプローチです。
まず、片方（book.csv）を結合キー LANGUAGE の辞書順でソートしてオンメモリに乗せておきます。
そして、もう片方（language.csv）を1行ずつ読んでいき結合キー LANGUAGE で走査していくのですが、辞書順でこれ以上後には存在し合えないところまで走査したら次の結合キーにスキップします。

実装はこんな感じになります。

List<BookLine> sortedBookLines = CsvStream.lines(reader::read)
        .map(BookLine::new)
        .sorted(Comparator.comparing(BookLine::getLanguage)) // LANGUAGE の辞書順でソート
        .collect(Collectors.toList());

CsvStream.lines(reader::read)
        .map(LanguageLine::new)
        .forEach(throwConsumer(language -> {
            for (BookLine book : sortedBookLines) {
                final int compared = language.getLanguage().compareTo(book.getLanguage());
                if (compared < 0) {
                    continue; // 辞書順で前の場合は次へ
                }
                if (compared > 0) {
                    break; // 辞書順で後の場合はその先に一致するものはないのでスキップ
                }

                writer.write(language + book);
            }
        }));

これにより 44時間→4時間 に処理が短縮 されました。
しかし、これでもまだ処理時間は長いです。数時間かかる処理ともなると、その間の負荷考慮やエラー時のハンドリングが煩雑になる可能性があります。

以上のことから、処理の抜本的な改善が必要と考えられました。

最終的な実装【 3分 で終わる処理】

さて、タイトルから分かる通り バイナリサーチ のアルゴリズムを組み込むことで、処理時間を 3 分に短縮 することができました。

念のため説明します。
バイナリサーチとは、2分探索とも呼ばれるデータ検索アルゴリズムのひとつです。ソートされたデータの探索範囲を半分に絞込みながら探索するアルゴリズムです。探索範囲を半分にしていくことから計算量は $O(log2n)$ とかなり小さくなります。

Java（SE8）では java.util.Arrays.binarySearch が用意されており、次のようにバイナリサーチを利用することができます。

int[] numbers = {1, 2, 3};

// 検索キーが合致したインデックス = binarySearch(探索対象のソートされたデータ, 探索キー)
int index = Arrays.binarySearch(numbers, 1); // index: 0

ただし、java.util.Arrays.binarySearch　は データの重複を考慮していません 。探索対象のソートされたデータに探索キーに合致するデータが複数存在する場合、合致するデータの いずれかのインデックス が返ってくることになります。
冒頭の要件に

画像のように結合キー LANGUAGE には重複がある

があるため、ちょっとした工夫が必要そうです。

そこで、バイナリサーチをしていずれかのインデックスが返ってきた場合には、マイナスとプラスの両方向にも操作する 独自ロジックを追加しました。

最終的な実装はこのようになりました。

List<BookLine> sortedBookLines = CsvStream.lines(reader::read)
        .map(BookLine::new)
        .sorted(Comparator.comparing(BookLine::getLanguage))
        .collect(Collectors.toList());

String[] sortedLanguage = sortedBookLines.stream() // ソートされた LANGUAGE 配列をバイナリサーチ用に準備しておく
        .map(ReviewFileLine::getLanguage)
        .toArray(String[]::new);

CsvStream.lines(reader::read)
        .map(LanguageLine::new)
        .forEach(language -> {
            final List<Integer> matchedBookLineIndexes = binarySearchIndexes(sortedLanguage, language.getLanguage());

            if (matchedBookLineIndexes.isEmpty()) {
                return;
            }

            matchedBookLineIndexes.forEach(throwConsumer(matchedIndex ->
                    writer.write(language + sortedBookLines.get(matchedIndex))
            ));
        });

// ↑メインの処理 ↓マイナスとプラスの両方向にも操作する独自ロジック

private List<Integer> binarySearchIndexes(String[] targets, String key) {
    final List<Integer> matchedIndexes = new ArrayList<>();
    final int index = Arrays.binarySearch(targets, key);

    if (index < 0) {
        return matchedIndexes;
    }

    matchedIndexes.add(index);

    // 重複レコードをマイナス方向に走査
    for (int i = index - 1; i >= 0; i--) {
        if (!targets[i].equals(key)) {
            break;
        }
        matchedIndexes.add(i);
    }
    // 重複レコードをプラス方向に走査
    for (int i = index + 1; i < targets.length; i++) {
        if (!targets[i].equals(key)) {
            break;
        }
        matchedIndexes.add(i);
    }

    return matchedIndexes;
}

まとめ

JVM メモリチューニングや分散処理実装をせずとも、アルゴリズム一本で処理時間を短縮できたのは爽快でした。今後なにかの参考になればと思います。