Polars #05

#Window Function #코테 #단골 #SQL

아래와 같은 SQL이 있다.

OVER (PARTITION BY ~ ORDER BY ~) 내용이 있고,

GROUP BY, WITH 등의 내용도 보인다.

이 내용의 핵심은 year 데이터가 3년 연속으로 등장하는 ROW 에서 NAME 을 추출한다.

이걸 Polars 에서 어떻게 만들어 볼 수 있을까...

Window Function 을 이해해야 알 수 있다는 가정하에,

예시 데이터를 간단하게 생성해보았다.

---

# 예제 데이터

data = {

    "name": ["Bob", "Bob", "Bob", "Charlie", "Charlie", "Charlie", "Charlie", "Charlie", "Charlie", "Charlie", "Alice", "Alice", "Alice", "Alice"],

    "year": [2017, 2018, 2019, 2014, 2015, 2016, 2018, 2019, 2020, 2021, 2018, 2019, 2020, 2021]

}

---

데이터에서 name 과 year 를 표현하고 있다.

일단,

해당 데이터(Dict Type)를 DataFrame 으로 표현하고 시작한다.

---

df = pl.DataFrame(data)

이어서,

해당 DataFrame 을 name, year 순으로 SORT 하고난 뒤에

cum_count() 또는 rank() 함수를 통해 새로운 COL 을 생성한다.

df = df.sort(["name", "year"]).with_columns([

    (pl.col("year") - pl.col("year").cum_count().over("name")).alias("grp"),

    (pl.col("year") - pl.col("year").rank().over("name")).alias("grp2")

])

---

grp와 grp2 는 같은 결과를 보여준다.

알아두면 좋은 함수이다.

이제 누가 몇번 연속으로 등장했는지 COUNT 를 해보는 함수를 작성해본다.

---

grouped = df.group_by(["name", "grp"]).agg([

    pl.count("year").alias("consecutive_years"),

    pl.min("year").alias("start_year"),

    pl.max("year").alias("end_year")

])

---

consecutive_years 를 주목해서 보아야 한다.

2000, 2002, 2003 으로 이어진 데이터는 consecutive_years 결과가 2로 나와야 한다.

하지만 2000, 2001, 2002, 2003 은 consecutive_years 결과가 4로 나와야 한다.

이제 3번 이상( '>=' ) 조건에 알맞는 데이터를 FILTER 해보고

---

filtered = grouped.filter(pl.col("consecutive_years") >= 3)

---

 바로 결과를 확인하면

---

result = filtered.select("name").unique().sort("name")

---

아래와 같다.

cum_count().over()

rank().over()

unique()

정도 기억하면 좋을 것 같다.