BigQueryのSQLでユーザーをランダム化比較試験のグループに振り分ける

ビジネスにおいて何らかの介入（ビジネス施策）による効果を測定するために、ユーザーを介入を行う実験グループ（Treatment Group）と介入を行わない対照グループ（Control Group）に分け、ランダム化比較試験（Randomized Controlled Trial）（または AB テスト）を行うことがあります。

https://en.wikipedia.org/wiki/Randomized_controlled_trial
https://en.wikipedia.org/wiki/Treatment_and_control_groups

今回は BigQuery でそのグループの振り分けを行う SQL を紹介します。
コントロールグループを全ユーザーのパーセンテージから指定し、複数のトリートメントグループを作成する仕様を想定しています。

/*
1. コントロールグループ比率とトリートメントグループ数の指定

はじめにコントロールグループの比率とトリートメントグループの数を指定します。
usersテーブルは単なるuser_idのみを持つサンプルのテーブルで、任意のテーブルに変更可能です。
*/
DECLARE control_group_percentage FLOAT64 DEFAULT 0;
DECLARE treatment_group_count INT64 DEFAULT 1;

SET control_group_percentage = 20.0;
SET treatment_group_count = 2;

ASSERT control_group_percentage >= 0 AND control_group_percentage < 100
  AS 'control_group_percentage >= 0 AND control_group_percentage < 100';
ASSERT treatment_group_count > 0 AS 'treatment_group_count > 0';

CREATE TEMPORARY TABLE users AS (
  SELECT '00000000' AS user_id UNION ALL
  SELECT '11111111' AS user_id UNION ALL
  SELECT '22222222' AS user_id UNION ALL
  SELECT '33333333' AS user_id UNION ALL
  SELECT '44444444' AS user_id UNION ALL
  SELECT '55555555' AS user_id UNION ALL
  SELECT '66666666' AS user_id UNION ALL
  SELECT '77777777' AS user_id UNION ALL
  SELECT '88888888' AS user_id UNION ALL
  SELECT '99999999' AS user_id
);

/*
2. ランダムな整数の振り分け

次にRAND関数の結果をuser_idごとに振り分け、一時テーブルを作成します。
このランダムな整数によって、コントロールグループとトリートメントグループへの振り分けが
クエリを実行されるたびにランダムに行われます。
これは後続のSQLのサブクエリには含められません。最終的にコントロールグループとトリートメントグループに
振り分けた結果をUNION ALLで結合しているため、仮にサブクエリにした場合は振り分けが
2回実行されてしまうためです。
*/
CREATE TEMPORARY TABLE users_with_random_number AS (
  SELECT
    user_id
    , RAND() AS random_number
  FROM
    users
);

WITH
  /*
  3. パーセンタイルランクの付与

  続いてPERCENT_RANK関数の結果に100をかけてパーセンタイルランクを付与します。
  PERCENT_RANK関数はその行が全行数に対する何パーセントの行であるかを示す、0.0から1.0までの小数を返します。
  */
  users_with_percent_rank AS (
    SELECT
      user_id
      , random_number
      , PERCENT_RANK() OVER(ORDER BY random_number) * 100 AS percent_rank
    FROM
      users_with_random_number
  )
  /*
  4. コントロールグループフラグの付与

  パーセンタイルランクを元にコントロールグループフラグを振り分けます。
  */
  , users_with_control_group_flag AS (
    SELECT
      user_id
      , random_number
      , IF(percent_rank < control_group_percentage, TRUE, FALSE) AS control_group_flag
    FROM
      users_with_percent_rank
  )
  /*
  5. コントロールグループの抽出

  コントロールグループフラグがTRUEのユーザーのみを抽出します。
  */
  , control_group_users AS (
    SELECT
      user_id
      , NULL AS treatment_group_number
      , control_group_flag
    FROM
      users_with_control_group_flag
    WHERE
      control_group_flag = TRUE
  )
  /*
  5. トリートメントグループの抽出

  コントロールグループフラグがFALSEのユーザーのみを抽出します。
  トリートメントグループ数を元にNTILE関数でトリートメントグループの番号を振り分けます。
  */
  , treatment_group_users AS (
    SELECT
      user_id
      , NTILE(treatment_group_count) OVER(ORDER BY random_number) AS treatment_group_number
      , control_group_flag
    FROM
      users_with_control_group_flag
    WHERE
      control_group_flag = FALSE
  )
  /*
  6. コントロールグループとトリートメントグループの結合

  最後に結合した結果を出力します。
  */
  , treatment_and_control_groups AS (
    SELECT
      user_id
      , treatment_group_number
      , control_group_flag
    FROM
      control_group_users

    UNION ALL

    SELECT
      user_id
      , treatment_group_number
      , control_group_flag
    FROM
      treatment_group_users
  )

SELECT
  user_id
  , treatment_group_number
  , control_group_flag
FROM
  treatment_and_control_groups;

上のクエリの出力では最低限のカラムのみですが、以下の設定で random_number と percent_rank も出力した場合、このような結果になります。

SET control_group_percentage = 20.0;
SET treatment_group_count = 2;

指定の通り、20%がコントロールグループになり、残りの 80%をふたつのトリートメントグループに等分しています。
試しにパラメータを変えて何度か実行してみてください。