BigQuery SQL でレイトレーシング


![bq-results-20201218-214658-cvi5aptd0lxq.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/57758/881642cd-97c9-75b7-f6b9-c7869625a16b.png)
BigQuery (Standard SQL) でレイトレーシングをしてみました。

## レイトレーシングとは

レイトレーシングとは、光の輸送(屈折や反射)を物理シミュレーションして現実的なCG画像を作りだす技術です。
最近では RTX や PS5 など、リアルタイムレイトレーシングが台頭してきています。

[レイ トレーシングとラスタライズの違い | NVIDIA](https://blogs.nvidia.co.jp/2018/03/26/whats-difference-between-ray-tracing-rasterization/)
[![Reflections Real-Time Ray Tracing Demo | Project Spotlight | Unreal Engine](http://img.youtube.com/vi/J3ue35ago3Y/0.jpg)](http://www.youtube.com/watch?v=J3ue35ago3Y)

レイトレーシングではピクセルごとにレイを飛ばして計算するため計算量が膨大になりがちですが、
ピクセルごとに独立に計算することができるので、処理の高速化が期待できます。
それなら BigQuery が得意分野じゃないか？と思い今回の挑戦をしてみました。

## BigQuery とは

超高速でSQLを分散実行し数秒でペタバイト級データに対しても結果が返ってくるデータ分析向けサーバーレス・データウェアハウスです。詳細は以下をごらんください。

[BigQuery とは | Google Cloud](https://cloud.google.com/bigquery/docs/introduction?hl=ja)

## BigQuery による画像生成

テキスト形式で画像を記述する [Portable Any Map](
https://ja.wikipedia.org/wiki/PNM_(%E7%94%BB%E5%83%8F%E3%83%95%E3%82%A9%E3%83%BC%E3%83%9E%E3%83%83%E3%83%88)) というフォーマットがあります。

BigQuery は CSV や JSON などのテキスト形式での出力をサポートしていますので、
今回は PNM として CSV を出力して画像を生成してみます。
画像出力の例としてマンデルブロ集合を描画してみます。

```sql
CREATE TEMPORARY FUNCTION F(s STRUCT<a FLOAT64, b FLOAT64, x FLOAT64, y FLOAT64, c INT64>)
    AS (
        IF(s.a*s.a + s.b*s.b <= 4, 
            STRUCT(s.a*s.a - s.b*s.b - s.x AS a,
                   2*s.a*s.b - s.y AS b, 
                   s.x AS x,
                   s.y AS y,
                   s.c+1 AS c),
        s)
    );
-- Fを4回呼び出す
CREATE TEMPORARY FUNCTION F4(s STRUCT<a FLOAT64, b FLOAT64, x FLOAT64, y FLOAT64, c INT64>)
    AS (F(F(F(F(s)))));

WITH data AS (
SELECT
  x, y,
  -- Fを16回呼び出す
  F4(F4(F4(F4(STRUCT(
    -1 + (j/400)*3 AS a, 
    -1.5 + (i/400)*3 AS b, 
    -1 + (j/400)*3 AS x, 
    -1.5 + (i/400)*3 AS y, 0 AS c))))) AS c
FROM
  -- (1,1) から (400,400) の座標を生成する
  UNNEST(GENERATE_ARRAY(1, 400)) AS i,
  UNNEST(GENERATE_ARRAY(1, 400)) AS j
)
-- PGMとして出力
SELECT "P2 400 400 16 #", -1 AS x, -1 AS y UNION ALL
SELECT FORMAT("%d #", c.c), i, j
FROM 
  data
ORDER BY
  i, j
```

各ピクセル i,j に対し、色(モノクロなので輝度) c を計算するコードになっています。
BigQuery SQL では UDF が使えますが、再帰はできないので不恰好な呼び出し方でループを実現しています。今回のレイトレーシングでもUDFを活用します。

この結果としてCSVとして結果をヘッダなしで出力すると以下のようなファイルになります。

```
P2 400 400 16 #,-1,-1
1 #,1,1
1 #,1,2
1 #,1,3
1 #,1,4
1 #,1,5
1 #,1,6
1 #,1,7
1 #,1,8
...
```

`#` 以降はコメントなのでこれは valid な pnm フォーマットです。
拡張子 pgm で保存すれば、Windows の場合は IfranView、macOS の場合は Preview.app で表示できます。

![bq-results-20201218-221255-nskv8mbyv01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/57758/bef1f965-8aa3-42e0-cf38-a704651c45ae.png)


これで BigQuery で画像を出力できることが確認できました。

## BigQuery によるレイトレーシング

というわけで、BigQueryでレイトレーシングをやってみましょう。
実際のSQLコードは以下のようになります。

```sql
-- Vec3のドット積
CREATE TEMPORARY FUNCTION DOT
  (a STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>,
   b STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>)
  AS (
    a.x*b.x + a.y*b.y + a.z*b.z
  )
;
-- 線形結合 aP + bQ
CREATE TEMPORARY FUNCTION COMB
  (a FLOAT64, p STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>,
   b FLOAT64, q STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>)
  AS (
    STRUCT( a*p.x+b*q.x AS x, a*p.y+b*q.y AS y, a*p.z+b*q.z AS z)
  )
;
-- 単位ベクトルを得る
CREATE TEMPORARY FUNCTION UNIT
  (r STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>)
  AS (
    STRUCT( r.x/SQRT(DOT(r,r)) AS x, r.y/SQRT(DOT(r,r)) AS y, r.z/SQRT(DOT(r,r)) AS z)
  )
;
-- [0,255]の範囲におさめる
CREATE TEMPORARY FUNCTION CLAMP255
  (f FLOAT64)
  AS (
    CASE WHEN f < 0 THEN 0
         WHEN f > 1 THEN 255
         ELSE CAST(ROUND(IFNULL(f, 0) * 255.) AS INT64)
    END
  )
;

-- レイの反射を計算する
CREATE TEMPORARY FUNCTION STEP
  (arg STRUCT<
    -- レイ情報
    rt ARRAY<STRUCT<
        i INT64, -- ピクセル座標
        j INT64, 
        r STRUCT<
          o STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 原点
          d STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 方向
          af STRUCT<x FLOAT64, y FLOAT64, z FLOAT64> -- 色の累積値
        >,
        final STRUCT<x FLOAT64, y FLOAT64, z FLOAT64> -- レイが無限遠方に飛んだときの色
    >>,
    world ARRAY<STRUCT< -- オブジェクト設定
      sc STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 球の中心
      sr FLOAT64, -- 球の半径
      att STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>, -- 色
      mat STRING -- 材質 (diffuse, metal)
    >>,
    sample_num INT64 -- サンプル番号
  >)
  AS ((
    -- (い) レイの衝突を計算
    WITH hitinfo AS (
      SELECT 
        i, j, r, s,
        IF(t IS NOT NULL,
          -- 衝突したので交点情報を返す
          STRUCT(
            COMB(1.0, r.o, t, r.d) AS p,
            COMB(1.0/s.sr, COMB(1.0, r.o, t, r.d), -1.0/s.sr, s.sc) AS n
          ),
          NULL
        ) AS hit
      FROM (
        SELECT
          i, j, r,
          (
            SELECT AS STRUCT
              -- 線分と球面の交点を計算する
              s, t
            FROM (SELECT
                s, (-b - SQRT(b*b - a*c))/a AS t
                FROM (SELECT 
                  s,
                  DOT(r.d, r.d) AS a,
                  DOT(oc, r.d) AS b,
                  DOT(oc, oc) - s.sr*s.sr AS c
                  FROM
                    (SELECT 
                        s, COMB(1.0, r.o, -1.0, s.sc) AS oc
                     FROM UNNEST(arg.world) AS s)
                )
                WHERE b*b - a*c > 0
            )
            WHERE t > 1e-7
            ORDER BY t
            LIMIT 1
          ).*
        FROM
          UNNEST(arg.rt)
        WHERE
          final IS NULL
      )
    )
    SELECT 
      STRUCT(ARRAY(
        SELECT AS STRUCT
          i, j,
          -- 衝突した場合材質に応じてレイを反射させる
          IF(hit IS NULL, NULL,
          CASE s.mat 
            WHEN 'diffuse' THEN
              -- Lambert則(厳密ではない)
              STRUCT(
                hit.p AS o,
                -- arg.sample_num/arg.sample_num はバグっぽい最適化がかかって RAND() が同じ値を返すのに対する workaround
                COMB(1.0, hit.n, 1.0,
                    UNIT(STRUCT(2*RAND() - arg.sample_num/arg.sample_num AS x,
                                2*RAND() - arg.sample_num/arg.sample_num AS y,
                                2*RAND() - arg.sample_num/arg.sample_num AS z))) AS d,
                STRUCT(r.af.x * s.att.x AS x, r.af.y * s.att.y AS y, r.af.z * s.att.z AS z) AS af
              )
            WHEN 'metal' THEN
              -- 鏡面反射
              STRUCT(
                hit.p AS o,
                COMB(1,UNIT(r.d),-2*DOT(UNIT(r.d), hit.n), hit.n) AS d,
                STRUCT(r.af.x * s.att.x AS x, r.af.y * s.att.y AS y, r.af.z * s.att.z AS z) AS af
              )
            -- (屈折は時間がなかったため省略)
          END) AS r, -- (う)
          -- 何にもヒットしなかったので背景を表示
          IF(hit IS NULL,
            (SELECT STRUCT(
              (1-t/2)*r.af.x AS x,
              (1-t*0.3)*r.af.y AS y,
              r.af.z AS z
            )
            FROM (SELECT (1+r.d.y/SQRT(DOT(r.d, r.d)))/2 AS t))
          , NULL) AS final
        FROM hitinfo
        UNION ALL
        -- 無限遠方にあるレイについては衝突判定をしない
        SELECT AS STRUCT i, j, r, final
        FROM UNNEST(arg.rt)
        WHERE final IS NOT NULL
      ) AS rt,
      arg.world AS world,
      arg.sample_num AS sample_num
    )
));

WITH pixels AS (
-- ピクセルを生成
SELECT
  i, j
FROM
  UNNEST(GENERATE_ARRAY(0, 399)) AS j,
  UNNEST(GENERATE_ARRAY(0, 399)) AS i
),
world AS (
  -- オブジェクトの定義
  -- c: 中心, r: 半径, att: 色, mat: 材質
  SELECT
    STRUCT(0.0 AS x, 0.0 AS y, -1.0 AS z) AS sc, 0.5 AS sr,
    STRUCT(0.8 AS x, 0.5 AS y, 0.2 AS z) AS att, "diffuse" AS mat
  UNION ALL SELECT
    STRUCT(1.0 AS x, 0.0 AS y, -1.0 AS z) AS sc, 0.5 AS sr,
    STRUCT(0.8 AS x, 0.7 AS y, 0.9 AS z) AS att, "metal" AS mat
  UNION ALL SELECT
    STRUCT(-1.0 AS x, 0.0 AS y, -1.0 AS z) AS sc, 0.5 AS sr,
    STRUCT(0.6 AS x, 0.5 AS y, 0.9 AS z) AS att, "metal" AS mat
  UNION ALL SELECT
    STRUCT(0.4 AS x, -0.4 AS y, -0.5 AS z) AS sc, 0.1 AS sr,
    STRUCT(0 AS x, 1 AS y, 0.5 AS z) AS att, "diffuse" AS mat
  UNION ALL SELECT
    STRUCT(-0.4 AS x, -0.4 AS y, -0.5 AS z) AS sc, 0.1 AS sr,
    STRUCT(1 AS x, 1 AS y, 1 AS z) AS att, "metal" AS mat
  UNION ALL SELECT
    STRUCT(0.0 AS x, -100.5 AS y, -1.0 AS z) AS sc, 100.0 AS sr,
    STRUCT(0.5 AS x, 0.5 AS y, 0.2 AS z) AS att, "diffuse" AS mat
),
result AS (
  WITH samples AS (
    SELECT 
      -- 8回反射させる (力技)
      -- BigQuery Scripting を使えばもっとスマートにできるが、スキャン料金がかかる
      STEP(STEP(STEP(STEP(STEP(STEP(STEP(STEP(
        STRUCT(
          ARRAY(SELECT AS STRUCT
            i, j,
            -- (あ) カメラからのレイを生成
            STRUCT(
              STRUCT(0.0 AS x, 0.0 AS y, 0.0 AS z) AS o,
              STRUCT(2.0*(i+2*RAND()-s/s)/400-1 AS x, 1.0-2.0*(j+2*RAND()-s/s)/400 AS y, -1.0 AS z) AS d,
              STRUCT(1.0 AS x, 1.0 AS y, 1.0 AS z) AS af
            ) AS r,
            CAST(NULL AS STRUCT<x FLOAT64, y FLOAT64, z FLOAT64>) AS final
          FROM pixels) AS rt,
          ARRAY(SELECT AS STRUCT * FROM world) AS world,
          s AS sample_num
        )
      )))))))).rt AS rt
    FROM
      -- 50サンプル取る
      UNNEST(GENERATE_ARRAY(1,50)) AS s
  )
  SELECT i, j, final
  FROM samples, UNNEST(rt)
)
-- PMX形式で出力 (ヘッダなしでCSV出力すればそのまま読み込めます)
SELECT
  "P3 400 400 255 #", -1 AS i, -1 AS j,
UNION ALL
(
  SELECT
    FORMAT("%d %d %d #",
      -- 各サンプルを加算してガンマ補正
      CLAMP255(SQRT(SUM(final.x)/COUNT(final))),
      CLAMP255(SQRT(SUM(final.y)/COUNT(final))),
      CLAMP255(SQRT(SUM(final.z)/COUNT(final)))), i, j
  FROM result
  GROUP BY i, j -- (え)
)
ORDER BY j, i
```

詳細については割愛しますが、流れとしては普通のレイトレと同様な流れになっています


- I. 各ピクセルごとに以下を S 回 (サンプル数) 実行する
    - 1. カメラから出るレイを生成 (アンチエイリアスのためにランダムに少しずらす) (あ)
    - 2. 各物体との当たり判定を計算する
        - 物体テーブル world と CROSS JOIN して t が最小となる物体を SQL で抽出しています
    - 3. 当たり判定がある場合は、物体の材質に応じて反射の方向を計算する (い)
    - 4. 当たり判定がない場合は、背景の色を参照する
        - final というカラムに格納し、これ以上の当たり判定をしないようにします。。
    - 5. 反射が起きた場合、2 に戻る
        - 繰り返しは STEP という関数を何回も呼びだすことで実現しています
        - コメントにもあるように BigQuery では SCRIPTING というプロシージャ機能にループがあるので、それを使うのもありです。
- II. ピクセルごとにサンプルを集計し、出力
    - GROUP BY で集計しています (え)

## 出力

上のコードを実行して ppm で保存したものが冒頭の画像になります。
なおこのスクリプトの実行には 10 分 24 秒かかってしまいました[^1]。手元の Python でのリファレンス実装では 10秒かからないのでちょっと遅いですね。

改善案としては、そもそも BigQuery では JavaScript で UDF が書けるため、それをレイの計算に利用するという方法があります。おそらくは普通の JavaScript で書くと同じかそれ以上のスピードは出ると予想されるため、数秒で終了するはずです。ただ今回はピュアな SQL で書いてみたいということで JavaScirpt は使いませんでした。[^2]

## 結論

RTX を使おう。

## 参考文献

参考文献としては以下を参照しました。

- https://raytracing.github.io/books/RayTracingInOneWeekend.html
    - レンダリングのアルゴリズムはほとんどこれを利用しています
- http://kagamin.net/hole/simple/index.htm
    - 非再帰のコードの参考にしました

[^1]: ちなみにこのクエリはスキャンバイトが 0B なので無料で実行できます。サンドボックスなら課金不要なので試してみてください。
[^2]: しかし、実はこれ以上複雑にすると `Resources exceeded during query execution: Not enough resources for query planning - too many subqueries or query is too complex.` エラーとなります。屈折の実装は断念しました。