polarsの処理をSQLで書くには

この記事は MicroAd Advent Calendar 2023 の6日目の記事です。

https://qiita.com/advent-calendar/2023/microad

この記事ではpolarsの処理をSQLで書けてしまう素敵機能、[polars SQL](https://pola-rs.github.io/polars/py-polars/html/reference/sql.html) を紹介しようと思います。

## polars SQLの使い方

はじめに、polars SQLの利用方法を簡単に紹介します。
1. SQLContext オブジェクトを作成
2. SQLContext にデータフレームをテーブルとして登録
3. SQLContext.exexute でクエリを実行
```
ctx = pl.SQLContext(films=df)
ctx.execute('SELECT * FROM films LIMIT 5', eager=True)
```

#### 1. SQLContext オブジェクトを作成
polars SQLではまず始めにデータフレームとクエリ内のテーブルを紐づけるために、[polars.SQLContext](https://pola-rs.github.io/polars/docs/python/version/0.18/reference/sql.html) オブジェクトを作成します。

#### 2. SQLContext にデータフレームをテーブルとして登録

基本的にはSQLContextの初期化時にテーブルを登録するのが楽でおすすめなのですが、[Python Polars 0.17.13](https://github.com/pola-rs/polars/releases/tag/py-0.17.13) で追加された機能で、それ以前のバージョンでは以下のSQLContextのメソッドで登録する必要があるため注意が必要です。
- [polars.SQLContext.registor](https://pola-rs.github.io/polars/docs/python/version/0.18/reference/api/polars.SQLContext.register.html)
- [polars.SQLContext.register_globals](https://pola-rs.github.io/polars/docs/python/version/0.18/reference/api/polars.SQLContext.register_globals.html)
- [polars.SQLContext.register_many](https://pola-rs.github.io/polars/docs/python/version/0.18/reference/api/polars.SQLContext.register_many.html)

#### 3. SQLContext.exexute でクエリを実行
その後、`execute` メソッドで所望のクエリをデータフレームに対して実行できます。
この時 `eager=True` を指定するとDataFrame、デフォルトでは遅延評価を行える[LazyFrame](https://pola-rs.github.io/polars/py-polars/html/reference/lazyframe/index.html)が返されます。

## polars SQLで利用可能な構文
以下は個人的によく使うSQL構文がサポートされているかどうかをまとめています。（[Python Polars 0.19.19](https://github.com/pola-rs/polars/releases/tag/py-0.19.19) 時点）
polarsでは日々アプデが行われていてpolars SQLについても頻繁に機能追加が行われているので、最新のサポート対象を知りたい場合は[リリースノート](https://github.com/pola-rs/polars/releases?page=1)を読むことをお勧めします。

##### 利用可
- 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
```python
%%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
```python
%%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
```python
%%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
```python
%%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
```python
%%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
```python
%%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
```python
%%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
```python
%%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)
```