Pythonでリストをgroupbyするときの注意点

Pythonでは標準関数として，リストの要素を特定のキーでグループ化するgroupby 関数がありますが，リスト内で同じキーが連続している場合に1つのグループとして分類されるので少し注意が必要でした．

SQLのGROUP BY句の動作とは異なる

itertools.groupby(iterable, key=None) は第一引数にリストとkeyを指定することでkeyでリストの要素をグループ化できます．
例えば，要素が辞書型のリストを特定のフィールドでグループ化したいときには以下のように簡単にキーとグループを返すイテレータを得られます

l = [
  {"group": 1, "id": "A"},
  {"group": 1, "id": "B"},
  {"group": 2, "id": "C"},
]

for key, values in groupby(l, key=lambda x: x["group"]):
	print(list(values))

# OUTPUT:
# [{"group": 1, "id": "A"}, {"group": 1, "id": "B"}]
# [{"group": 2, "id": "C"}]

ただここで注意が必要なのが，key 関数の値が変わるたびに休止または新しいグループを生成する点です．以下の例のリストのように group=1 の要素が連続してない場合，キーが同じでも id が A, B のグループとDのグループに分かれてしまいます．

l = [
  {"group": 1, "id": "A"},
  {"group": 1, "id": "B"},
  {"group": 2, "id": "C"},
  {"group": 1, "id": "D"},
]

for key, values in groupby(l, key=lambda x: x["group"]):
	print(list(values))

# OUTPUT:
# [{"group": 1, "id": "A"}, {"group": 1, "id": "B"}]
# [{"group": 2, "id": "C"}]
# [{"group": 1, "id": "D"}]

公式のドキュメントにも

groupby() の操作は Unix の uniq フィルターと似ています。 key 関数の値が変わるたびに休止または新しいグループを生成します (このために通常同じ key 関数でソートしておく必要があるのです)。この動作は SQL の入力順に関係なく共通の要素を集約する GROUP BY とは違います。

https://docs.python.org/ja/3.11/library/itertools.html#itertools.groupby

と書いてあるので，groupbyを使う際は先にソートするようにしましょう．

ll = [
  {"group": 1, "id": "A"},
  {"group": 1, "id": "B"},
  {"group": 2, "id": "C"},
  {"group": 1, "id": "D"},
]

for key, values in groupby(sorted(l, key=lambda x: x["group"]), key=lambda x: x["group"]):
	print(list(values))

# OUTPUT:
# [{"group": 1, "id": "A"}, {"group": 1, "id": "B"}, {"group": 1, "id": "D"}]
# [{"group": 2, "id": "C"}]

終わりに

イテレータベースの挙動なので納得できる反面，知らないとバグを発生捺せかねないので，なんだかな～という思いです．おそらくそこまで出番が多い関数でもないので問題になることもそうないのでしょう．