![見出し画像](https://assets.st-note.com/production/uploads/images/97834260/rectangle_large_type_2_a3137e89bca4515e3c11fdeee5838609.png?width=1200)
高速データフレームライブラリ「Polars」のススメ - 特許データの組織集計を例に
はじめに
下記の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に書いた。
![](https://assets.st-note.com/img/1676169983648-ZFHihj1f98.png?width=1200)
処理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を使って、クロス集計表を整えている。
![](https://assets.st-note.com/img/1676170781533-va8rO3oBpe.png?width=1200)
![](https://assets.st-note.com/img/1676170851923-j4dRszXk7Z.png?width=1200)
#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を使う。出力結果は下記の通りである。
![](https://assets.st-note.com/img/1676171196464-vqSCHfJHN3.png?width=1200)
処理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を使う必要がある。出力結果は以下のようになる。
![](https://assets.st-note.com/img/1676171537683-HyYDmz6gVw.png?width=1200)
計算時間
14万件の特許データ(出願人は34,000)での計算時間は下記の通り。処理1は約1/2、処理2は約1/10になった。
![](https://assets.st-note.com/img/1676173398723-U01X8A2AGr.png)
PandasとPolarsの文法の違いをどう調べるか
公式ドキュメント・参考記事
確かにPolarsはPandasライクではあるが、細かいところで文法・メソッド・引数が異なったり、同じようなメソッドでも出力結果が異なるケースがよくある。この点は下記のようなドキュメントや記事を参照しながら、試していくと良いだろう。
また、色々と調べても分からない場合や、あまり時間がないときは、処理が重いと考えられるところだけをPolarsに書き換え、他はPandasのままにしておくようなハイブリッド形式も良いと思う。
Twitterで呟くと公式アカウントが教えてくれる
また、わからないことがあれば、ツイートしてみるのも良い。Polarsの公式アカウントが教えてくれた。
![](https://assets.st-note.com/img/1676172473494-0xNyKEsMJp.png?width=1200)
`https://t.co/2YHlhw8QQJ(sorted(df.columns))`
— polars data (@DataPolars) February 11, 2023
ChatGPTに聞いてみる
最初、下記のように質問すると、Pandasのコードが返答してきた。しかも、やりたいことは列名で列をソートすることだったので、コードのロジックも違う(これは私の聞き方も悪かったと思う)。さらにpolarsで書いて欲しいという依頼をしたが、polarsは分からないと言われてしまった。
![](https://assets.st-note.com/img/1676172743198-437hyVYkn7.png?width=1200)
そこで、自分のやりたいことをPandasのコードで示し、これをPolarsに変える方法を聞いてみるとうまくいった。何らかのコード変換をChatGPTに聞く際は、このように具体的なコードを示して尋ねると良いだろう。
![](https://assets.st-note.com/img/1676172708126-Y8dmkprtsS.png?width=1200)
全コード
最後に、全コードを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')