db-benchmarkのデータ生成ちょっとだけ速くした話．

はじめに

アドベントカレンダー21日目の記事です．
この記事ではデータフレームライブラリの性能測定のためのベンチマークの一つであるdb-benchmarkのデータ生成をちょっとだけ高速化した話をします．
このベンチマークはAWS上で指定されたしたマシン環境で，pandasやpolarsといったオープンソースのデータフレームライブラリの性能を測定しているのですが，ベンチマーク測定のためのデータ生成が4時間くらいかかって非常に遅いです．
データ生成に時間がかかる理由としては50Gとかのサイズのデータを作ることが最も大きな理由ですが，それ以外にも様々な特性を持ったデータを1個ずつ生成していることも時間がかかる要因の一つです．
以下は，データ生成を並列で行い10%くらい速くできたという内容です．

方針と実装

方針はシンプルで
1.Rで書かれたソースコードをpythonから呼び出す．
2.pythonを並列化する
という方針です．
単純に乱数列を作るならpythonからRを呼び出す必要は全くないのですが，元のソースコードを実行した場合と全く同じ結果を出力する必要があるため，このような方針を取りました．
ソースコードは以下です．

import rpy2.robjects as robjects
from rpy2.robjects.packages import importr
from concurrent.futures import ProcessPoolExecutor
import warnings

warnings.filterwarnings("ignore")
data_table = importr('data.table')

def get_formatted_number(number):
    formatted_number = "{:.0e}".format(number)
    formatted_number = formatted_number.replace('e+0', 'e').replace('e+','e').replace('e0','e')
    return formatted_number

args_list = [(1e7, 1e2, 0, 0), (1e7, 1e1, 0, 0), (1e7, 2e0, 0, 0),(1e7, 1e2, 0, 1),(1e7, 1e2, 5, 0),
             (1e8, 1e2, 0, 0), (1e8, 1e1, 0, 0), (1e8, 2e0, 0, 0),(1e8, 1e2, 0, 1),(1e8, 1e2, 5, 0),
             (1e9, 1e2, 0, 0), (1e9, 1e1, 0, 0), (1e9, 2e0, 0, 0),(1e9, 1e2, 0, 1),(1e9, 1e2, 5, 0)]
def data_gen(i):
    r = robjects.r
    N = args_list[i][0]
    K = args_list[i][1]
    nas = args_list[i][2]
    is_sort = args_list[i][3]
    file_name = 'G1_'+ get_formatted_number(N) +  '_' + get_formatted_number(K) + '_' + str(nas) +  '_' +str(is_sort) + '.csv'
    r_code = f'''
    library(data.table)
    N <- {N}
    K <- {K}
    nas <-{nas}
    is_sort <-{is_sort}
    output <- "{file_name}"
    set.seed(108)
    DT = list()
    DT[["id1"]] = sample(sprintf("id%03d",1:K), N, TRUE)
    DT[["id2"]] = sample(sprintf("id%03d",1:K), N, TRUE)
    DT[["id3"]] = sample(sprintf("id%010d",1:(N/K)), N, TRUE) 
    DT[["id4"]] = sample(K, N, TRUE)                          
    DT[["id5"]] = sample(K, N, TRUE)                         
    DT[["id6"]] = sample(N/K, N, TRUE)                       
    DT[["v1"]] =  sample(5, N, TRUE)                          
    DT[["v2"]] =  sample(15, N, TRUE)                        
    DT[["v3"]] =  round(runif(N,max=100),6)
    setDT(DT)
    if (nas>0L) {{
      cat("Inputting NAs\n")
      for (col in paste0("id",1:6)) {{
        ucol = unique(DT[[col]])
        nna = as.integer(length(ucol) * (nas/100))
        if (nna)
          set(DT, DT[.(sample(ucol, nna)), on=col, which=TRUE], col, NA)
        rm(ucol)
      }}
      nna = as.integer(nrow(DT) * (nas/100))
      if (nna) {{
        for (col in paste0("v",1:3))
          set(DT, sample(nrow(DT), nna), col, NA)
      }}
    }}
    if (is_sort==1L) {{
      cat("Sorting data\n")
      setkeyv(DT, paste0("id", 1:6))
    }}
    write.csv(DT, file = output, row.names = FALSE)
    '''
    r(r_code)

num = len(args_list)
if __name__ == "__main__":
    with ProcessPoolExecutor() as executor:
         executor.map(data_gen, range(num))

結果

手元で動かしたところだいたい10%くらい速くなった感じです(もっと速くしたかったが断念)．
一番のボトルネックは50Gのデータ生成で，乱数列そのものの並列化とかしないと速くならなさそう.