この記事は MicroAd Advent Calendar 2023 の6日目の記事です。
この記事ではpolarsの処理をSQLで書けてしまう素敵機能、polars SQL を紹介しようと思います。
polars SQLの使い方
はじめに、polars SQLの利用方法を簡単に紹介します。
- SQLContext オブジェクトを作成
- SQLContext にデータフレームをテーブルとして登録
- SQLContext.exexute でクエリを実行
ctx = pl.SQLContext(films=df)
ctx.execute('SELECT * FROM films LIMIT 5', eager=True)
1. SQLContext オブジェクトを作成
polars SQLではまず始めにデータフレームとクエリ内のテーブルを紐づけるために、polars.SQLContext オブジェクトを作成します。
2. SQLContext にデータフレームをテーブルとして登録
基本的にはSQLContextの初期化時にテーブルを登録するのが楽でおすすめなのですが、Python Polars 0.17.13 で追加された機能で、それ以前のバージョンでは以下のSQLContextのメソッドで登録する必要があるため注意が必要です。
3. SQLContext.exexute でクエリを実行
その後、execute メソッドで所望のクエリをデータフレームに対して実行できます。
この時 eager=True を指定するとDataFrame、デフォルトでは遅延評価を行えるLazyFrameが返されます。
polars SQLで利用可能な構文
以下は個人的によく使うSQL構文がサポートされているかどうかをまとめています。（Python Polars 0.19.19 時点）
polarsでは日々アプデが行われていてpolars SQLについても頻繁に機能追加が行われているので、最新のサポート対象を知りたい場合はリリースノートを読むことをお勧めします。
利用可
- SELECT
- DESTINCT
-
*, EXCLUDE - CASE
- WHERE
- (LEFT, INNER, CROSS) JOIN
- GROUP BY
- ORDER BY
- HAVING
- LIMIT
- OFFSET
- WITH句での一時テーブル作成
利用不可
- サブクエリ
- 配列の要素へのアクセス
- 配列の展開
- WINDOW関数
SELECTやGRPUBY BY、JOINなど基本的な構文はサポートしているようです。また、WITH句での一時テーブルもサポートしているため、込み入ったクエリでも問題なく実行できます。
一方で、配列や構造体へのアクセスは一切サポートしていないようでした。これらを扱う場合は事前に展開してから扱う必要がありそうです。
polars SQLの実行速度
polars SQLを利用してしまうとpolarsの利点である高速処理が失われてしまうのではないかという不安があり、polars Expressionsで書いた場合と比較してみました。
結論からになるのですが、確認した範囲ではpolars SQLとpolars Expressionsには明確な実行速度の差はありませんでした。
| polars SQL | polars Expressions | |
|---|---|---|
| SUM | 934 µs ± 215 µs | 934 µs ± 83.5 µs |
| GROUP BY | 564 µs ± 57.9 µs | 528 µs ± 47.6 µs |
| ORDER BY | 65.6 ms ± 3.39 ms | 65.2 ms ± 2.61 ms |
| LEFT JOIN | 8.4 ms ± 493 µs | 8.05 ms ± 717 µs |
以下、確認を行なった環境とソースコードです。
環境
- MacBook Air Apple M2
- Python 3.11
- polars 0.19.17
比較に利用したデータフレーム
# numpyでランダムなデータフレームを作成
df = pl.DataFrame(np.random.randint(0, 10, (1_000_000, 4)))
shape: (1_000_000, 4)
┌──────────┬──────────┬──────────┬──────────┐
│ column_0 ┆ column_1 ┆ column_2 ┆ column_3 │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 │
╞══════════╪══════════╪══════════╪══════════╡
│ 0 ┆ 4 ┆ 2 ┆ 5 │
│ 0 ┆ 7 ┆ 5 ┆ 2 │
│ 1 ┆ 6 ┆ 9 ┆ 6 │
│ 8 ┆ 2 ┆ 2 ┆ 8 │
│ … ┆ … ┆ … ┆ … │
│ 0 ┆ 9 ┆ 0 ┆ 1 │
│ 6 ┆ 2 ┆ 0 ┆ 8 │
│ 8 ┆ 6 ┆ 1 ┆ 3 │
│ 0 ┆ 1 ┆ 2 ┆ 0 │
└──────────┴──────────┴──────────┴──────────┘
SUM
polars SQL
%%timeit
ctx = pl.SQLContext(t=df)
ctx.execute(
'''
SELECT
SUM(column_0),
SUM(column_1),
SUM(column_2),
SUM(column_3)
FROM t
''',
eager=True,
)
564 µs ± 57.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
polars Expressions
%%timeit
df.select([
pl.sum('column_0'),
pl.sum('column_1'),
pl.sum('column_2'),
pl.sum('column_3'),
])
528 µs ± 47.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
GROUP BY
polars SQL
%%timeit
ctx = pl.SQLContext(t=df)
ctx.execute(
'''
SELECT
column_0,
COUNT(*)
FROM t
GROUP BY column_0
''',
eager=True,
)
2.22 ms ± 191 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
polars Expressions
%%timeit
df.group_by('column_0').count()
2.24 ms ± 221 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
ORDER BY
polars SQL
%%timeit
pl.SQLContext(t=df).execute(
'''
SELECT
*
FROM t
ORDER BY column_0, column_1, column_2, column_3
''',
eager=True,
)
65.2 ms ± 2.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
polars Expressions
%%timeit
df.sort(['column_0', 'column_1', 'column_2', 'column_3'])
65.6 ms ± 3.39 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
LEFT JOIN
polars SQL
%%timeit
pl.SQLContext(t=df).execute(
'''
WITH t2 AS (
SELECT
column_0,
SUM(column_1)
FROM t
GROUP BY column_0
)
SELECT
*
FROM t LEFT JOIN t2
ON t.column_0 = t2.column_0
''',
eager=True,
)
8.4 ms ± 493 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
polars Expressions
%%timeit
df2 = df.group_by('column_0').agg(pl.sum('column_1'))
df.join(df2, on='column_0')
8.05 ms ± 717 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)