윈도우 함수
윈도우 함수 (Window functions)
윈도우 함수는 초능력을 가진 표현식이에요. select 컨텍스트 안에서 그룹 단위의 집계를 수행할 수 있게 해 주죠. 그게 무슨 뜻인지 감을 잡아 볼게요. 먼저 포켓몬 데이터셋을 불러오겠습니다.
출처: 공식문서
import polars as pl
types = (
"Grass Water Fire Normal Ground Electric Psychic Fighting Bug Steel "
"Flying Dragon Dark Ghost Poison Rock Ice Fairy".split()
)
type_enum = pl.Enum(types)
# then let's load some csv data with information about pokemon
pokemon = pl.read_csv(
"docs/assets/data/pokemon.csv",
).cast({"Type 1": type_enum, "Type 2": type_enum})
print(pokemon.head())
shape: (5, 13)
┌─────┬───────────────────────┬────────┬────────┬───┬─────────┬───────┬────────────┬───────────┐
│ # ┆ Name ┆ Type 1 ┆ Type 2 ┆ … ┆ Sp. Def ┆ Speed ┆ Generation ┆ Legendary │
│ --- ┆ --- ┆ --- ┆ --- ┆ ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ enum ┆ enum ┆ ┆ i64 ┆ i64 ┆ i64 ┆ bool │
╞═════╪═══════════════════════╪════════╪════════╪═══╪═════════╪═══════╪════════════╪═══════════╡
│ 1 ┆ Bulbasaur ┆ Grass ┆ Poison ┆ … ┆ 65 ┆ 45 ┆ 1 ┆ false │
│ 2 ┆ Ivysaur ┆ Grass ┆ Poison ┆ … ┆ 80 ┆ 60 ┆ 1 ┆ false │
│ 3 ┆ Venusaur ┆ Grass ┆ Poison ┆ … ┆ 100 ┆ 80 ┆ 1 ┆ false │
│ 3 ┆ VenusaurMega Venusaur ┆ Grass ┆ Poison ┆ … ┆ 120 ┆ 80 ┆ 1 ┆ false │
│ 4 ┆ Charmander ┆ Fire ┆ null ┆ … ┆ 50 ┆ 65 ┆ 1 ┆ false │
└─────┴───────────────────────┴────────┴────────┴───┴─────────┴───────┴────────────┴───────────┘
그룹 단위 연산 (Operations per group)
윈도우 함수는 그룹 안에서 연산을 수행하고 싶을 때 이상적이에요. 예를 들어 "Speed" 열로 포켓몬 순위를 매기고 싶다고 해 봅시다. 그런데 전역 순위가 아니라 "Type 1" 열로 정의된 각 그룹 안에서 속도 기준 순위를 매기고 싶어요. "Speed" 열로 데이터를 순위 매기는 표현식을 작성한 다음 함수 over를 붙여, 이것이 "Type 1" 열의 고유값들에 걸쳐 일어나야 한다고 지정하면 됩니다.
result = pokemon.select(
pl.col("Name", "Type 1"),
pl.col("Speed").rank("dense", descending=True).over("Type 1").alias("Speed rank"),
)
print(result)
shape: (163, 3)
┌───────────────────────┬─────────┬────────────┐
│ Name ┆ Type 1 ┆ Speed rank │
│ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ u32 │
╞═══════════════════════╪═════════╪════════════╡
│ Bulbasaur ┆ Grass ┆ 6 │
│ Ivysaur ┆ Grass ┆ 3 │
│ Venusaur ┆ Grass ┆ 1 │
│ VenusaurMega Venusaur ┆ Grass ┆ 1 │
│ Charmander ┆ Fire ┆ 7 │
│ … ┆ … ┆ … │
│ Moltres ┆ Fire ┆ 5 │
│ Dratini ┆ Dragon ┆ 3 │
│ Dragonair ┆ Dragon ┆ 2 │
│ Dragonite ┆ Dragon ┆ 1 │
│ Mewtwo ┆ Psychic ┆ 2 │
└───────────────────────┴─────────┴────────────┘
이 연산을 시각화하자면, Polars가 "Type 1" 열에서 같은 값을 공유하는 데이터 부분집합을 골라 그 값들에 대해서만 순위 표현식을 계산한다고 상상하면 돼요. 그러면 그 특정 그룹의 결과가 원래 행들에 다시 투영되고, Polars는 존재하는 모든 그룹에 대해 이 작업을 수행합니다.
여기서 눈여겨볼 점이 있어요. "Golbat" 포켓몬 행의 "Speed" 값은 90으로 "Venusaur"의 80보다 크지만, "Golbat"과 "Venusaur"는 "Type 1" 열에서 같은 값을 공유하지 않기 때문에 후자가 1위로 랭크됐어요.
함수 over는 계산을 수행할 그룹을 지정하는 표현식을 임의의 개수로 받을 수 있어요. 위 순위 계산을 "Type 1"과 "Type 2" 열의 조합에 걸쳐 반복하면 더 세밀한 순위를 얻을 수 있습니다.
result = pokemon.select(
pl.col("Name", "Type 1", "Type 2"),
pl.col("Speed")
.rank("dense", descending=True)
.over("Type 1", "Type 2")
.alias("Speed rank"),
)
print(result)
shape: (163, 4)
┌───────────────────────┬─────────┬────────┬────────────┐
│ Name ┆ Type 1 ┆ Type 2 ┆ Speed rank │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ enum ┆ u32 │
╞═══════════════════════╪═════════╪════════╪════════════╡
│ Bulbasaur ┆ Grass ┆ Poison ┆ 6 │
│ Ivysaur ┆ Grass ┆ Poison ┆ 3 │
│ Venusaur ┆ Grass ┆ Poison ┆ 1 │
│ VenusaurMega Venusaur ┆ Grass ┆ Poison ┆ 1 │
│ Charmander ┆ Fire ┆ null ┆ 7 │
│ … ┆ … ┆ … ┆ … │
│ Moltres ┆ Fire ┆ Flying ┆ 2 │
│ Dratini ┆ Dragon ┆ null ┆ 2 │
│ Dragonair ┆ Dragon ┆ null ┆ 1 │
│ Dragonite ┆ Dragon ┆ Flying ┆ 1 │
│ Mewtwo ┆ Psychic ┆ null ┆ 2 │
└───────────────────────┴─────────┴────────┴────────────┘
일반적으로 함수 over로 얻는 결과는 집계 다음에 함수 explode를 호출해서도 얻을 수 있는데, 행의 순서는 달라져요.
result = (
pokemon.group_by("Type 1")
.agg(
pl.col("Name"),
pl.col("Speed").rank("dense", descending=True).alias("Speed rank"),
)
.select(pl.col("Name"), pl.col("Type 1"), pl.col("Speed rank"))
.explode("Name", "Speed rank")
)
print(result)
shape: (163, 3)
┌───────────────────────┬────────┬────────────┐
│ Name ┆ Type 1 ┆ Speed rank │
│ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ u32 │
╞═══════════════════════╪════════╪════════════╡
│ Sandshrew ┆ Ground ┆ 5 │
│ Sandslash ┆ Ground ┆ 3 │
│ Diglett ┆ Ground ┆ 2 │
│ Dugtrio ┆ Ground ┆ 1 │
│ Cubone ┆ Ground ┆ 6 │
│ … ┆ … ┆ … │
│ GyaradosMega Gyarados ┆ Water ┆ 5 │
│ Lapras ┆ Water ┆ 13 │
│ Vaporeon ┆ Water ┆ 11 │
│ Clefairy ┆ Fairy ┆ 2 │
│ Clefable ┆ Fairy ┆ 1 │
└───────────────────────┴────────┴────────────┘
이것이 보여 주듯 보통 group_by와 over는 서로 다른 모양의 결과를 만들어 냅니다.
group_by는 보통 집계에 사용된 그룹 수만큼의 행을 가진 결과 데이터프레임을 만들고;over는 보통 원본과 같은 행 수의 데이터프레임을 만듭니다.
그런데 함수 over가 항상 원본 데이터프레임과 같은 행 수의 결과를 만드는 것은 아니에요. 그것을 다음에서 살펴봅시다.
결과를 데이터프레임 행에 매핑하기 (Mapping results to dataframe rows)
함수 over는 mapping_strategy 파라미터를 받는데, 이 값이 그룹에 대한 표현식 결과를 데이터프레임의 행에 어떻게 다시 매핑할지를 결정합니다.
group_to_rows
기본 동작은 "group_to_rows"예요. 그룹에 대한 표현식의 결과는 그룹과 같은 길이여야 하며, 결과는 그 그룹의 행에 다시 매핑됩니다. 행의 순서가 중요하지 않다면 "explode" 옵션이 더 성능이 좋아요. Polars는 결과 값을 원래 행에 매핑하는 대신, 같은 그룹의 값들이 서로 옆에 놓인 새 데이터프레임을 만들어요. 그 차이를 이해하기 위해 아래 데이터프레임을 보세요.
athletes = pl.DataFrame(
{
"athlete": list("ABCDEF"),
"country": ["PT", "NL", "NL", "PT", "PT", "NL"],
"rank": [6, 1, 5, 4, 2, 3],
}
)
print(athletes)
shape: (6, 3)
┌─────────┬─────────┬──────┐
│ athlete ┆ country ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 │
╞═════════╪═════════╪══════╡
│ A ┆ PT ┆ 6 │
│ B ┆ NL ┆ 1 │
│ C ┆ NL ┆ 5 │
│ D ┆ PT ┆ 4 │
│ E ┆ PT ┆ 2 │
│ F ┆ NL ┆ 3 │
└─────────┴─────────┴──────┘
선수들을 자기 나라 안에서 순위로 정렬할 수 있어요. 그렇게 하면 네덜란드 선수들은 두 번째, 세 번째, 여섯 번째 행에 있었는데 그 자리에 그대로 남습니다. 바뀌는 것은 선수 이름의 순서로, "B", "C", "F"에서 "B", "F", "C"로 바뀌어요.
result = athletes.select(
pl.col("athlete", "rank").sort_by(pl.col("rank")).over(pl.col("country")),
pl.col("country"),
)
print(result)
shape: (6, 3)
┌─────────┬──────┬─────────┐
│ athlete ┆ rank ┆ country │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ str │
╞═════════╪══════╪═════════╡
│ E ┆ 2 ┆ PT │
│ B ┆ 1 ┆ NL │
│ F ┆ 3 ┆ NL │
│ D ┆ 4 ┆ PT │
│ A ┆ 6 ┆ PT │
│ C ┆ 5 ┆ NL │
└─────────┴──────┴─────────┘
explode
mapping_strategy 파라미터를 "explode"로 설정하면 같은 나라의 선수들이 함께 묶이지만, 국가 기준의 최종 행 순서는 같지 않아요. Polars가 각 그룹 행의 위치를 추적할 필요가 없기 때문에 "explode"는 보통 "group_to_rows"보다 빠릅니다. 다만 "explode"는 유지하고 싶은 다른 열들의 순서를 재배열한다는 뜻이라 더 주의가 필요해요. 이 결과를 만드는 코드는 다음과 같습니다.
result = athletes.select(
pl.all()
.sort_by(pl.col("rank"))
.over(pl.col("country"), mapping_strategy="explode"),
)
print(result)
shape: (6, 3)
┌─────────┬─────────┬──────┐
│ athlete ┆ country ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 │
╞═════════╪═════════╪══════╡
│ E ┆ PT ┆ 2 │
│ D ┆ PT ┆ 4 │
│ A ┆ PT ┆ 6 │
│ B ┆ NL ┆ 1 │
│ F ┆ NL ┆ 3 │
│ C ┆ NL ┆ 5 │
└─────────┴─────────┴──────┘
join
mapping_strategy 파라미터의 또 다른 가능한 값은 "join"이에요. 이 값은 결과 값을 리스트로 집계하고 같은 그룹의 모든 행에 그 리스트를 반복합니다.
result = athletes.with_columns(
pl.col("rank").sort().over(pl.col("country"), mapping_strategy="join"),
)
print(result)
shape: (6, 3)
┌─────────┬─────────┬───────────┐
│ athlete ┆ country ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ list[i64] │
╞═════════╪═════════╪═══════════╡
│ A ┆ PT ┆ [2, 4, 6] │
│ B ┆ NL ┆ [1, 3, 5] │
│ C ┆ NL ┆ [1, 3, 5] │
│ D ┆ PT ┆ [2, 4, 6] │
│ E ┆ PT ┆ [2, 4, 6] │
│ F ┆ NL ┆ [1, 3, 5] │
└─────────┴─────────┴───────────┘
윈도우 집계 표현식 (Windowed aggregation expressions)
그룹 값들에 적용된 표현식이 스칼라 값을 만들어 내는 경우, 그 스칼라는 그룹의 행들에 걸쳐 브로드캐스트됩니다.
result = pokemon.select(
pl.col("Name", "Type 1", "Speed"),
pl.col("Speed").mean().over(pl.col("Type 1")).alias("Mean speed in group"),
)
print(result)
shape: (163, 4)
┌───────────────────────┬─────────┬───────┬─────────────────────┐
│ Name ┆ Type 1 ┆ Speed ┆ Mean speed in group │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ i64 ┆ f64 │
╞═══════════════════════╪═════════╪═══════╪═════════════════════╡
│ Bulbasaur ┆ Grass ┆ 45 ┆ 54.230769 │
│ Ivysaur ┆ Grass ┆ 60 ┆ 54.230769 │
│ Venusaur ┆ Grass ┆ 80 ┆ 54.230769 │
│ VenusaurMega Venusaur ┆ Grass ┆ 80 ┆ 54.230769 │
│ Charmander ┆ Fire ┆ 65 ┆ 86.285714 │
│ … ┆ … ┆ … ┆ … │
│ Moltres ┆ Fire ┆ 90 ┆ 86.285714 │
│ Dratini ┆ Dragon ┆ 50 ┆ 66.666667 │
│ Dragonair ┆ Dragon ┆ 70 ┆ 66.666667 │
│ Dragonite ┆ Dragon ┆ 80 ┆ 66.666667 │
│ Mewtwo ┆ Psychic ┆ 130 ┆ 99.25 │
└───────────────────────┴─────────┴───────┴─────────────────────┘
더 많은 예시 (More examples)
연습을 위해 아래 몇 가지 윈도우 함수를 계산해 볼게요.
- 모든 포켓몬을 타입으로 정렬하고;
- 타입별로 처음
3마리 포켓몬을"Type 1"로 선택하고; - 타입 안에서 속도 내림차순으로 포켓몬을 정렬해 처음
3마리를"fastest/group"으로 선택하고; - 타입 안에서 공격력 내림차순으로 포켓몬을 정렬해 처음
3마리를"strongest/group"으로 선택하고; - 타입 안에서 이름순으로 포켓몬을 정렬해 처음
3마리를"sorted_by_alphabet"으로 선택한다.
result = pokemon.sort("Type 1").select(
pl.col("Type 1").head(3).over("Type 1", mapping_strategy="explode"),
pl.col("Name")
.sort_by(pl.col("Speed"), descending=True)
.head(3)
.over("Type 1", mapping_strategy="explode")
.alias("fastest/group"),
pl.col("Name")
.sort_by(pl.col("Attack"), descending=True)
.head(3)
.over("Type 1", mapping_strategy="explode")
.alias("strongest/group"),
pl.col("Name")
.sort()
.head(3)
.over("Type 1", mapping_strategy="explode")
.alias("sorted_by_alphabet"),
)
print(result)
shape: (43, 4)
┌────────┬───────────────────────┬───────────────────────┬─────────────────────────┐
│ Type 1 ┆ fastest/group ┆ strongest/group ┆ sorted_by_alphabet │
│ --- ┆ --- ┆ --- ┆ --- │
│ enum ┆ str ┆ str ┆ str │
╞════════╪═══════════════════════╪═══════════════════════╪═════════════════════════╡
│ Grass ┆ Venusaur ┆ Victreebel ┆ Bellsprout │
│ Grass ┆ VenusaurMega Venusaur ┆ VenusaurMega Venusaur ┆ Bulbasaur │
│ Grass ┆ Victreebel ┆ Exeggutor ┆ Exeggcute │
│ Water ┆ Starmie ┆ GyaradosMega Gyarados ┆ Blastoise │
│ Water ┆ Tentacruel ┆ Kingler ┆ BlastoiseMega Blastoise │
│ … ┆ … ┆ … ┆ … │
│ Rock ┆ Kabutops ┆ Kabutops ┆ Geodude │
│ Ice ┆ Jynx ┆ Articuno ┆ Articuno │
│ Ice ┆ Articuno ┆ Jynx ┆ Jynx │
│ Fairy ┆ Clefable ┆ Clefable ┆ Clefable │
│ Fairy ┆ Clefairy ┆ Clefairy ┆ Clefairy │
└────────┴───────────────────────┴───────────────────────┴─────────────────────────┘