[Python] Pandas보다 빠른 Polars

일봉, 1시간봉 정도의 큰 timeframe을 다룰 땐 데이터 사이즈가 그렇게 크지 않습니다. 한 종목의 일봉은 10년동안 2500개만 생길테니까요. 이정도 데이터면 pandas도 충분합니다.

하지만, 수천 종목의 1분봉/틱 데이터를 pandas로 다루긴 힘듭니다.

 

이에 대한 해결책으로 polars라는 패키지가 있습니다.(polars 설명: https://medium.com/cuenex/pandas-2-0-vs-polars-the-ultimate-battle-a378eb75d6d1)

위 설명을 한 장의 그림으로 요약하면, 아래와 같습니다.

Pandas와 Polars 속도 비교 (출처: Pandas 2.0 vs Polars)

모든 면에서 Polars가 Pandas보다 빠르죠. 때문에, 저는 원래 pandas로 하던 작업을 모두 polars로 갈아타고 있습니다.

속도가 빠른 이유는, rust 코드로 병렬 계산을 하기 때문이라고 합니다. 실제로 polars로 작성한 코드를 실행시키면, CPU 코어를 한번에 다 활용하는 것을 알 수 있습니다.

 

때문에 코어 수가 많은 컴퓨터일 경우 속도가 더 빨라질 것으로 생각되지만.. 코어가 적더라도 확실히 polars가 pandas보다는 빠를 것 같습니다.

pandas로는 부족하다 하시는 분들을 위해, 이번 포스팅에선 polars의 기초 사용 방법을 소개해보고자 합니다.

 

1. 데이터 불러오기

import polars as pl

# csv, parquet등 파일 읽기
data = pl.read_csv('data.csv')

# pandas DataFrame을 polars로 가져오고 싶을 경우
data = pl.from_pandas(pandas_dataframe)

polars는 이미 사용자에게 익숙한 pandas와 사용법이 유사하도록 설계되었습니다. 많은 부분에서 pandas와 비슷한데요, 데이터를 불러오는 부분은 pandas와 거의 동일합니다.

 

위 코드로 불러온 데이터는 아래와 같이 출력됩니다. 출력하는 속도도 pandas보다 훨씬 빠르네요...

 

특정 row를 불러오기 위해 pandas에선 .iloc을 사용했는데요, polars에선 아래처럼 list다루듯이 []를 사용하면 됩니다.

data[100:120]

 

 

2. 컬럼 추가

Polars의 컬럼 추가는 처음엔 어색하고 다루기 복잡한데요, 제 경우 지금은 적응되어 polars가 오히려 더 편한 듯 합니다.

컬럼추가를 위해 with_columns()라는 함수를 사용합니다.

data = data.with_columns(pl.col('marketcap').shift(1).over('ticker').alias('lastcap'))

data = data.with_columns(pl.col('매수대금(개인)(만원)').rolling_sum(window_size=10).shift(1).over('ticker').alias('1d_indiv_buy'),
                    (pl.col('매수대금(개인)(만원)') - pl.col('매도대금(개인)(만원)')).alias('순매수대금(개인)(만원)'),

 

세 줄의 예시 코드를 해석해봅시다.

 

첫 번째 라인을 보시면,

-> pl.col('marketcap')은 data에서 'marketcap' 컬럼을 선택한다는 의미입니다.

-> shift(1)은 이를 한 row씩 shift한다는 의미입니다.(pandas의 shift와 동일)

-> over('ticker')는 이를 ticker를 기준으로 그루핑해 실행한다는 의미입니다. 그냥 shift를 하면 종목 A의 마지막 날짜 가격이 종목 B의 첫 번째 날짜 가격으로 들어가게 될 테니까요.

-> alias('lastcap')은 이렇게 만든 데이터를 'lastcap' 이라는 컬럼으로 추가하라는 의미입니다.

이렇게 해서 lastcap이라는 컬럼을 만들었고, 이 컬럼은 각 종목별 전일 시가총액이 입력됩니다.

 

두 번째 라인을 해석해봅시다

-> pl.col()로 매수대금(개인)(만원) 컬럼을 불러오고, 이를 rolling_sum합니다(pandas의 rolling().sum()과 동일). rolling_sum(window_size=10)이므로, 10개 row의 sum을 사용합니다.

-> 이를 shift(1)로 한 칸씩 밀어주고, over('ticker')를 활용해 ticker별로 그루핑합니다.

-> 이렇게 만든 컬럼을 alias('1d_indiv_buy')로 '1d_indiv_buy' 컬럼에 입력해줍니다.

 

세 번째 라인은 이제 좀 이해가 되시겠죠?

-> pl.col()로 매수대금, 매도대금 두 개의 컬럼을 한번에 불러오고, 이 두 값의 차이를 구합니다.

-> alias()로 이를 새로운 컬럼 값으로 입력합니다.

 

3. 그루핑(groupby)

이번엔 groupby()를 배워봅시다.

사용 예로, 각 종목별 평균거래대금 등을 구할 때 활용됩니다.

data.groupby('ticker').count()

data.groupby('ticker').sum()

위 코드처럼, data.groupby('ticker')로 ticker에 대해 그루핑하고,

count(), sum() 등을 통해 카운트, 합계 등을 계산합니다.

 

4. 조건문

pandas에서 where(~~ ,other=~~)로 조건에 따른 지정값을 넣어주기도 하는데요, polars에선 이를 아래와 같이 구현합니다.

data = data.with_columns(pl.when(pl.col('indiv_net_buy') > 0).then(1)
                    .when(pl.col('indiv_net_buy') < 0).then(-1)
                    .otherwise(0).alias('indiv_net_buy_encode'))

직관적인 코드라 금방 이해가 되실텐데요, 빠르게 해석해보면

indiv_net_buy > 0일땐 1, indiv_net_buy < 0 일땐 -1, 그 외 경우엔 0을 입력한 데이터를 Indiv_net_buy_encode라는 컬럼으로 입력하라고 해석 가능합니다.

 

5. join, concat

pandas의 merge와 동일한 join기능을 알아보겠습니다.

data = data.join(data2,how='left',left_on=['ticker','timestamp'],right_on=['ticker','public_date'])

pandas의 merge와 비슷하죠? data와 data2를 left join으로 합치는 코드입니다. pandas와 유사하므로 쉽게 해석 가능합니다.

 

concat 역시 pandas의 concat과 동일합니다.

data_concat = pl.concat([data1,data2])

 

 

6. 데이터 저장

마지막으로 데이터 저장입니다. 마찬가지로 pandas와 똑같지만, pandas가 읽을때 read_~, 쓸때 to_~를 사용하는 것에 반해, polars는 아마 워딩의 일관성을 지키려고 조금 노력했던 것 같은데요. Polars에선 읽을때 read_~, 쓸때 write_~를 사용합니다.

data.write_parquet('data_written.parquet')

워딩의 차이가 있을 뿐, 기능은 pandas와 동일합니다.

 

 

자, 간단하게 polars의 사용법을 좀 배워봤는데요,

기본적인 사용법만 익히면 pandas보다 더 빠른 속도와 나름대로의 직관성.. 으로 인해 개인적으론 굉장히 만족하고 있습니다.

 

큰 데이터프레임을 다룰 땐 이제 pandas가 아닌 polars를 다뤄보면 어떨까요? (익숙해지다보면 pandas를 잊어버린다는 문제가...)