高速データフレームライブラリ「Polars」のススメ - 特許データの組織集計を例に

2023年2月12日 12:41

はじめに

　下記のnoteで書いた集計について、高速化する方法を考えていた。

　Pythonでデータフレームを扱う際、Pandasを使うことが多いと思う。ただ、データ量が多くなってくると処理が遅く感じる時がある。そもそも自分の書き方が悪い部分も大いにあるのだが、プロダクト開発ではなくアドホックにデータ処理・分析を回すプログラムを書くことが多いので、高速化する労力はあまりかけず、「まぁ、プログラムを回して放置して、その間に違う作業をしていよう」、というケースが多いのが正直なところ。
　とはいえ、たまには高速化することも考えたいと思い調べてみると、Polarsというデータフレームがあることを知った。Rust製で、Pandasに比べと省メモリで高速らしい。また文法もPandasライクなので移行しやすそうだ。実際、14万件の特許データに対して526.2秒かかっていた処理が、ロジックは変えずにただPolarsに書き換えただけで53.5秒程度に短縮できた(約1/10)。
　本noteでは、Polarsを試した内容をメモしておく。また、Pandasとは文法・メソッド・引数が異なることも多いため、試した調べ方も記録しておく。

PandasをPolarsに修正したコード

　処理したいデータは下記表のような特許情報である。このデータから、処理1. applicantsの年次推移、処理2. applicants別上位keywords集計をしたい。Pandasで使ったコードは、以前のnoteに書いた。

処理1. applicantsの年次推移集計

Pandas部分をコメントアウトし、その下にPolarsによる処理を書いていく。

#import pandas as pd
import polars as pl
import sys
import collections # 処理2で利用

# ----- 処理1. applicantsの年次推移集計 ----- #
input_file = 'sample_data.tsv'
#df = pd.read_csv(input_file, delimiter='\t').fillna('')
df = pl.read_csv(input_file, sep='\t').fill_null('').fill_nan('')

PandasとPolarsはどちらもread_csvでファイルを読み込める。ただし、Polarsはdelimiterではなくsepだけであり、fillnaは使えない。その代わり、fill_null()とfill_nan()を使う。

# 区切り文字で分割
tmp = [x.split('|') for x in df['applicants']]

# 前後スペース削除
tmp2 = []
for x in tmp:
    tmp2.append([y.strip() for y in x if y])

# 空白要素削除 & 集合化(重複があれば一意化)
tmp3 =[]
for t in tmp2:
     tmp3.append(list(set([x for x in t if x])))

num = sum(len(v) for v in tmp3)

# tmp3の1要素ごとにyear、id、keywordsを突合(処理1だけであればkeywordsは不要)
i = 0
j = 0
records = []
for a in tmp3:
    for b in a:
        #record = [df.id[j], df.year[j], b, df.keywords[j]]
        record = [df['id'][j], df['year'][j], b, df['keywords'][j]]
        records.append(record)
        sys.stdout.write('\r%d/%d' % (i+1,num))
        i = i + 1
    j = j + 1

#df2 = pd.DataFrame(records, columns=['id', 'year', 'applicants','keywords'])
df2 = pl.DataFrame(records, columns=['id', 'year', 'applicants','keywords'])

最後にデータフレームかしている。データフレームの作り方はPandasもPolarsも同様である。

# applicantsとyearでクロス集計
#df_count = pd.crosstab(df2['applicants'], df2['year'])
df_count = df2.pivot('id', 'applicants', 'year', 'count').fill_null(0)
col_names = sorted(df_count.columns[1:])
col_names.insert(0, df_count.columns[0])
df_count = df_count[col_names]

applicantsとyearのクロス集計をする処理だが、Polarsではcrosstabは使えない。そこでpivotを使って「df_count = df2.pivot('id', 'applicants', 'year', 'count')」としている。引数は左からvalues(集計する対象)、index(行：集計軸1)、columns(列：集計軸2)、aggregate_fn(出力内容)である。なお、pivotをしただけでは、0はnullとなり、出力結果のカラムはソートされていない。そこで、.fill_null(0)やsortedを使って、クロス集計表を整えている。

#df_count.insert(0, 'sum', list(df_count.sum(axis=1)))
df_count = df_count.with_columns(df_count[:,1:].sum(axis=1).alias('sum'))
df_count = df_count[['applicants','sum'] + df_count.columns[1:14]]

#df_count = df_count.sort_values(by='sum', ascending=False)
df_count = df_count.sort(by='sum', reverse=True)

合計値をクロス集計表に挿入しているが、Pandasのような書き方はできない。Polarsではwith_columnsを使ってカラム(sum)を挿入する。そして最後にsumで降順にソートしている。ちなみに、Polarsにはsort_valuesやsort_indexはない。また引数もascendingではなく、reverseである。

#df_count.to_csv('applicants_year_count.tsv', sep='\t')
df_count.write_csv('applicants_year_count.tsv', sep='\t')

最後にクロス集計表を出力している。to_csvはなく、write_csvを使う。出力結果は下記の通りである。

処理2. applicants別上位keywords集計

次に、applicants別にkeywordsを集計し、上位を出力するコードを示す。上位10のkeywordsを区切り文字「|」で結合して出力する。

# ----- 処理2. applicants別上位keywords集計 ----- #
applicants_li = list(df_count['applicants'])

records2 = []
i = 0
num2 = len(applicants_li)
for x in applicants_li:
    #tmp4 = '|'.join(df2[df2['applicants'] == x].keywords)
    tmp4 = '|'.join(df2.filter(pl.col('applicants') == x)['keywords'])
    tmp5 = [x.strip() for x in tmp4.split('|')]
    tmp6 = collections.Counter(tmp5).most_common(10)
    values, counts = zip(*tmp6)
    records2.append([x, '|'.join(values)])
    sys.stdout.write('\r%d/%d' % (i + 1,num2))
    i = i + 1

#df_out = pd.DataFrame(records2, columns=['applicants', 'keywords'])
#df_out.to_csv('applicants_keywords.tsv', sep='\t', index=False)
df_out = pl.DataFrame(records2, columns=['applicants', 'keywords'])
df_out.write_csv('applicants_keywords.tsv', sep='\t')

データフレームを任意の条件でフィルタリングする際、Pandasでは「df2[df2['applicants'] == x].keywords」という書き方ができるが、Polarsではできない。「df2.filter(pl.col('applicants') == x)['keywords']」のように、filterを使う必要がある。出力結果は以下のようになる。

計算時間

14万件の特許データ(出願人は34,000)での計算時間は下記の通り。処理1は約1/2、処理2は約1/10になった。

PandasとPolarsの文法の違いをどう調べるか

公式ドキュメント・参考記事

　確かにPolarsはPandasライクではあるが、細かいところで文法・メソッド・引数が異なったり、同じようなメソッドでも出力結果が異なるケースがよくある。この点は下記のようなドキュメントや記事を参照しながら、試していくと良いだろう。
　また、色々と調べても分からない場合や、あまり時間がないときは、処理が重いと考えられるところだけをPolarsに書き換え、他はPandasのままにしておくようなハイブリッド形式も良いと思う。

Twitterで呟くと公式アカウントが教えてくれる

また、わからないことがあれば、ツイートしてみるのも良い。Polarsの公式アカウントが教えてくれた。

`https://t.co/2YHlhw8QQJ(sorted(df.columns))`
— polars data (@DataPolars) February 11, 2023

ChatGPTに聞いてみる

　最初、下記のように質問すると、Pandasのコードが返答してきた。しかも、やりたいことは列名で列をソートすることだったので、コードのロジックも違う(これは私の聞き方も悪かったと思う)。さらにpolarsで書いて欲しいという依頼をしたが、polarsは分からないと言われてしまった。

　そこで、自分のやりたいことをPandasのコードで示し、これをPolarsに変える方法を聞いてみるとうまくいった。何らかのコード変換をChatGPTに聞く際は、このように具体的なコードを示して尋ねると良いだろう。

全コード

最後に、全コードを1つにまとめたものを下記に示す。

#import pandas as pd
import polars as pl
import sys
import collections # 処理2で利用

# ----- 処理1. applicantsの年次推移集計 ----- #
input_file = 'sample_data.tsv'
#df = pd.read_csv(input_file, delimiter='\t').fillna('')
df = pl.read_csv(input_file, sep='\t').fill_null('').fill_nan('')

# 区切り文字で分割
tmp = [x.split('|') for x in df['applicants']]

# 前後スペース削除
tmp2 = []
for x in tmp:
    tmp2.append([y.strip() for y in x if y])

# 空白要素削除 & 集合化(重複があれば一意化)
tmp3 =[]
for t in tmp2:
     tmp3.append(list(set([x for x in t if x])))

num = sum(len(v) for v in tmp3)

# tmp3の1要素ごとにyear、id、keywordsを突合(処理1だけであればkeywordsは不要)
i = 0
j = 0
records = []
for a in tmp3:
    for b in a:
        #record = [df.id[j], df.year[j], b, df.keywords[j]]
        record = [df['id'][j], df['year'][j], b, df['keywords'][j]]
        records.append(record)
        sys.stdout.write('\r%d/%d' % (i+1,num))
        i = i + 1
    j = j + 1

#df2 = pd.DataFrame(records, columns=['id', 'year', 'applicants','keywords'])
df2 = pl.DataFrame(records, columns=['id', 'year', 'applicants','keywords'])

# applicantsとyearでクロス集計
#df_count = pd.crosstab(df2['applicants'], df2['year'])
df_count = df2.pivot('id', 'applicants', 'year', 'count').fill_null(0)
col_names = sorted(df_count.columns[1:])
col_names.insert(0, df_count.columns[0])
df_count = df_count[col_names]

#df_count.insert(0, 'sum', list(df_count.sum(axis=1)))
df_count = df_count.with_columns(df_count[:,1:].sum(axis=1).alias('sum'))
df_count = df_count[['applicants','sum'] + df_count.columns[1:14]]

#df_count = df_count.sort_values(by='sum', ascending=False)
df_count = df_count.sort(by='sum', reverse=True)

#df_count.to_csv('applicants_year_count.tsv', sep='\t')
df_count.write_csv('applicants_year_count.tsv', sep='\t')

# ----- 処理2. applicants別上位keywords集計 ----- #
applicants_li = list(df_count['applicants'])

records2 = []
i = 0
num2 = len(applicants_li)
for x in applicants_li:
    #tmp4 = '|'.join(df2[df2['applicants'] == x].keywords)
    tmp4 = '|'.join(df2.filter(pl.col('applicants') == x)['keywords'])
    tmp5 = [x.strip() for x in tmp4.split('|')]
    tmp6 = collections.Counter(tmp5).most_common(10)
    values, counts = zip(*tmp6)
    records2.append([x, '|'.join(values)])
    sys.stdout.write('\r%d/%d' % (i + 1,num2))
    i = i + 1

#df_out = pd.DataFrame(records2, columns=['applicants', 'keywords'])
#df_out.to_csv('applicants_keywords.tsv', sep='\t', index=False)
df_out = pl.DataFrame(records2, columns=['applicants', 'keywords'])
df_out.write_csv('applicants_keywords.tsv', sep='\t')