[Julia]データフレームのフィルターを早くする方法

Question

実現したいこと

データフレームからリストに入っている値を取り出すコードを早くしたいと考えています。

a = DataFrame(a = rand(1:10, 1000000), b = rand(1:3, 1000000))
testList = [1,2,3,4]
a[in.(a.a, [testList]), :]

発生している問題・分からないこと

自身でいろいろ試した結果これが一番早かったのですが、実際のユースケースではこれがかなりの時間を消費しており改善したいと思っています。
実際は30個の値が入っている辞書をループし毎回255,856行, 12列のデータフレームから値を取り出す作業を行っています。

for (key, val) in diffChecked
    group = df[in.(df.val, [val]), :]
end

Profviewを使用して使用時間を見てみたところ、構造としては上から
broadcast.jlのmaterialize→copy→copyto_nonleaf!→getindex→_broadcast_getindex→_broadcast_getindex_evalf→in
となっていました。

該当のソースコード

特になし

試したこと・調べたこと

• teratailやGoogle等で検索した
• ソースコードを自分なりに変更した
• 知人に聞いた
• その他

上記の詳細・結果

試しにfilter関数に変更してみたりsetを使用してみたりと変更を加えたのですがやはり元の形が一番早いという結果でした。

a = DataFrame(a = rand(1:10, 1000000), b = rand(1:3, 1000000))
testList = [1,2,3,4]
testList2 = Set([1,2,3,4])

a[in.(a.a, [testList]), :]
filter(:a => in(testList2), a, view=true)

テスト1
@benchmark a[in.(a.a, [testList]), :]
テスト2
@benchmark filter(:a => in($testList2), $a, view = true)

テスト1では平均時間が8.184msだったのがテスト2では8.395msという結果でした。

補足

特になし

Answer 1

僕の質問への回答としてはujimushi_sradjpさんのものが最速です。
しかしそもそもデータフレームに対してinを使う事自体が遅かったので、先に条件に合うindicesをリストとして作成しそれを用いて新しいデータフレームを作成する事で数倍早くなりました。
※僕の例はdfが数百万行あり、辞書のループ毎に数十万行を抜き出す際の話です。

↓元のコード
for (key, val) in diffChecked
    group = df[in.(df.val, [val]), :]
end

↓改善版
for (key, val) in diffChecked
    inList = sorted_intersect2(df.val, sort(val))
    group = df[inList, :]
end

上記ではリンクで提示されている関数の戻り値に変更を及ぼしたものです。
しかしfindallを使用するやり方でも同様に早くなりました。

Answer 2

ほぼ最速だと思いますが，実環境ではデータフレームの行数・列数が多いので，
result = df[条件,:]だと，条件に合致した行を抽出したデータフレームの新規作成
になり要素のコピーが発生し，そこで時間がかかっているのではないでしょうか?

例えば同じコードで@viewマクロを使ってみてはどうでしょう?
これなら抽出したデータフレームは元のaのビュー扱いとなるので，
要素のコピーが発生せず時間が減るのではないかと思います。

一応手持ちの環境での@benchmarkの例です。

julia> @benchmark @view a[in.(a.a, [testList]), :]
BenchmarkTools.Trial: 1013 samples with 1 evaluation.
 Range (min … max):  4.733 ms …   5.780 ms  ┊ GC (min … max): 0.00% … 13.86%
 Time  (median):     4.896 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   4.934 ms ± 163.691 μs  ┊ GC (mean ± σ):  0.52% ±  2.28%

         ▃▄▅▇██▄▁                                              
  ▂▁▂▃▃▅▆████████▅▃▃▂▂▁▁▂▂▂▁▁▂▁▁▂▂▁▁▁▁▁▂▁▁▂▂▁▃▃▃▂▃▂▂▁▁▂▂▂▂▃▃▃ ▃
  4.73 ms         Histogram: frequency by time        5.59 ms <

 Memory estimate: 3.17 MiB, allocs estimate: 11.

julia> @benchmark a[in.(a.a, [testList]), :]
BenchmarkTools.Trial: 808 samples with 1 evaluation.
 Range (min … max):  5.918 ms …   8.107 ms  ┊ GC (min … max): 0.00% … 10.39%
 Time  (median):     6.076 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   6.175 ms ± 285.954 μs  ┊ GC (mean ± σ):  1.12% ±  2.83%

       ▄▄█▆▅▁                                                  
  ▂▂▂▅███████▆▄▄▂▂▂▂▂▂▂▁▂▁▁▁▁▁▁▁▂▁▁▂▁▂▁▁▂▁▁▁▂▁▂▂▃▃▃▃▄▄▂▂▁▂▁▁▂ ▃
  5.92 ms         Histogram: frequency by time        7.05 ms <

 Memory estimate: 9.27 MiB, allocs estimate: 23.

追記

結果はDataFrameではなくSubDataFrameになるので，大元のDataFrameを変更すると変になったりするので，少し注意が必要かもしれません。

再追記

なお，テストについては，実際のDataFrameと同じ12列でやった方がいいと思います。
例えば次のように12列のデータフレームを設定した場合，

a = DataFrame(
    NamedTuple(map(x -> x => rand(1:10, 1000000), Symbol.('a':'l')))
)

手元の環境だと

@benchmark a[in.(a.a, [testList]), :] : 12.929 ms
@benchmark filter(:a => in($testList2), $a, view = true) : 9.107 ms

と実行時間が逆転します。これはコピー有無の差ですね。

@benchmark @view a[in.(a.a, [testList]), :] : 4.746 ms

なので，これが一番速そうです。なお，蛇足ですが

df = @view a[in.(a.a, [testList]), :]

のような感じで使います。

Comments

[akirayoshimura]

回答いただきありがとうございます！
@viewを用いたところ144.562424sから144.024033sに短縮されました！
更に大きなデータを扱うときもあるので、そちらなどでは大きく改善しそうです。