집계와 그룹바이 (Aggregation and Group By)
집계와 그룹바이 (Aggregation and Group By)
Polars의 컨텍스트 중 하나인 group_by는 열을 기준으로 그룹을 나눈 뒤, 그 부분집합(부분 열)에 표현식(expression)을 적용하게 해줘요. 그룹핑 기준이 되는 열의 유일한 값(unique value) 으로 데이터가 나뉘는 구조죠. 이번 섹션에서 이 강력한 기능을 하나씩 살펴볼게요.
먼저 미국 의회 데이터셋을 읽어볼게요.
import polars as pl
url = "hf://datasets/nameexhaustion/polars-docs/legislators-historical.csv"
schema_overrides = {
"first_name": pl.Categorical,
"gender": pl.Categorical,
"type": pl.Categorical,
"state": pl.Categorical,
"party": pl.Categorical,
}
dataset = (
pl.read_csv(url, schema_overrides=schema_overrides)
.with_columns(pl.col("first", "middle", "last").name.suffix("_name"))
.with_columns(pl.col("birthday").str.to_date(strict=False))
)
기본 집계 (Basic aggregations)
집계 값에 여러 표현식을 동시에 적용하기는 아주 쉽습니다. agg 함수 안에 적용하고 싶은 표현식을 나열만 하면 되는데, 집계 개수에 상한은 없고 조합도 자유로워요. 아래 예시에서는 first_name 열을 기준으로 그룹을 나눈 뒤 다음 집계를 적용할게요.
- 그룹의 행 개수 세기 — 즉 데이터셋에서 각 유일한 이름이 몇 명인지 세는 거예요.
- 집계 함수 없이 열을 참조해
gender열의 값들을 리스트로 묶기. - 그룹 안에서
last_name열의 첫 번째 값 가져오기.
집계를 계산한 뒤 바로 결과를 정렬하고 상위 5개 행으로 잘라서 한눈에 보기 좋게 정리해볼게요.
q = (
dataset.lazy()
.group_by("first_name")
.agg(
pl.len(),
pl.col("gender"),
pl.first("last_name"), # `pl.col("last_name").first()`를 줄여 쓴 것
)
.sort("len", descending=True)
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 4)
┌────────────┬──────┬───────────────────┬───────────┐
│ first_name ┆ len ┆ gender ┆ last_name │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ u32 ┆ list[cat] ┆ str │
╞════════════╪══════╪═══════════════════╪═══════════╡
│ John ┆ 4227 ┆ ["M", "M", … "M"] ┆ Walker │
│ William ┆ 3309 ┆ ["M", "M", … "M"] ┆ Few │
│ James ┆ 2414 ┆ ["M", "M", … "M"] ┆ Armstrong │
│ Charles ┆ 1514 ┆ ["M", "M", … "M"] ┆ Carroll │
│ Thomas ┆ 1502 ┆ ["M", "M", … "M"] ┆ Tucker │
└────────────┴──────┴───────────────────┴───────────┘
이 정도면 꽤 간단하죠? 이제 조금 더 나아가볼게요.
조건부 집계 (Conditionals)
각 주(state)의 의원이 "Pro(찬성)"인지 "Anti(반대)"인지 행정부를 기준으로 몇 명인지 알고 싶다고 해볼게요. 이런 질문은 lambda를 쓰거나 데이터프레임을 별도로 손질할 필요 없이 집계 안에서 바로 처리할 수 있어요.
q = (
dataset.lazy()
.group_by("state")
.agg(
(pl.col("party") == "Anti-Administration").sum().alias("anti"),
(pl.col("party") == "Pro-Administration").sum().alias("pro"),
)
.sort("pro", descending=True)
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 3)
┌───────┬──────┬─────┐
│ state ┆ anti ┆ pro │
│ --- ┆ --- ┆ --- │
│ cat ┆ u32 ┆ u32 │
╞═══════╪══════╪═════╡
│ CT ┆ 0 ┆ 5 │
│ NJ ┆ 0 ┆ 3 │
│ DE ┆ 1 ┆ 3 │
│ MD ┆ 0 ┆ 2 │
│ MA ┆ 0 ┆ 2 │
└───────┴──────┴─────┘
그룹 안에서 필터링 (Filtering)
그룹 자체를 필터링할 수도 있어요. 그룹별 평균을 계산하고 싶은데, 그 그룹의 모든 값을 포함하고 싶지 않은 경우를 생각해볼게요. 그리고 그렇게 필터링한 행을 데이터프레임에서 실제로 지우고 싶지도 않아요. 그 행들이 다른 집계를 위해 또 필요할 수 있거든요.
아래 예시가 바로 그 방법을 보여줘요.
Note: 코드를 명확하게 만들기 위해 Python 함수를 정의할 수도 있어요. 이런 함수는 Polars 표현식을 돌려주는 함수라서 실행 비용이 들지 않아요. 쿼리 실행 중에 Series에 커스텀 함수를 적용하는 게 아니거든요. Rust에서도 표현식을 반환하는 함수를 만들 수 있고요.
from datetime import date
def compute_age():
return date.today().year - pl.col("birthday").dt.year()
def avg_age(gender: str) -> pl.Expr:
return (
compute_age()
.filter(pl.col("gender") == gender)
.mean()
.alias(f"avg {gender} age")
)
q = (
dataset.lazy()
.group_by("state")
.agg(
avg_age("M"),
avg_age("F"),
(pl.col("gender") == "M").sum().alias("# male"),
(pl.col("gender") == "F").sum().alias("# female"),
)
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 5)
┌───────┬────────────┬───────────┬────────┬──────────┐
│ state ┆ avg M age ┆ avg F age ┆ # male ┆ # female │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ f64 ┆ f64 ┆ u32 ┆ u32 │
╞═══════╪════════════╪═══════════╪════════╪══════════╡
│ ND ┆ 133.792683 ┆ 87.5 ┆ 164 ┆ 2 │
│ HI ┆ 116.321839 ┆ 88.6 ┆ 87 ┆ 25 │
│ IN ┆ 154.328571 ┆ 93.285714 ┆ 1137 ┆ 28 │
│ PA ┆ 167.008592 ┆ 92.636364 ┆ 3205 ┆ 22 │
│ VI ┆ 96.0 ┆ 81.0 ┆ 12 ┆ 9 │
└───────┴────────────┴───────────┴────────┴──────────┘
평균 나이가 말도 안 되게 높아 보이죠? 그건 이 데이터가 1800년대로 거슬러 올라가는 역사 데이터인데, 데이터셋에 등장하는 모든 사람이 지금도 살아 있다고 가정하고 계산하고 있기 때문이에요.
중첩 그룹바이 (Nested grouping)
앞의 두 질문은 중첩된 group_by로도 처리할 수 있어요. 다만 그렇게 하면 방금 보여드린 유용한 기능 몇 가지를 소개할 수 없었을 거예요. 😉 중첩 group_by를 하려면 그룹핑에 사용할 열을 여러 개 나열하면 됩니다.
먼저 중첩 group_by로 각 주의 "Pro"·"Anti" 행정부 의원 수를 세어볼게요.
q = (
dataset.lazy()
.group_by("state", "party")
.agg(pl.len().alias("count"))
.filter(
(pl.col("party") == "Anti-Administration")
| (pl.col("party") == "Pro-Administration")
)
.sort("count", descending=True)
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 3)
┌───────┬─────────────────────┬───────┐
│ state ┆ party ┆ count │
│ --- ┆ --- ┆ --- │
│ cat ┆ cat ┆ u32 │
╞═══════╪═════════════════════╪═══════╡
│ CT ┆ Pro-Administration ┆ 5 │
│ VA ┆ Anti-Administration ┆ 5 │
│ PA ┆ Anti-Administration ┆ 3 │
│ NJ ┆ Pro-Administration ┆ 3 │
│ DE ┆ Pro-Administration ┆ 3 │
└───────┴─────────────────────┴───────┘
이번에는 중첩 group_by로 주별·성별 의원 평균 나이를 계산해볼게요.
q = (
dataset.lazy()
.group_by("state", "gender")
.agg(
# `avg_age` 함수는 필요 없어요:
compute_age().mean().alias("avg age"),
pl.len().alias("#"),
)
.sort("#", descending=True)
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 4)
┌───────┬────────┬────────────┬──────┐
│ state ┆ gender ┆ avg age ┆ # │
│ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ cat ┆ f64 ┆ u32 │
╞═══════╪════════╪════════════╪══════╡
│ NY ┆ M ┆ 165.204634 ┆ 3965 │
│ PA ┆ M ┆ 167.008592 ┆ 3205 │
│ OH ┆ M ┆ 157.579961 ┆ 2142 │
│ IL ┆ M ┆ 146.069482 ┆ 1895 │
│ CA ┆ M ┆ 115.400464 ┆ 1725 │
└───────┴────────┴────────────┴──────┘
결과는 같지만 데이터의 형태(format)가 다르다는 점을 눈여겨보세요. 상황에 따라 한쪽 형식이 더 알맞을 수 있어요.
정렬 (Sorting)
데이터프레임을 정렬하는 일은, 사실 그룹핑 작업 중 정렬 순서를 관리하기 위해서만 하는 경우가 흔해요. 예를 들어 각 주에서 가장 나이가 많은 의원과 어린 의원의 이름을 얻고 싶다고 해볼게요. 정렬부터 하고 그룹핑하는 방식으로 시작할 수 있어요.
def get_name() -> pl.Expr:
return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")
q = (
dataset.lazy()
.sort("birthday", descending=True)
.group_by("state")
.agg(
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
)
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 3)
┌───────┬───────────────────┬───────────────┐
│ state ┆ youngest ┆ oldest │
│ --- ┆ --- ┆ --- │
│ cat ┆ str ┆ str │
╞═══════╪═══════════════════╪═══════════════╡
│ PI ┆ Carlos Romulo ┆ Pablo Ocampo │
│ PA ┆ Thomas Fitzsimons ┆ Israel Jacobs │
│ IN ┆ Waller Taylor ┆ John Test │
│ MO ┆ Spencer Pettis ┆ Rufus Easton │
│ LA ┆ Jean Destréhan ┆ Thomas Posey │
└───────┴───────────────────┴───────────────┘
그런데 이름을 알파벳 순으로도 정렬하고 싶다면 별도의 정렬 연산이 하나 더 필요해요. 다행히 group_by 컨텍스트 안에서 정렬할 수 있고, 이때 기본 데이터프레임의 정렬 순서는 바뀌지 않아요.
q = (
dataset.lazy()
.sort("birthday", descending=True)
.group_by("state")
.agg(
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
get_name().sort().first().alias("alphabetical_first"),
)
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 4)
┌───────┬───────────────────┬──────────────────────┬────────────────────┐
│ state ┆ youngest ┆ oldest ┆ alphabetical_first │
│ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ str ┆ str ┆ str │
╞═══════╪═══════════════════╪══════════════════════╪════════════════════╡
│ PA ┆ Thomas Fitzsimons ┆ Israel Jacobs ┆ Aaron Kreider │
│ NE ┆ Samuel Daily ┆ Experience Estabrook ┆ Albert Jefferis │
│ MA ┆ William Widgery ┆ Artemas Ward ┆ Aaron Hobart │
│ KS ┆ Jacob LaTurner ┆ James Lane ┆ Abel Wilder │
│ IA ┆ Bernhart Henn ┆ Timothy Davis ┆ Abby Finkenauer │
└───────┴───────────────────┴──────────────────────┴────────────────────┘
심지어 한 열을 다른 열이 만든 순서로 정렬할 수도 있는데, 이 정렬도 group_by 컨텍스트 안에서 동작해요. 위 쿼리를 조금 바꾸면 성(first name)이 맨 앞에 오는 의원이 남성인지 여성인지 확인할 수 있어요.
q = (
dataset.lazy()
.sort("birthday", descending=True)
.group_by("state")
.agg(
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
get_name().sort().first().alias("alphabetical_first"),
pl.col("gender").sort_by(get_name()).first(),
)
.sort("state")
.limit(5)
)
df = q.collect()
print(df)
shape: (5, 5)
┌───────┬──────────────────┬────────────────┬────────────────────┬────────┐
│ state ┆ youngest ┆ oldest ┆ alphabetical_first ┆ gender │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ str ┆ str ┆ str ┆ cat │
╞═══════╪══════════════════╪════════════════╪════════════════════╪════════╡
│ AK ┆ Mary Peltola ┆ Thomas Cale ┆ Anthony Dimond ┆ M │
│ AL ┆ John McKee ┆ Israel Pickens ┆ Albert Goodwyn ┆ M │
│ AR ┆ Archibald Yell ┆ James Bates ┆ Albert Rust ┆ M │
│ AS ┆ Eni Faleomavaega ┆ Fofó Sunia ┆ Eni Faleomavaega ┆ M │
│ AZ ┆ Ben Quayle ┆ Coles Bashford ┆ Ann Kirkpatrick ┆ F │
└───────┴──────────────────┴────────────────┴────────────────────┴────────┘
병렬 처리를 죽이지 마세요 (Do not kill parallelization)
Python 사용자만: 이 섹션은 Python 전용이고 Rust에는 해당하지 않아요. Rust에서는 블록과 클로저(lambda)가 동시에 실행될 수 있고 실제로 그렇게 돼요.
Python은 보통 Rust보다 느려요. "느린" 바이트코드를 실행하는 오버헤드도 있지만, Python은 전역 인터프리터 락(Global Interpreter Lock, GIL) 의 제약 안에서 동작해야 해요. 이 말은 곧 병렬 처리 단계에서 lambda나 커스텀 Python 함수를 적용하면, Polars의 속도가 Python 코드 실행만큼으로 제한되고 여러 스레드가 그 함수를 동시에 실행할 수 없게 된다는 뜻이에요.
Polars는 그룹들에 대해 집계 함수 계산을 병렬화하려고 시도해요. 그래서 가능하면 lambda와 커스텀 Python 함수 사용을 피하라고 권장해요. 대신 Polars 표현식 API의 범위 안에서 머물도록 해보세요. 물론 항상 가능한 건 아니에요. lambda 사용을 더 알고 싶다면 사용자 정의 함수 관련 사용자 가이드 섹션을 참고해주세요.