【Julia】Mmap.mmapで「配列の配列」を作った話

## はじめに
Juliaでバイナリデータを読み込む手順がよくわからなかったので、色々と調べた結果をメモする。
特にメモリに乗らない大きなデータを処理したいので、[Mmap.mmap](https://docs.julialang.org/en/v1/stdlib/Mmap/)が使えないか試してみた。
環境は以下の通り。

+ macOS Mojave 10.14.6
+ JupyterLab 1.1.4
+ Julia 1.3.1

**※下記の内容は随所でテキトーなことをいっているので、情報源としては信頼しないようご注意ください。**

## 目的
やりたいことは大規模な時系列データの読み込みであり、以下のような状況を想定している。

+ 複数の連番ファイルに各時刻のデータが1次元配列としてバイナリで保存されている
+ データをベクトルや行列として読み込み、色々な線形計算を行いたい
+ ただし、データ量が大きく、いくつかのファイルをメモリに読み込むことはできるが、すべてのファイルを1つの行列にまとめるとメモリ不足になる
+ 各データは密なベクトルであり、構成される行列も密行列となるため、疎行列用のライブラリは適さない

もちろん、データを逐一読み込みながら処理すれば一応メモリ不足は回避できるが、それだとあからさまに遅そうなのでもう少しスマートにやりたい。そこで、ファイルをメモリにマッピングして利用可能な[Mmap.mmap](https://docs.julialang.org/en/v1/stdlib/Mmap/)で処理できないか試してみる。

## テストコード
### 問題設定
$n \gg m$の縦長の乱数行列$\mathbf{A} \in \mathbb{R}^{n \times m}$を生成し、メモリに保持可能な小さな行列になるよう$\mathbf{A}^{\top} \mathbf{A} \in \mathbb{R}^{m \times m}$を計算してみる。
ここでは、$n = 2,000,000$、$m = 100$とし、数値は正規乱数から生成した。倍精度実数を使う場合、トータルのデータ量は約1.6GBになる。

以下の3つの場合で計算時間とメモリ量を見てみる。

1. mmapを使わず、普通の行列として計算
2. 行列全体を1つのファイルに保存し、mmapで読み込む
3. 行列を列ごとに個別のファイルに保存し、mmapで読み込む

計算時間とメモリ量は[BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl)の`@benchmark`マクロを使用して確認した。
なお、行列の生成やファイルの書き込み部分は計測の対象外としている。

### データの準備
まずは下記のコードでテストデータを作成する。

```julia
using Random

# データ行列
n = 2_000_000
m = 100
A = randn(Float64,(n,m))

# 行列を1つのファイルに書き込む
# ヘッダーとして行数と列数を付加
open("test_mat.bin", "w") do io
    write(io, n)
    write(io, m)
    write(io, A)
end

# 行列を列ごとに個別のファイルに書き込む
# ヘッダーとして行数を付加
for i in 1:m
    open("test_mat$i.bin", "w") do io
        write(io, n)
        write(io, A[:,i])
    end
end
```

これにより、$n \times m$の行列が保存された1つのファイルと、$n$行のデータベクトルが保存された$m$個のファイルが作成される。

### mmapを使わず、普通の行列として計算
まずは普通に計算した場合。行列$\mathbf{A}$はすべてメモリに保持されている。

```julia
using LinearAlgebra

function test_mat(A)
    A'*A
    nothing
end
```

計測は以下のようにするだけ。

```julia
@benchmark test_mat(A)
```

結果は以下の通り。

```
BenchmarkTools.Trial: 
  memory estimate:  78.20 KiB
  allocs estimate:  2
  --------------
  minimum time:     756.787 ms (0.00% GC)
  median time:      759.918 ms (0.00% GC)
  mean time:        762.470 ms (0.00% GC)
  maximum time:     767.908 ms (0.00% GC)
  --------------
  samples:          7
  evals/sample:     1
```

当然だが、メモリ割り当てはほとんど生じておらず、後述の結果と比べても計算時間は一番小さかった。ただし、関数の外側で行列をメモリに入れているので、Juliaのプロセスを見ると1.8GB程度のメモリ消費になっていた。そのため、データ量が増えた場合にはやはりこの手法は使えない（実際には数十GB以上のデータを処理したいので）。

### 行列全体を1つのファイルに保存し、mmapで読み込む
次に行列を保存したファイルからデータを読み込んだ場合を見てみる。

```julia
function test_mmap_mat()
    f = open("test_mat.bin", "r")
    row = read(f, Int)
    col = read(f, Int)
    A = Mmap.mmap(f, Matrix{Float64},(row,col))
    close(f)
    A'*A
    nothing
end
```

結果は以下の通り。

```
BenchmarkTools.Trial: 
  memory estimate:  79.09 KiB
  allocs estimate:  19
  --------------
  minimum time:     1.146 s (0.00% GC)
  median time:      1.190 s (0.00% GC)
  mean time:        1.183 s (0.00% GC)
  maximum time:     1.214 s (0.00% GC)
  --------------
  samples:          5
  evals/sample:     1
```

メモリに保持した場合に比べて1.5倍程度の時間がかかっているが、それ以外は同様のように見える。実際のメモリ消費はこちらの方が圧倒的に小さいため、この程度の速度低下は許容範囲ではないだろうか。計算時間の増加割合はデータ量に応じて大きくなるかもしれないが、数倍までならメモリ不足にならない方がいいと思う。

### 行列を列ごとに個別のファイルに保存し、mmapで読み込む
最後に本命の分割されたファイルを読み込む場合。

```julia
function test_mmap_vec(m)
    A = []
    for i in 1:m
        f = open("test_mat$i.bin", "r")
        row = read(f, Int)
        vec = Mmap.mmap(f, Vector{Float64},row)
        push!(A,vec)
        close(f)
    end
    
    # 通常の行列ではないため、愚直に計算
    B = zeros(m,m)
    for j in 1:m
        for i in 1:m
            B[i,j] = dot(A[i],A[j])
        end
    end
    nothing
end
```

上記のコードでは、ファイルから読み出したベクトルを1つにまとめるために`push!`を使っている。このとき、`A`のタイプは`Array{Float64,2}`ではなく、`Array{Any,1}`となり、その要素として`Array{Float64,1}`が格納されている形になる。扱いに注意が必要な形ではあるが、これであればマッピングそのものはかなり高速に実行できる。[^1]

[^1]: 後で確認したら`A = Vector{Float64}[]`として初期化すると結果は`Array{Array{Float64,1},1}`になるので、型を指定した方が分かりやすいとは思う。

ちなみに`hcat`や`append!`を使うとあからさまに遅く、メモリ消費もかなり大きかった。これはおそらくベクトルを追加する度にメモリを取り直しているためだと思われる。結果は行列になるので扱いは容易だが、mmapの利点が失われるのでやる意味は薄そうだ。

```
BenchmarkTools.Trial: 
  memory estimate:  507.03 KiB
  allocs estimate:  22210
  --------------
  minimum time:     8.489 s (0.00% GC)
  median time:      8.489 s (0.00% GC)
  mean time:        8.489 s (0.00% GC)
  maximum time:     8.489 s (0.00% GC)
  --------------
  samples:          1
  evals/sample:     1
```

行列積の計算は他と比べてかなり遅かった。これは結果が行列になっておらず、ベタ書きのアルゴリズムを使っているせいだと考えられる。ちゃんとした行列積のアルゴリズムを使えば1つのファイルから読み込んだ場合と同程度になるのではないだろうか。

そこで、列ごとに処理できるよう、列ベクトルのノルムを計算したときの時間を見てみる。

```julia
function test_mmap_vec2(n,m)
    A = []
    for i in 1:m
        f = open("test_mat$i.bin", "r")
        row = read(f, Int)
        vec = Mmap.mmap(f, Vector{Float64},row)
        push!(A,vec)
        close(f)
    end
    
    # 1列ごとにアクセスする計算に変更
    norm.(A)
    nothing
end
```

その結果、下記のようにかなり高速化された。

```
BenchmarkTools.Trial: 
  memory estimate:  118.86 KiB
  allocs estimate:  2313
  --------------
  minimum time:     522.845 ms (0.00% GC)
  median time:      641.839 ms (0.00% GC)
  mean time:        641.067 ms (0.00% GC)
  maximum time:     765.137 ms (0.00% GC)
  --------------
  samples:          8
  evals/sample:     1
```

そのため、既存のライブラリが有効利用できるベクトル単位での処理なら問題なく使えるのだと思われる。とはいえ、汎用的ではないため、扱いやすい小規模な行列に落とし込むための工夫が必要だろう。この辺は今後の課題だ。

## おわりに
あんまり良さげな結果でもないが、mmapで分割されたファイルを読み込むことはできたので、ひとまずはこれでよしとしよう。
`push!`で作った「配列の配列」は正直使い勝手が悪いので、素直に中間ファイルに書き出してから改めて行列としてマッピングするのもありなんじゃないかと思ったが、数十GB単位のファイルコピーなんてさすがにやりたくない。
なので、列ベクトルだけが参照できる状態でも計算可能な低ランク近似アルゴリズムを適用し、早々に生データにアクセスしないで済むようにすべきだろう。これはこれで興味があるので、また色々調べてやってみようと思う。

<!--
※以下の内容はやり直したら差が出なかったのでなし！
関数を同名で再作成してたので、おそらく初回のコンパイル時のメモリ割り当てと比較してたっぽい。
ただまぁ、型を指定した方が配列の中身が見えるので、方針に変更はなし。

追記：配列の初期化
mmapの配列を初期化する際、空の配列の型を指定することでメモリと計算時間を小さくできたのでメモしておく。

下記のように`A = []`と初期化すると以下のようになる。

```
# 初期化を
# A = [] 
# とした場合の実行結果

  0.044276 seconds (161.73 k allocations: 8.306 MiB)
 typeof(A) : Array{Any,1}
 sizeof(A) : 1208
```

このデータは列数が151のものを使用しているので、データ量は$8\ \mathrm{byte} \times 151 = 1208\ \mathrm{byte}$になっている。
一方、型を明示的に宣言してやると以下のようになる。

```
# 初期化を
# A = Vector{Float64}[] 
# とした場合の実行結果

  0.018078 seconds (43.12 k allocations: 2.256 MiB)
 typeof(A) : Array{Array{Float64,1},1}
 sizeof(A) : 1208
```

当然といえば当然だが、抽象型の`Any`より具象型`Float64`などにする方がメモリ割り当てが少なくなる。
計算時間も約半分になっているので、今後はこちらを使っていこうと思う。

なお、配列の配列は`hcat(A...)`で2次元配列に変換できるが、このときはやはりメモリが取り直されて普通の行列になってしまい、メモリ消費が大きくなる。mmapのまま行列にできればそれが一番なのだが、現状その方法がわからないため、配列の配列で使えるアルゴリズムを考えてみるという方針には変わりなさそうだ。
-->