集合演算のnumpyとpandasの速度比較

概要

先日、numpyのsetdiff1dをpandasのisinを使って書き直したらかなり速くなった、という経験をしました。一般的に成り立つのか調べてみました。

結論としては、和集合、積集合、差集合の3種類の集合演算に関して、調べたデータに関して、Google Colab上ではnumpyとpandasで大差ありませんでした。しかし、データサイズと計算時間の関係が両者で違うことが示唆されました。従って、計算環境やデータによってはnumpyの方が大幅に速かったりpandasの方が大幅に速かったりする場合もあるのではないかと思います。

比較対象のコード

• 和集合 (numpy)

  np.union1d(arr1, arr2)
  

• 積集合 (numpy)

  np.intersect1d(arr1, arr2)
  

• 差集合 (numpy)

  np.setdiff1d(arr1, arr2)
  

• 和集合 (pandas)

  pd.concat([sr1, sr2[~sr2.isin(sr1)]])  # 他にも書き方は色々
  

• 積集合 (pandas)

  sr1[sr1.isin(sr2)]
  

• 差集合 (pandas)

  sr1[~sr1.isin(sr2)]
  

比較条件

• Google Colabを使用
• ランダムな文字列を各要素とする2つの集合を使用
• 文字列はアルファベットと数字のみで構成
• 文字列の長さは10で固定 (可変長も試しましたが同様の結果だったので省略します)
• 集合の要素数は100から10Mまで
• 2つの集合の重複率は0.2
• 各集合は予めソート済み
• 各集合は重複なし
• np.Array型とpd.Series型に予め変換済み
• 試行回数は10回ずつ (ランダムデータは同じ)

結果

横軸は各集合の大きさ、縦軸は計算時間です。集合サイズ1M以上では、numpyとpandasの計算時間はほぼ同程度でした。なお、ソートのみの計算時間も参考に示しています。

よく見ると、集合が小さい場合はnumpyの方が速く、中程度ではpandasの方が速いことが分かります。全体的に見ると、numpyでは両対数プロットでほぼ直線になっているのに対し、pandasでは途中までは下に凸の曲線を描いています。

もちろん、0.001秒程度の差は、多数回繰り返す場合を除き、仕事の効率に全然影響ありません。しかし、両者の振舞いに違いがあることから、環境やデータによっては大きな差となる可能性もあると考えられます。

ソースコード全体

!pip install perfplot
import numpy as np
import pandas as pd
import perfplot, string

def gen_data(n):
  length =10
  overlap_ratio = 0.2
  random_cut_ratio = 0
  chars = np.array(list(string.ascii_letters + string.digits))
  words = np.random.choice(chars, size=(int((2-overlap_ratio)*n), length))
  if random_cut_ratio == 0:
    words = words.astype(object).sum(axis=1).astype(str)
  else:
    words[np.random.rand(*words.shape) < random_cut_ratio] = ' '
    words = words.astype(object).sum(axis=1).astype(str)
    words = np.char.replace(words, ' ', '')
  words = np.unique(words)
  arr1, arr2 = words[:n], words[-n:]
  arr1, arr2 = np.sort(arr1), np.sort(arr2)
  sr1, sr2 = pd.Series(arr1), pd.Series(arr2)
  return arr1, arr2, sr1, sr2

n_repeat = 10
def repeat(func):
  def wrapper(*args, **kws):
     return [func(*args, **kws) for _ in range(n_repeat)]
  return wrapper

@repeat
def numpy_union(arr1, arr2, sr1, sr2):
  return np.union1d(arr1, arr2)

@repeat
def numpy_intersect(arr1, arr2, sr1, sr2):
  return np.intersect1d(arr1, arr2)

@repeat
def numpy_setdiff(arr1, arr2, sr1, sr2):
  return np.setdiff1d(arr1, arr2)

@repeat
def pandas_union(arr1, arr2, sr1, sr2):
  return pd.concat([sr1, sr2[~sr2.isin(sr1)]])

@repeat
def pandas_intersect(arr1, arr2, sr1, sr2):
  return sr1[sr1.isin(sr2)]

@repeat
def pandas_setdiff(arr1, arr2, sr1, sr2):
  return sr1[~sr1.isin(sr2)]

@repeat
def sort_only(arr1, arr2, sr1, sr2):
  return np.sort(arr1), np.sort(arr2)

pp = perfplot.bench(setup = gen_data,
                    kernels = [numpy_union, numpy_intersect, numpy_setdiff,
                               pandas_union, pandas_intersect, pandas_setdiff, sort_only],
                    labels = ['numpy_union', 'numpy_intersect', 'numpy_setdiff',
                               'pandas_union', 'pandas_intersect', 'pandas_setdiff', 'sort_only'],
                    n_range=np.logspace(2,7,15), equality_check=None, xlabel='Length')
pp.show()
pp.save('tmp.png')

補足: numpyとpandasで結果が同じ (ソートの有無は除く) ことを確認する部分は省略しました。