ndjsonをソートするプログラムを書いた

github.com

経緯

qrnを作ってからndjsonとかJSON Linesとか呼ばれているJSONを一行にして改行で並べたフォーマット*1がなかなか合理的で便利なことに気づいた。

たとえば

  • 複数行のデータ(SQLなど)をエスケープして一行にまとめられる
  • jqで加工・フィルタリングしやすい

とか。

一方でcoreutilsみたいな基本的なツールはなかなか見つからなくて、複数ファイルをJSONのキーでJOINしてキーの値でソートする、みたいなことをやろうとすると、ツール探す・ツール作る・sed/awk/ruby/etc..とひと手間かけなくてはいけなくて、なにかそういうツールが欲しいと思っていた。

sortを作る

GUN sortやBSD sortの良いところは、どれだけ大きいファイルの入力しても、時間はかかってもソートしてくれるところで、あれはどうなっているのかなと調べたら外部マージソートを使っている…らしい。

なるほどと思って、Rustでファイルを指定されたキャパシティでチャンクファイルに分けて、ファイルとメモリでマージソートするcrateを愚直に書いてみた。

github.com

このcrateを使ってsortプログラムを書いてみたところ、目論見通りメモリ使用量は抑えつつそこそこの速度でファイルをソートできたので「じゃあ次はndjsonやるか」と意気揚々としていたのだが。

ndjsonをソートする

serde_jsonを使って実装した。

が…遅い…。219MBぐらいのndjsonのソートで2分半ぐらいかかる。

何が遅いかというと実装していて薄々気づいていたが、比較用の関数内でJSONのデシリアライズを行っており、それが比較のたびに毎行毎行 行われるため、処理に時間がかかってしまっている。

JSONの指定されたキーの値でソートしているので、デシリアライズしないわけにはいかない。

sort_by_cached_keyを使う

そういうソートキーの抽出に馬鹿みたいに時間がかかることを見越してかRustにはすでにsort_by_cached_key()というおあつらえ向きのメソッドが用意されていたので、それを使うように変更する。

ex_merge_sort crateに新しく機能を追加するにはけっこうな改修が必要そうだったので、新しくex_merge_sort_by_keyというcrateを作成した。

github.com

jlsortのライブラリをex_merge_sort_by_keyに変更したところ、現実的な時間でソートが行えるようになった。 キャッシュにメモリを使っていることもあって--capacityオプションで指定した値と2倍ぐらい乖離があるが、メモリ使用量を制限できているし、まあ…

--capacity 4gと指定して3.4GBのndjsonをソートしたところ、実行時間が2分24秒、メモリ使用量が8GBいかないくらいだったので、巨大なndjsonをソートするときはまあまあ実用的に使えると思う。

% ls -lah 16xsalaries.ndjson
-rw-r--r--  1 sugawara  staff   3.4G  8 22 16:07 16xsalaries.ndjson

% time -f "Time:%E, Memory:%M KB" jlsort -k to_date -c 4g 16xsalaries.ndjson > /dev/null
Time:2:42.51, Memory:7729932 KB

ちなみにさっきはてブで見つけた超便利そうなCSV/TSV/JSONのスイスアーミーナイフmillerで上記のndjsonをソートしたところ、30GB以上のメモリを使いつつ20分たっても結果は返ってこなかったので、現時点ではjlsortは巨大ndjsonに対してはそこそこ有用なツールな気がする。

さらに速度を上げるためには

これ以上速度を上げるとなるとすっきりとしたやり方は難しそうで、デシリアライズが速いフォーマットにいったんまとめて変換してからソートすることを思いついたが、Apache ArrowとかFlexBuffersを横目に見つつ、めんどくさそうなので手が動いていない。

追記メモ

サンプルのsalaries.ndjsonだとsort -t, -k4でもソートできそうなので手元のMacで少し検証してみた。

sort(たぶんBSD?)

% gtime -f "Time:%E, Memory:%M KB" sort -t, -k4 16xsalaries.ndjson > /dev/null
^CCommand exited with non-zero status 255
Time:15:41.08, Memory:17817920 KB

15分待って結果は返ってこず。

GNU sort

% gtime -f "Time:%E, Memory:%M KB" gsort -t, -k4 16xsalaries.ndjson > /dev/null
Time:4:23.87, Memory:4195664 KB
% gtime -f "Time:%E, Memory:%M KB" gsort -t, -k4 -S $((1024 * 1024 * 16)) 16xsalaries.ndjson > /dev/null
Time:4:45.35, Memory:9277104 KB
% gtime -f "Time:%E, Memory:%M KB" env LANG=C gsort -t, -k4 -S $((1024 * 1024 * 16)) 16xsalaries.ndjson > /dev/null
Time:0:29.30, Memory:9276800 KB

LANG=C速い。