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

動機

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

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

完成形

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テーブル

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

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

num
2
3
4
5

lower_setテーブル

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テーブル

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テーブル

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))に減らせます。

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

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桁前半くらいにするのが良さそうです。

感想

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

注意

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

SELECT * FROM answer ORDER BY num ASC