無から素数列挙 on Athena (Presto)

# 動機

Athenaを扱うスクリプトを書いていたときに、クエリキャンセル時のテストを行う必要がありました。
ただ、Athenaは検索が早いため、処理に時間かかるクエリを実データに対して投げると、それなりにデータスキャンが行われ、それなりに課金が発生する可能性があります。

この理由から、`Data scanned: 0 KB` で任意の長さの時間がかかりそうな処理を書けるようになろうと考えました。
ついでに自分の理解度や表現力でどのような処理ができるか試すため、プログラミングの定番の素数列挙をテーマにしました。

# 完成形

```sql
WITH seq AS (
  SELECT
    *
  FROM
    UNNEST (
      sequence(2,5)
    ) AS _(num)
),
lower_set AS (
  SELECT
    *
  FROM
    seq
  CROSS JOIN UNNEST (
    sequence(1, num)
  ) AS _(lower_num)
),
cnt AS (
  SELECT
    num,
    COUNT(*) AS divisible_count
  FROM
    lower_set
  WHERE
    num % lower_num = 0
  GROUP BY
    num
),
answer AS (
  SELECT
    num AS prime
  FROM
    cnt
  WHERE
    divisible_count = 2
  ORDER BY
    num ASC
)

SELECT * FROM answer
```

# 解説

### `seq`テーブル

```sql
WITH seq AS (
  SELECT
    *
  FROM
    UNNEST (
      sequence(2,5)
    ) AS _(num)
),
```

`sequence`関数を使って連番の配列を生成します。
生成した配列を`UNNEST`句を使って展開し、`num`カラムに格納したテーブルのように扱います。

|num|
|:-:|
|2|
|3|
|4|
|5|

### `lower_set`テーブル

```sql
lower_set AS (
  SELECT
    *
  FROM
    seq
  CROSS JOIN UNNEST (
    sequence(1, num)
  ) AS _(lower_num)
),
```

`seq`テーブルで生成した数字に対して、それより小さい値を全て組み合わせたテーブルを作ります。
`UNNEST`句は`JOIN`とセットで使うことで、行ごとの値を参照した配列を作ることができます。
今回は`sequence(1, num)`なので、`num=2`の行には`ARRAY[1,2]`をUNNESTしたものが、`num=3`の行には`ARRAY[1,2,3]`をUNNESTしたものが`CROSS JOIN`されることになります。

|num|lower_num|
|:-:|:-:|
|2|1|
|2|2|
|3|1|
|3|2|
|3|3|
|4|1|
|4|2|
|4|3|
|4|4|
|5|1|
|5|2|
|5|3|
|5|4|
|5|5|


### `cnt`テーブル

```sql
cnt AS (
  SELECT
    num,
    COUNT(*) AS divisible_count
  FROM
    lower_set
  WHERE
    num % lower_num = 0
  GROUP BY
    num
),
```

`lower_set`テーブルで生成した`num >= lower_num, lower_num >= 1`となる`(num,lower_num)`の全組について、
`num`が`lower_num`で割り切れる行の数を`num`ごとに集計します。

|num|divisible_count|
|:-:|:-:|
|2|2|
|3|2|
|4|3|
|5|2|

## `answer`テーブル

```sql

answer AS (
  SELECT
    num AS prime
  FROM
    cnt
  WHERE
    divisible_count = 2
  ORDER BY
    num ASC
)
```

`cnt`テーブルから素数を抽出します。
素数であればその数字自体と1でのみ割り切れるので、`divisible_count = 2`になっているはずです。

|prime|
|:-:|
|2|
|3|
|5|

これで列挙完了です。

# WITHの使いどころ

この例のように`WITH`を使って数珠繋ぎに処理を進めていくことで、行数やメモリ消費と引き換えに複数の処理をシンプルに保つことができます。

# 計算量の改善

n以下の素数を列挙するとき、`lower_set`テーブルにて`(num,lower_num)`の組が`num^2`通りあるので`O(n^2)`の計算量になってしまいます。
`sqrt(num) < lower_num`となる組が割り切れる場合、必ず`sqrt(num) > lower_num`となる割り切れる組が存在するので、
この対称性を利用してクエリを一部見直すことで`O(n^(3/2))`に減らせます。

```sql
lower_set AS (
  SELECT
    *
  FROM
    seq
  CROSS JOIN UNNEST (
    sequence(1, cast( sqrt(num) AS integer ) )
  ) AS _(lower_num)
),
```

```sql
answer AS (
  SELECT
    num AS prime
  FROM
    cnt
  WHERE
    divisible_count = 1
  ORDER BY
    num ASC
)
```

素数列挙の計算量はもう少し少なくできるはずですが、現状良い実装が思いつかなかったので、
これより少ない計算量の実装が出来た人がいればお知らせください。

# 時間計測

`O(n^(3/2))`の方のクエリを使って、nまでの素数を列挙する処理時間を計測しました。

|n|平均(s)|1回目(s)|2回目(s)|3回目(s)|
|:-:|:-:|:-:|:-:|:-:|
|1000|0.30|0.5|0.21|0.19|
|5000|0.62|0.56|0.71|0.59|
|10000|0.79|0.83|0.79|0.74|
|50000|1.52|1.58|1.5|1.49|
|100000|2.91|2.78|3.02|2.92|
|500000|22.62|23.79|23|21.06|
|1000000|57.41|54.56|51.67|66|

少しRUNNINGにさせる目的ならnを6桁前半くらいにするのが良さそうです。

# 感想

意外にも列挙できることがわかった。
いい教材になりそう。

# 注意

もし素数列を使用する場合は、ソートされているとは限らないので必要があればソートしましょう。

```sql
SELECT * FROM answer ORDER BY num ASC
```