LegalOn Technologies Engineering Blog

LegalOn Technologies 開発チームによるブログです。

Rustで標準ライブラリのHashMapの代わりにhashbrownを使う

Tips

こんにちは。株式会社LegalOn Technologies でエンジニアをしている赤部です。

自然言語処理をしていると、単語の出現回数を数えたり、単語と何らかのデータを紐付けたりすることが頻繁に必要になります。これらのことを簡単に行える、最も一般的でよく知られたデータ構造はハッシュマップではないでしょうか？

プログラミング言語 Rust では、標準ライブラリの std::collections::HashMap を用いればハッシュマップを簡単に導入できます。しかし、場合によっては標準ライブラリの HashMap ではなく hashbrown クレートを利用したほうが、シンプルに効率的なコードを実装できるかもしれません。

この記事では、まず hashbrown クレートを紹介し、コード例とともに hashbrown クレートの使いどころを紹介します。

hashbrown クレート

Rust 1.66 時点の標準ライブラリの HashMap は、 hashbrown クレートをラップしたものとなっています。

標準ライブラリと hashbrown クレートの大きな違いは、hashbrown クレートではデフォルトで aHash という高速なハッシュアルゴリズムが使われていますが、標準ライブラリでは Python や Ruby などの他のプログラミング言語と同様に SipHash 1-3 という強力な HashDoS 対策が施されたアルゴリズムがデフォルトとなっています。また、 hashbrown クレートは no_std 環境でも alloc クレートが有効であれば利用できますが、標準ライブラリの HashMap は環境ノイズを利用してハッシュ値にランダム性を与えるため、 no_std では利用することができません。

この他に、 hashbrown クレートではハッシュマップの低レベルな API や、標準ライブラリでまだ採用されていない API が提供されています。

標準ライブラリの HashMap は最も広く最初に使われるものなので、それを踏まえれば適切な実装でしょう。しかし、速度が重視される場面や組み込みデバイス等での利用を考えた場合は、状況に応じて最適なクレートを選択する必要があります。

std::collections::HashMap で所有権を持つ型をキーにする際の問題点

Rust では、 HashMap のキーとして、データへの所有権を持つ型を指定することも、参照を指定することもできます。参照を利用すれば、データのコピーにかかるコストを削減できますが、参照先のデータが生存している期間しか利用することができません。

ではここで、入力された文字列を空白で区切り、単語の出現回数を HashMap で数え上げるプログラムを考えてみましょう。

HashMap::entry() を利用する方法

まずは entry() を利用する方法です。コードは以下のようになります。

use std::collections::HashMap;

fn count_words(text: &str) -> HashMap<String, usize> {
    let mut result = HashMap::new();
    for word in text.split(' ') {
        // ハッシュマップの指定されたキーの値をインクリメント
        // キーが登録されていなければ事前に0を追加する
        *result.entry(word.to_string()).or_insert(0) += 1;
    }
    result
}

fn main() {
    let result = count_words("I have a pen . You have a ball .");
    for (word, cnt) in result {
        println!("{word}: {cnt}");
    }
}

実行結果:

.: 2
You: 1
ball: 1
a: 2
have: 2
I: 1
pen: 1

このコードでは、引数として与えられた text の参照先が破棄されても戻り値を使い続けられるように、 HashMap のキーの型を文字列への所有権を持つ String 型としています*1。 HashMap::entry() の引数にはキーの型である String をムーブして渡す必要があります。しかし、 word は text を参照する文字列スライス (&str) であり、 String 型ではありません。このため、 str::to_string() を呼び出して word から String を生成する必要があります。 String の生成にはヒープ上の領域確保と文字列のコピーが必要なため、それなりに大きなコストとなります。

entry() を用いる場合、単語が既に登録されている場合であっても毎回 String を生成する必要がありますが、それらは明らかに無駄な動作です。

HashMap::get() と HashMap::insert() を利用する方法

そこで、単語が登録されているかを事前に確認し、必要な場合のみ to_string() を呼び出すことを考えます。コードは以下のようになります。

fn count_words(text: &str) -> HashMap<String, usize> {
    let mut result = HashMap::new();
    for word in text.split(' ') {
        if let Some(c) = result.get_mut(word) {
            // 登録されている場合はインクリメント
            *c += 1;
        } else {
            // 登録されていない場合は新しいキーに1を追加
            result.insert(word.to_string(), 1);
        }
    }
    result
}

これで無駄に String を生成する必要が無くなりました。 しかし、このコードも無駄な処理が完全に無くなったわけではありません。ハッシュ値の計算などを含めた HashMap の探索を、(1) 単語の有無を確認するときと、(2) 単語の登録を行うとき、の2回行う必要があるからです。

HashMap::raw_entry_mut() を利用する方法

そこで、現時点の Nightly 版の Rust では raw_entry_mut() というAPIが用意されています。 raw_entry_mut() は、与えられたキーから HashMap の特定のエントリへのポインタのようなものを返します。これを用いると、探索とデータの挿入を分割することができるため、必要最小限の処理で上述のコードと同様の処理を実現できます。コードは以下のようになります。

#![feature(hash_raw_entry)] // Nightly の機能を使用するために必要

use std::collections::HashMap;

fn count_words(text: &str) -> HashMap<String, usize> {
    let mut result = HashMap::new();
    for word in text.split(' ') {
        *result
            .raw_entry_mut()
            .from_key(word)
            .or_insert_with(|| (word.to_string(), 0))
            .1 += 1;
    }
    result
}

しかし、Nightly 版の Rust は API が変更される可能性があり、使い続けられる保証はありません。実際に、この API に対しては問題点が指摘がされ、見直しが行われています。

主な指摘は以下の2つです。

  • or_insert_with() で from_key() で指定したキーとは全く関係のないキーを登録できてしまう
  • API が複雑

この API の使用目的として最も考えられるものは、本記事で取り上げているような「キーが登録されていないときだけ登録する」という処理を効率的に行うことです。であるならば、機能を今よりも簡略化し、より単純な API で目的を達成できるようにする必要があります。

hashbrown::HashMap::entry_ref() を利用した解決策

上述の raw_entry_mut() の問題点を受けて、 hashbrown クレートには entry_ref() という API が追加されました。 entry_ref() ではキーの参照を受け取り、必要に応じて内部で所有権を持つ型を生成する仕組みとなっています。

entry_ref() を用いると、たった1行で、 entry() を用いた場合よりも短い文字数で単語の出現回数を数えることができます。コードは以下のようになります。

use hashbrown::HashMap;

fn count_words(text: &str) -> HashMap<String, usize> {
    let mut result = HashMap::new();
    for word in text.split(' ') {
        *result.entry_ref(word).or_insert(0) += 1;
    }
    result
}

将来的には Rust の標準ライブラリでも entry_ref() を利用できるようになるかもしれません。

速度比較

ここまで hashbrown クレートを用いたシンプルで効率的なコードを紹介しましたが、実際に各コードでどれほどの速度差があるかを比較します。入力するテキストは、現代日本語書き言葉均衡コーパス (BCCWJ) のコア約6万文を短単位分割した約127万語です。

実行時間の比較 [ms]
関数std::collectionshashbrown
HashMap::entry()38.738.5
HashMap::get() + insert()26.619.8
HashMap::raw_entry_mut()26.519.7
HashMap::entry_ref()―20.0

各ライブラリの中で比較すると、文字列のコピーを抑制することで、 entry() を利用する方法に比べて2倍程度高速化されることが確認できます。標準ライブラリと hashbrown クレートで比較すると、デフォルトのハッシュアルゴリズムが高速な hashbrown が速度面で勝っていることが確認できます。

表を見る限り、 get() + insert() でハッシュマップの探索が2回行われることによる速度低下は確認できませんが、コードの単純さも鑑みれば entry_ref() を使うメリットは大きいように思います。

おわりに

この記事では、Rust の hashbrown クレートを紹介し、標準ライブラリの HashMap との比較を行いました。

実際には、処理するデータの特徴も考慮しつつ、適切なデータ構造やライブラリを選択する必要があります。しかし、標準ライブラリの HashMap が扱いづらいと感じたとき、標準ライブラリの中で使われている hashbrown クレートは有力な候補の1つになるのではないでしょうか。

メンバー募集中！

株式会社LegalOn Technologies では、SREや、バックエンドエンジニア、検索システム、研究開発に興味のあるインターン生など様々なポジションを募集しています。

ご興味がある方は以下の求人ページから求人要項をご確認いただけますので、お気軽にご応募ください！

https://herp.careers/v1/legalforce/requisition-groups/d2e157cc-120b-4ade-8879-0326c32127bd

*1:このコード例の場合、キーの型は &str でも十分です。しかし、キーを String にする必要がある状況は往々にして存在します。