Haskellに再入門して100万エッジのグラフネットワークでDijkstra探索するまで

これは何

以前、様々な言語でDijkstraによるグラフネットワーク探索処理の実行速度を比較した のですが、これに長年の夢だったHaskell版を追加したときにかなり苦労しました。

ほぼポエムなので技術記事として有用とは言いにくいのですが、Haskellでメモリ利用効率を気にしたプログラミングするときの参考程度になれば、と思い記事にしています。

結論だけ先に述べると、Data.Vector.Growable モジュールを使って変更可能な配列を作り、その副作用をIOモナドに閉じ込める で解決しています。

Dijkstraベンチの実コードは GitHub にあります。

ハマったこと時系列

かなり前にHaskell入門してからかなり間が空いてしまいほとんど記憶にないので、ほぼゼロから再入門しました。dataとnewTypeとclassの違いが、読んで理解はした気になるものの覚えきれずおぼつかないという入門レベルです。

開発環境は Windows 11 + WSL2 ubuntu 22.04 です。

Haskell開発環境の作成方法

調べると色々出てくるのですが、おおよそ cabal か stack が流行っているという理解に達しました。

様々な言語のバージョン管理を一元化する上で asdf を使い慣れているため、Haskellもこれで導入できるものから探しました。

が、asdf plugin list all | grep haskell で見つかった asdf-haskell をみると stack のラッパーとのこと。これを入れた場合、そもそも asdf でやりたかった「asdfでバージョンを直接指定しての導入」ができなさそう。いったん asdf 経由での導入は見送ることに。

そうこうするうちに ghcup を見つけたのでひとまず ghcup tui 経由で cabal と ghc を導入することにしました。よさそう

※以下画像は公式サイトより引用

VSCode 拡張

Haskell の拡張 を入れておく。ここはあまり悩まず進められました。

検出されるエラーがコンパイルした時と違うことが多々あるが、ひとまず気にしないことにします。

本題：大きめのデータの扱い方に悩む

前提とするグラフネットワークのデータ構造

今回のベンチマークでは100万エッジのグラフネットワークをCSVから読み込んでメモリ内に保持し、それに対してDijkstraアルゴリズムで探索処理を行っています。

このグラフネットワークデータをメモリに保持するうえで、Goのコードを例にすると以下のような構造にしています。特にこれでなければならないというわけではないのですが、メモリ内への格納効率をそこそこ意識して、無駄になりにくいデータ構造を採用しています。

type NodeId int
type NodeIndex int
type Distance int32

type Edge struct {
	first  NodeIndex
	second Distance
}

type G struct {
	id2idx map[NodeId]NodeIndex
	idx2id []NodeId
	idx    NodeIndex
	edge   [][]Edge
}

ポイントは、元のグラフネットワークで使われている各ノードの番号 NodeId を そのままエッジの両端データとしては利用せず、一度連番で NodeIndex に採番しなおしているところです。

そうすることでグラフのエッジのデータを密な配列として持つことができ、O(1)の中でも最速である配列の「要素を特定する」という処理のみでエッジデータを取得することができます。

ただ、最終的に結果を見るときには元の NodeId と NodeIndex の相互変換も必要になるので、以下の対応表も用意しています。

• idx2id: NodeIndex から NodeId を逆に引くための配列（これもO(1)で参照できる)
• id2idx: 過去に採番した NodeIndex を覚えておいて NodeId からたどるためのハッシュマップ。ここはベンチマーク本体のDijkstra探索中は一切参照しないので、妥協してハッシュマップにしています

この対応表も含めると、元のグラフネットワークデータでのノード番号がintの範囲なら普通にハッシュマップで持っても大して容量的・性能的には変わらないような気もします。

が、ここはこのベンチマークのレギュレーション上のこだわりとして

• 元のノード番号が仮に文字列でも性能容量面のインパクトが少ない
• 言語によっては標準のハッシュマップの性能に差が出てしまってベンチマークの支配的な要素になると、言語ごとの差が把握しにくい

などの要件を想定して、この構造を採用しています。

Haskellで大きな配列をどう持つか

さて、このグラフデータ構造にでてくる配列をHaskellでどう実装するかというところで大変ハマったので、順に内容を見ていきます

Listは使えない子

Haskell入門の記事ではListが頻出しているので、単純にListで実装しました。

[] で書けるのも直感的だし、++ での気軽なサイズ拡張や パターンマッチングも使いやすくて便利です。

先のグラフデータを以下のように実装してみました。

Graph.hs

import qualified Data.IntMap as IntMap

type NodeIndex = Int -- IdとIndexを明確に区別するため型に別名をつけておく
type NodeId = Int 
type Distance = Int -- 距離も実体がIntなので区別しておく

data Edge = Edge { nodeIndex :: NodeIndex, distance :: Distance } deriving Show

data G = G {
    id2idx :: IntMap.IntMap Int,
    idx2id :: [NodeId],
    edge :: [[Edge]],
    idx :: NodeIndex
}

しかしHaskellのListは使えない子でした。本当に単方向リストとして実装されているので、n番目の要素の参照ですら先頭から順にたどってしまいO(n)。O(1)なのはheadとtailくらいです。

これではせっかくエッジ情報を密な配列に詰めなおしたのに全く意味がない……

また、Dijkstraでは 全ノードに対して「その時点での開始ノードからの最短距離」を保持する配列を使い、その内容を更新しながら探索を進める処理があります。

その「更新」をListで実現するには

let newList = take (n-1) oldList ++ [newValue] ++ drop n oldList

のようにする必要があり、これもO(n)で、しかも全ノード分の情報を持つ配列を、更新処理のたびに新たに生成することになりかなり処理負荷がありそうです。

ともかく実行してみましたが、100万行の元データは読み込めるものの、探索処理が何分待っても帰ってこなくなりました。

Strictに手を出すのは早すぎた

ここでHaskelは遅延評価ということが気になり始めます。

正格評価してやれば計算量が減るかもしれない! と思い、.cabal ファイルのオプションを設定してみます。

    default-extensions:  Strict StrictData

しかしこれは逆効果で、遅延評価の時はすぐに終了していたCSV読み込みが返ってこなくなりました。ファイルサイズを1/10くらいにするとようやく読み込める。

これまではグラフデータを探索処理から参照する段になって遅延評価が始まって帰ってこなくなっていたのが、ロード時の正格評価になったから、処理が遅くなるポイントが前倒しになったのかな……など想像はしますが、証拠の押さえ方も対処もわからない。

指定をStrictだけにしたりStrictDataだけにしたり、あるいはソースコードの let の変数に ! をつけたりを繰り返しましたがいまいち挙動が把握できない。

一旦この線は撤退することにしました。

Vectorではまだ不十分

いろいろググると「Listは使えない子なのでData.Vectorを使いなさい」という記事が見つかります。

こちらは内部では配列を持っており、要素の参照はO(1)でできます。良いですね。

グラフの構造をVectorに差し替えます。

Graph.hs

import Data.Vector as V

data G = G {
    id2idx :: IntMap.IntMap Int,
    idx2id :: V.Vector NodeId,
    edge :: V.Vector (V.Vector Edge),
    idx :: NodeIndex
}

要素の参照はリストの list !! nと違って vec ! n だったり、要素の追加は list ++ [a] ではなく snoc vec a (snocはconsの逆)など細かい違いをひとつひとつ覚えるのがつらい。

また、初期化をどうするんだ？というのもHaskellに慣れていないうちはハマりました。 APIマニュアルのコンストラクタ部分 を読みなさいということを理解。

let a = V.replicate 1 0
let b = V.singleton 0

など少しずつ理解が進む。

次に気になるのは探索で使う要素の更新はどうするのか？？です。

take と drop はあるが、結合がない……? 存在しないのか、見つけられていないだけなのかもわからない……

そうする中、// を発見。

let newVec = V.// [(n,newVal)] oldVec

で要素を更新できそう。

こちらに差し替えて実行してみるも、まだ検索で固まるのは改善せず。

やはり内容を更新するたびにVector全部を再度つくりなおすようなままでは無理なのか……

Data.Vector.Unboxed.Mutable で沼に

さらにググると、Vector 配下の Mutable というモジュール名が目につき始めます。

MVector という型クラスがあり、以下のようなことができそう。

import Data.Vector.Unboxed.Mutable MV

-- mvec は (Distance, NodeIndex) を格納する可変配列
let mvec = MV.何らかの初期化 :: MVector (Distance, NodeIndex) 

-- n番目の要素をnewValで更新
...
MV.write mvec n newVal 
...

ただ、Unboxedとは……というあたりをあまりちゃんと理解しないまま 機械的に V.Vector を MV.MVector に置き換えて実装してみるが、全くコンパイルが通らなくなりました。

初期化をどうすればよいのか？関数に引き渡すときにそもそもなんという型で定義すればよい? MVector Edge と書いてもダメなようで、MVector とは何者？？ 「MVectorのMVector」（エッジデータの部分に使う）はどういう型を宣言すればよい？？ record型は使えるのか？？ など疑問百出。

ひとつひとつなんとか理解していくことにします

Unboxedとは

冷静にData.Vectorのドキュメントを読み、値をそのまま配列に格納するものがUnboxと理解。そうでないものは対象を指すポインタ的なものを配列に格納する。

• Intなど基本的な型について、Unboxed用の実装がそれぞれ用意されている
• タプルも中で２配列に分解して構成してくれる

などを理解。C言語のFFIで使うStorableとの違いも表面的には把握。

Data.Vector.Generic モジュールの役割も理解はしたものの、どう委譲されていくのか具体的には想像がまだ追いつかない。いったん先に進めます。

MVector の初期化はどうするのか

• MV.singleton はなさそう
• MV.new はあるけど最初の1つ値を追加するのはどうすれば?
• MV.replicate n initialValue を発見。探索用の配列はサイズも初期値も決まっているのでこれで初期化できる

そして、初期化した結果は let ではなく <- で「アクションから値を取り出す」必要がある……! 単に let hoge を hoge <- に書き換えると既存の関数がコンパイルエラーに。

過去Scalaをやっていた名残で 「<- がモナドのflatMapの糖衣構文」という認識はあるものの、do の中の各行も同じくflatMapの糖衣構文のはずで、それらとの違いが理解できていなかったり、MVectorを別の関数に引き渡したりする中で結局どう書けばよい、というレベルでも理解できていない。

少し基本的なHaskellの学習に戻ることに。

do記法や<-やIO周りを復習

こちらの記事 で改めて勉強。かなりクリアになりました。関数内で <- でとりだしたあと最後に関数が返す値をモナドにしておき、return hogehoge で返せばよいと。

do記法とセットでこのような処理を書いていると、pure相当の処理なのにわざわざreturnという名前になっている気持ちか少し理解できた。

MVector (PrimState m) a

ようやくMVector の定義の意味が分かってきて、MVector 中身の型 だけではなく、もう一つの型パラメータ PrimState に状態をどこに保持するかを定義してあげないといけないことを理解。

ここに MV.RealWorld と書けばIOモナドとセットで状態を扱ってくれる（副作用を扱えるようにしてくれる）と思えばよいのか。※この記述はかなりマサカリをもらいそうなことは承知の上で書いています……

探索処理の中で純粋関数で書いていた処理をIOモナドを返すように修正すると、型定義 MVector RealWorld (Distance,NodeIndex) を引数に含む関数のコンパイルが通るように。

ここが今回一番大きな山を越えた瞬間だった気がします。

しかし、まだ探索処理が返ってこない状況は改善せず。

Data.Vector.GrowからData.Vector.Growableへ

これは探索用の配列だけではなく、グラフネットワーク側のVectorもMVector化しなければならないだろうと考えてそちらも対応することに。

ここで、グラフ構築処理ではCSVを読みながら順次内容を追記していくので、MVectorをどうやってサイズ伸長させていけばよいかわからない問題にぶつかる。

ググると こちらの記事 が見つかる。2020年になってこのレベルのモジュールがまだ公開されてなかったのか…… 意外とHaskell若いという気になってきます。

この記事の最後に grow-vector: Mutable vector with efficient appends が出ているとあるので、先にそちらを試してみるが、依存ライブラリのバージョンが合わなかったため断念。

この記事の data-vector-growable を試す。こちらのモジュールは push があるので大変使いやすい印象。

最終的にグラフデータ部分は以下になりました

Graph.hs

import qualified Data.Vector.Growable as GV

data G = G {
    id2idx :: IntMap.IntMap Int,
    idx2id :: GV.GrowableUnboxedIOVector NodeId,
    edge :: GV.GrowableIOVector [Edge],
    idx :: NodeIndex
}

CSVロード処理を書き直したところ、無事ベンチマーク動作！

あわせて、リベンジで StrictData のみ指定しておきました。

最終的に Haskellのベンチ結果は3秒程度で、最速のC++の1.5倍程度。意外と(?)ほかの言語とも十分に健闘している速度。コンパイルできるまでは長いが、できてしまえば処理速度は速めの言語だったのかという認識に。

最後に

意外と速度は出た、とはいえ、メモリ上のデータを想像しながらのプログラミングには、Haskellは相当辛い印象ではありました。やりたいことに対して知るべきことが多すぎる……

Haskellは趣味やコンピュータサイエンス系の学習としてはかなり面白いですが、いろいろな点でかなり高いレベルまで理解しないと、ほかの言語では普通にできることでもすぐにハマり脱出に苦労するため、今のところは業務で使える気が全くしないのがつらいところです。

が、そこまて悲観せずとも、考えようによっては Data.Vector.Growable の使い方を覚えたということは、今後メモリ管理でハマっても、全部自分で管理しているヒープでなんとかすることもできなくはないということで、実は大きな山を越えたのかも？

本記事に誤認識や、よりよい参考資料などありましたら、ぜひ指摘・紹介ください！