성능 최적화 및 SQL 활용 🚀 - 자주 쓰는 명령어로 배우는 Polars #6

안녕하세요! '자주 쓰는 명령어로 배우는 Polars' 시리즈의 마지막 글입니다. 이번에는 Polars의 성능 최적화와 SQL 기능 활용에 대해 자세히 알아보도록 하겠습니다.

지금까지 Polars의 기본적인 데이터 처리 방법부터 그룹화, 결합, 재구조화 등 다양한 기능들을 살펴보았습니다. 오늘은 Polars의 성능을 극대화할 수 있는 최적화 기법들을 자세히 알아보면서, LazyFrame을 활용한 쿼리 최적화, 데이터 타입 최적화, 병렬 처리 등 실무에서 꼭 필요한 성능 개선 방법들을 살펴보도록 하겠습니다. 또한 SQL에 익숙한 사용자들을 위해 Polars에서 제공하는 SQL 기능도 함께 알아보겠습니다.

💡 여기서 사용하는 Polars 버전은 1.5입니다. polars 버전이 1 미만이면 예시 코드가 제대로 동작하지 않을 수 있습니다.

먼저 필요한 라이브러리를 임포트 하고 예제 데이터를 불러와 보겠습니다. (아직 설치를 안 하신 분들이 있다면 지난 글을 참고해서 먼저 설치해 주세요 🛠️)

import polars as pl

# 데이터 불러오기
df = pl.read_csv("Pokemon.csv")

1. LazyFrame을 통한 연산 최적화 🔄

LazyFrame은 Polars의 가장 강력한 기능 중 하나입니다. 실제 연산을 즉시 수행하지 않고 필요한 시점까지 지연시켜, 최적의 실행 계획을 수립할 수 있게 해 줍니다.

LazyFrame의 작동 원리

LazyFrame은 다음과 같은 최적화를 자동으로 수행합니다:

1. Filter Pushdown: 필터링 연산을 가능한 한 데이터 읽기 단계로 이동시켜 처리할 데이터량을 줄입니다.
2. Projection Pushdown: 필요한 칼럼만 선택하여 메모리 사용을 최적화합니다.
3. Join Optimization: 여러 테이블을 조인할 때 최적의 순서와 방식을 선택합니다.
4. Predicate Optimization: 복잡한 필터 조건을 단순화하고 최적화합니다.

lazy_query = (
    df.lazy()
    .filter(pl.col("HP") > 100)
    .group_by("Type 1")
    .agg(
        [
            pl.col("Attack").mean().alias("avg_attack"),
            pl.col("Defense").mean().alias("avg_defense"),
        ]
    )
)

print("최적화된 실행 계획:")
print(lazy_query.explain())

print("실행결과:")
print(lazy_query.collect())

실행결과

위 예시에서 lazy_query를 정의할 때는 실행 계획만 수립되고, collect()를 호출하는 시점에 실제 데이터 처리가 이루어집니다.

LazyFrame vs DataFrame 성능 비교

실제로 LazyFrame이 얼마나 효율적인지 살펴보겠습니다:

import time

def measure_performance(query_func):
    start_time = time.time()
    result = query_func()
    end_time = time.time()
    return result, end_time - start_time

eager_result, eager_time = measure_performance(
    lambda: df.filter(pl.col("HP") > 100)
        .group_by("Type 1")
        .agg(pl.col("Attack").mean())
)

lazy_result, lazy_time = measure_performance(
    lambda: df.lazy()
        .filter(pl.col("HP") > 100)
        .group_by("Type 1")
        .agg(pl.col("Attack").mean())
        .collect()
)

print(f"Eager 실행 시간: {eager_time:.4f}초")
print(f"Lazy 실행 시간: {lazy_time:.4f}초")

실행결과

2. 메모리 최적화 기법 💾

대규모 데이터셋을 다룰 때 메모리 관리는 매우 중요합니다. Polars는 다양한 메모리 최적화 기법을 제공합니다.

Scan 함수를 활용한 효율적인 데이터 읽기

Polars는 scan_csv(), scan_parquet(), scan_ipc() 등 다양한 scan 함수를 제공합니다. 이러한 함수들은 지연 실행(lazy evaluation) 방식으로 데이터를 처리하며 다음과 같은 상황에서 특히 유용합니다:

• 여러 개의 대용량 파일 처리
• 복잡한 데이터 처리 파이프라인 최적화
• 메모리 사용량을 최소화해야 하는 환경
• 쿼리 실행 계획 최적화가 필요한 경우

 

💡 아래 예시는 설명을 위한 것으로, Pokemon.csv 데이터셋에서는 실행되지 않습니다.

 

from pathlib import Path

# 다양한 형식의 파일을  읽기
parquet_files = list(Path("data/").glob("*.parquet"))
csv_files = list(Path("data/").glob("*.csv"))

lazy_frames = [
    *[pl.scan_parquet(f) for f in parquet_files],  # parquet 파일
    *[pl.scan_csv(f) for f in csv_files]           # csv 파일
]

# 데이터 처리 파이프라인 구성
query = (
    pl.concat(lazy_frames)
    .filter(
        (pl.col("date").str.contains("2024")) &     # 필터링은 파일 읽기 시점에서 수행
        (pl.col("value") > 1000)
    )
    .select([                                       # 필요한 컬럼만 선택
        "date",
        "category",
        "value",
        "region"
    ])
    .group_by(["category", "region"])
    .agg([
        pl.col("value").mean().alias("avg_value"),
        pl.count().alias("count")
    ])
    .sort("avg_value", descending=True)
)

print("\n처리 결과:")
print(query.collect())

데이터 타입 최적화

적절한 데이터 타입을 선택하면 메모리 사용량을 크게 줄일 수 있습니다:

optimized_df = df.with_columns(
    [
        pl.col("HP").cast(pl.Int16),
        pl.col("Attack").cast(pl.Int16),
        pl.col("Defense").cast(pl.Int16),
        pl.col("Sp. Atk").cast(pl.Int16),
        pl.col("Sp. Def").cast(pl.Int16),
        pl.col("Speed").cast(pl.Int16),
        pl.col("Generation").cast(pl.Int8),
    ]
)

print("메모리 사용량 비교:")
print(f"최적화 전: {df.estimated_size() / 1024:.2f} KB")
print(f"최적화 후: {optimized_df.estimated_size() / 1024:.2f} KB")

실행결과

3. SQL 기능 활용하기 🔍

Polars는 SQL 쿼리를 직접 실행할 수 있는 기능을 제공합니다. 이를 통해 SQL에 익숙한 사용자들도 쉽게 Polars를 활용할 수 있어 진입장벽을 크게 낮출 수 있습니다.

with pl.SQLContext(pokemon=df) as ctx:
    result = ctx.execute(
        """
        SELECT
            "Type 1" as type,
            AVG(HP) as avg_hp,
            COUNT(*) as count
        FROM pokemon
        WHERE Attack > 100
        GROUP BY "Type 1"
        HAVING COUNT(*) > 5
        ORDER BY avg_hp DESC
    """
    )
    print("실행계획: ")
    print(result)
    print("실행결과: ")
    print(result.collect())

실행결과

SQL과 Polars Expr 조합하기

SQL 쿼리 결과는 LazyFrame을 반환하며 이후 polars 문법을 활용하여 추가 처리를 할 수 있습니다:

with pl.SQLContext(pokemon=df) as ctx:
    sql_result = ctx.execute(
        """
        SELECT "Type 1", HP, Attack, Defense
        FROM pokemon
        WHERE HP > 100
    """
    )

    # LazyFrame으로 추가 처리
    final_result = sql_result.with_columns(
        [
            (pl.col("Attack") + pl.col("Defense")).alias("total_power"),
            pl.col("HP").rank().over("Type 1").alias("hp_rank"),
        ]
    ).collect()

    print("실행결과:")
    print(final_result)

🎯 마무리

이번 글에서는 Polars의 성능 최적화와 SQL 기능 활용에 대해 자세히 알아보았습니다. LazyFrame을 활용한 쿼리 최적화를 통해 어떻게 데이터 처리 성능을 향상할 수 있는지 살펴보았고, SQL 친화적인 인터페이스를 통해 기존 SQL 사용자들도 쉽게 Polars를 활용할 수 있다는 것을 확인했습니다.

지금까지 6편에 걸쳐 Polars의 기본부터 고급 기능까지 살펴보았습니다. Polars는 빠른 속도와 직관적인 API, 메모리 효율성 등 많은 장점을 가진 라이브러리입니다. 특히 대용량 데이터 처리에서 그 진가를 발휘하며, pandas 사용자들도 쉽게 전환할 수 있도록 설계되어 있습니다. 이번 기회를 통해 polars를 직접 사용해 보시고 그 매력을 느껴보셨으면 좋겠습니다.

지금까지 긴 글 읽어주셔서 감사합니다 🙇‍♂️

📚 참고자료

• Polars Lazy API
• Polars SQL Context