SQL: 組み合わせにビット演算を使ったGROUP BY

はじめに

こんにちは、データアナリティクスチームのHuangです。

SQL関連について記事を書きました。

データの抽出となるとSQLはほぼ毎日のように使われている言語ではありますが、実際の仕事現場では教科書に載っているようなケースがほとんど存在してません。実際の業務では教科書には載っていないような要件が多くあり、印象に残っているケースがあります。

その中でも印象的なものを要件と解決方法について例を交えながら紹介していきます。

背景

膨大なデータを集計するときに、GROUP BY文はよく使われます。

例えばこのテーブルがあるとします。

とあるチェーン店のデータベースに、ある会員が登録した店舗の情報がある。同じ会員が複数の店舗を登録することができる。

shop_member

shop_id	member_id
A	1
A	2
B	1
B	3
B	4
C	3
C	4
C	5

^{shop_id: 店舗のID、 member_id: 会員のID。}

よくある集計のアプローチのひとつは各店舗に登録した会員数をカウントすることです。

SELECT
  shop_id,
  COUNT(DISTINCT member_id) AS member_count
FROM shop_member
GROUP BY 1
;

結果のテーブルはこのようになります。

shop_id	member_count
A	2
B	3
C	3

もしmember_idが最初からユニークだったら、集計の関数をCOUNT(*)にすることで、さらに短くできます。

要件

各店舗の集計も自然だが、店舗の組み合わせで集計することが場合によってあると思います。

例えば、店舗ごとではなく、都道府県という粒度で集計したい場合、ほかのテーブルから位置情報をJOINすることで、上のセクションと同じ課題になります。

shop_info

shop_id	pref_id	shop_name
A	12	XX店
B	12	YY店
C	13	ZZ店

^{shop_id: 店舗のID、 pref_id: 県のID、 shop_name: 店の名前。}

SELECT
  pref_id,
  COUNT(DISTINCT member_id) AS member_count
FROM shop_member
INNER JOIN shop_info
  USING (shop_id)
GROUP BY 1
;

結果は

pref_id	member_count
12	4
13	3

になる。

しかし、組み合わせを決めるデータが存在していないか、そもそも決めていないかのどちらであれば、すべての組み合わせに対して集計しなければいけないことになる。

余談ですが、実際のケースとしてあったのでこの記事が生まれることになりました。

この要件を整理すると、AとBとCを集計するだけではなくて、AB、AC、BCの組み合わせとABCでも集計しないといけない。

上記の課題を解決した際の望ましい結果としては、以下のデータになる。

これをSQLで実現したい。

shop_ids	member_count
A	2
B	3
C	3
AB	4
BC	3
AC	5
ABC	6

直感的なアプローチ

普通にテーブルを横にマージする場合はJOIN、縦にマージする場合はUNION ALLなどの集合操作を使います。

その考えに沿っていくと、ここは自分と自分をJOINしてから集計するというのが自然である。

例えば２店舗の場合は、

WITH combinations AS (
  SELECT
    CONCAT(t_m.shop_id, t_s.shop_id) AS shop_ids,
    t_m.shop_id AS shop_id_m,
    t_s.shop_id AS shop_id_s
  FROM shop_info AS t_m
  INNER JOIN shop_info AS t_s
    ON t_m.shop_id = t_s.shop_id
  WHERE
    -- 組み合わせがユニークになるように
    t_m.shop_id < t_s.shop_id
)
SELECT
  shop_ids,
  COUNT(DISTINCT member_id) AS member_count
FROM combinations
INNER JOIN shop_member
  ON (
    combinations.shop_id_m = shop_member.shop_id
    OR combinations.shop_id_s = shop_member.shop_id
  )
GROUP BY 1

という風に集計できるが、このアプローチを３つ以上の組み合わせには展開するのは非効率の上に難しい。

そこで汎用的なやり方を考えていきたい。

汎用的なやり方

ここからの内容はビット演算をメインに展開していきます。

二進数やビット演算に詳しくない方もいると思いますので、事前に簡単に読んでいたら後半の部分の理解が早くなります。

求めたいすべての組み合わせを数値で表して、数値のビット(0か1か)はそれぞれの店舗を表す。

上記の例に店舗が３つある場合、1 (001) - 7 (111) の数値ですべての組み合わせを表すことができます。

WITH t_bit AS (
  SELECT
    shop_id,
    -- bitの位置をidのランキングで決める
    -- 毎回の結果が同じになるようにORDER BYが必要
    ROW_NUMBER() OVER (ORDER BY shop_id) AS rank,
    -- シフト操作
    1 << (rank::integer - 1) AS shop_bit
  FROM shop_info
), series AS (
  SELECT
    -- データベースによって使えない場合がある。
    -- 7 を SELECT MAX(shop_bit) * 2 - 1 FROM t_bit にすると
    -- 動的にすることができる
    GENERATE_SERIES(1, 7) AS shop_bits
), combinations AS (
  SELECT
    shop_bits,
    LISTAGG(shop_id) WITHIN GROUP ( ORDER BY shop_id ) AS shop_ids
  FROM series
  INNER JOIN t_bit
    ON series.shop_bits & t_bit.shop_bit > 0
  GROUP BY
    1
)
SELECT
  shop_ids,
  COUNT(DISTINCT member_id) AS member_count
FROM combinations
INNER JOIN t_bit
  ON combinations.shop_bits & t_bit.shop_bit > 0
INNER JOIN shop_member
  ON t_bit.shop_id = shop_member.shop_id
GROUP BY
  1
ORDER BY
  1

上記のクエリを途中経過も含めて解説しながらいくと、

各店舗のビットを決める

SELECT
  shop_id,
  -- bitの位置をidのランキングで決める
  -- 毎回の結果が同じになるようにORDER BYが必要
  ROW_NUMBER() OVER (ORDER BY shop_id) AS rank,
  -- シフト操作
  1 << (rank::integer - 1) AS shop_bit
FROM shop_info

rank列は見ての通り1から各店舗のランキングを計算する。 <<は1(001)を左にシフトする。例えば $1 \ll 1$ は $001$ $を $010$ にする、つまり2にした。表の数値では2倍になった。

t_bit

shop_id	rank	shop_bit
A	1	1(001)
B	2	2(010)
C	3	4(100)

すべての組み合わせを連続の数値で表す

SELECT
  -- データベースによって使えない場合がある。
  -- 7 を SELECT MAX(shop_bit) * 2 - 1 FROM t_bit にすると
  -- 動的にすることができる
  GENERATE_SERIES(1, 7) AS shop_bits

1から7までは 001 - 111 となるので、ビットですべての組み合わせを表すことができた。

series

shop_bits
1
2
3
4
5
6
7

bitの組み合わせから店舗名(shop_id)の組み合わせに

SELECT
  shop_bits,
  LISTAGG(shop_id) WITHIN GROUP ( ORDER BY shop_id ) AS shop_ids
FROM series
INNER JOIN t_bit
  ON series.shop_bits & t_bit.shop_bit > 0
GROUP BY
  1

JOINの条件に注目してほしい。 &は同じ位置のbitで論理的ANDを行う

例えば AC(101) と A(001)、B(010)それぞれで計算すると、

$101 \& 001 = 001 \\ 101 \& 010 = 000$

このように&の左側の組み合わせに右側の店舗が入っているかいないかは、計算が0になっているかいないかで判別することができます。

combinations

shop_bits	shop_ids
1	A
2	B
3	AB
4	C
5	AC
6	BC
7	ABC

^{LISTAGGはCOUNTなどの関数と同じGROUP BYには共存できない(AWS Redshift)}

最後に集計する

上のセクション同様に、作成したcombinationsをJOINして集計する。

最後に

この記事はSQLのGROUP BY集計時に、任意の組み合わせをつくる課題をビット操作で効率的に解決できました。

探索的な課題(どの組み合わせが一番良いとか)にはちょうどいいかもしれないので、機会があればぜひ活用していただけると嬉しいです。