파티셔닝과 셔플 최적화

Spark 성능 문제의 절반은 셔플에서 시작된다

Spark 작업이 느릴 때 원인은 다양하다. CPU가 부족할 수도 있고, Executor 메모리가 모자랄 수도 있고, 파일이 너무 작아 스케줄링 오버헤드가 커질 수도 있다. 하지만 실무에서 가장 자주 비용을 폭발시키는 지점은 셔플(Shuffle)이다.

셔플은 한 파티션에 있던 데이터를 다른 파티션, 보통 다른 Executor로 다시 나누어 보내는 과정이다. groupBy, join, distinct, orderBy, repartition, 윈도우 집계처럼 “같은 키의 데이터가 한곳에 모여야 하는” 연산에서 발생한다. 이때 Spark는 네트워크로 데이터를 옮기고, 중간 shuffle file을 로컬 디스크에 쓰고, 다음 stage에서 다시 읽는다. 그래서 셔플은 CPU, 네트워크, 디스크 I/O, 메모리를 동시에 건드린다.

성능 튜닝의 핵심은 셔플을 무조건 없애는 것이 아니다. 분산 처리에서 셔플은 필요한 순간이 있다. 중요한 판단은 세 가지다.

1. 불필요한 셔플을 만들지 않는 것
2. 필요한 셔플의 파티션 크기를 적정하게 유지하는 것
3. 데이터 스큐(skew) 때문에 일부 태스크만 오래 끌지 않게 하는 것

파티션은 병렬 처리 단위다

Spark의 파티션은 “데이터 조각”이면서 동시에 “태스크 실행 단위”다. 하나의 stage에서 파티션이 1,000개라면 대체로 태스크도 1,000개 생긴다. 클러스터 전체 코어 수가 200개라면 200개 정도의 태스크가 동시에 돌고, 나머지는 대기한다.

파티션 수가 너무 적으면 코어가 놀고, 각 태스크가 처리할 데이터가 커져 메모리 spill이 늘어난다. 반대로 너무 많으면 태스크 스케줄링 오버헤드가 커지고, 작은 output file이 많이 생긴다. 즉 파티션 튜닝은 “크게 나눌까, 작게 나눌까”의 문제가 아니라 병렬성과 오버헤드의 균형 문제다.

Spark SQL/DataFrame에서 자주 만지는 설정은 다음과 같다.

기본값 200은 작은 개발 환경에는 충분할 수 있지만, 대규모 클러스터에서는 너무 작을 수 있다. 예를 들어 1,000코어 클러스터에서 셔플 파티션이 200개면 한 stage에서 동시에 200개 태스크만 실행된다. 반대로 8코어 개발 환경에서 10,000개 파티션을 만들면 작업보다 스케줄링이 더 바빠진다.

Narrow dependency와 Wide dependency

Spark 실행 계획을 볼 때 가장 먼저 구분해야 하는 것은 narrow dependency와 wide dependency다.

• Narrow dependency: 부모 파티션 하나가 자식 파티션 하나 또는 소수에만 필요하다. map, filter, select, withColumn 같은 연산이 대표적이다. 보통 셔플이 없다.
• Wide dependency: 여러 부모 파티션의 데이터가 여러 자식 파티션으로 재분배된다. groupByKey, reduceByKey, join, distinct, orderBy, repartition 등이 대표적이다. 보통 stage 경계와 셔플을 만든다.

Spark UI에서 stage가 갑자기 나뉘고 Shuffle Read, Shuffle Write가 크게 보이면 wide dependency가 생긴 것이다. 이때부터는 코드 한 줄의 의미보다 물리 실행 계획이 더 중요하다.

repartition과 coalesce를 구분해야 한다

파티션 수를 조정할 때 가장 많이 쓰는 API는 repartition과 coalesce다. 이름은 비슷하지만 비용은 다르다.

# 전체 데이터를 다시 섞어 400개 파티션으로 균등하게 분배한다.
df2 = df.repartition(400, "customer_id")

# 파티션 수를 줄일 때 주로 사용한다. 기본적으로 full shuffle을 피한다.
df3 = df.coalesce(40)

주의할 점은 coalesce가 항상 빠른 선택은 아니라는 것이다. 10,000개 파티션을 10개로 줄이면 셔플은 피할 수 있지만, 10개 태스크가 너무 많은 데이터를 떠안아 느려질 수 있다. 반대로 repartition은 비싸지만 데이터를 균등하게 섞어 뒤 stage의 병목을 줄일 수 있다.

실무 기준은 단순하다. 줄이기만 하면 coalesce, 균등하게 다시 나눠야 하면 repartition이다.

셔플 파티션 수를 어떻게 잡을까

정답 숫자는 없다. 데이터 크기, 압축률, 연산 종류, Executor 메모리, 클러스터 코어 수에 따라 달라진다. 그래도 출발점은 있다.

1. 전체 shuffle read/write 크기를 Spark UI에서 확인한다.
2. 한 태스크가 처리할 목표 크기를 대략 64MB~256MB 범위로 둔다.
3. shuffle_data_size / target_partition_size로 필요한 파티션 수를 계산한다.
4. 클러스터 전체 코어 수보다 너무 작지 않게 맞춘다.
5. AQE가 켜져 있다면 초기값은 넉넉히 주고, post-shuffle coalesce가 줄이게 한다.

예를 들어 어떤 join stage의 shuffle write가 800GB이고 목표 파티션 크기를 128MB로 둔다면 계산상 파티션 수는 약 6,400개다.

800GB / 128MB ≈ 6,400 partitions

이 값은 “반드시 이렇게 하라”가 아니라 관측을 시작하는 기준이다. 실제로는 압축된 파일 크기와 메모리상 row 크기가 다르고, join이나 aggregation 중간에 데이터가 부풀 수 있다. 그래서 최종 판단은 Spark UI의 stage summary, spill size, task duration 분포를 보고 해야 한다.

spark.conf.set("spark.sql.shuffle.partitions", "6400")
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128MB")

대규모 작업에서는 spark.sql.shuffle.partitions=200을 그대로 두는 것보다, 클러스터 규모와 shuffle 데이터 크기에 맞춰 초기값을 크게 잡고 AQE가 줄이도록 하는 편이 안전한 경우가 많다.

AQE는 튜닝을 없애지 않고, 피드백 루프를 짧게 만든다

Adaptive Query Execution(AQE)은 Spark SQL이 런타임 통계를 보고 실행 계획을 조정하는 기능이다. Spark 3.2부터 기본 활성화되어 있다. AQE가 특히 셔플 튜닝에 도움 되는 지점은 세 가지다.

• Post-shuffle partition coalescing: 너무 작게 쪼개진 셔플 파티션을 런타임 통계 기반으로 합친다.
• Skew join optimization: 특정 파티션만 비정상적으로 큰 경우 쪼개서 처리한다.
• Join strategy conversion: 실행 중 실제 크기를 보니 한쪽이 작으면 sort-merge join을 broadcast hash join으로 바꿀 수 있다.

하지만 AQE가 만능은 아니다. 초기 파티션 수가 너무 작으면 이미 큰 파티션이 만들어져 spill이 난 뒤일 수 있다. 통계가 부정확하거나 UDF가 많아 Catalyst가 구조를 파악하지 못하면 최적화 여지도 줄어든다. AQE는 “아무 설정 없이 빠르게 해주는 마법”이 아니라 나쁜 초기 계획을 런타임에 보정하는 안전장치에 가깝다.

데이터 스큐: 평균이 아니라 꼬리를 봐야 한다

셔플 튜닝에서 가장 위험한 착시는 평균 파티션 크기다. 평균이 128MB여도 하나의 key가 데이터의 40%를 차지하면 한 파티션만 수십 GB가 된다. 이때 대부분의 태스크는 끝났는데 하나의 태스크만 오래 남는다. Spark UI에서는 stage의 task duration max가 median보다 압도적으로 크고, 특정 task의 shuffle read/spill이 튀어 보인다.

스큐를 다루는 방법은 원인에 따라 다르다.

Salting은 hot key에 임의 bucket 값을 붙여 임시로 key를 나누는 방식이다. 예를 들어 customer_id='anonymous'가 너무 많다면 (customer_id, salt)로 먼저 나눠 집계한 뒤, 마지막에 다시 customer_id 기준으로 합친다. 단, salting은 결과 정확성을 직접 책임져야 하므로 모든 join/aggregation에 무작정 적용하면 안 된다.

Join에서 셔플을 줄이는 사고방식

Join은 셔플 비용이 가장 크게 터지는 연산 중 하나다. 기본적으로 두 큰 테이블을 같은 join key 기준으로 재분배해야 하기 때문이다. 최적화는 다음 순서로 생각한다.

1. 필터와 컬럼 선택을 join 전에 끝낸다. 필요한 행과 컬럼만 남겨 shuffle payload를 줄인다.
2. 작은 테이블은 broadcast 가능한지 확인한다. Spark 기본 자동 broadcast 임계값은 spark.sql.autoBroadcastJoinThreshold로 제어된다.
3. 통계를 최신화한다. Catalog 통계가 틀리면 Catalyst가 잘못된 join 전략을 고른다. SQL 테이블은 ANALYZE TABLE을 검토한다.
4. 반복 join의 key 분배를 맞춘다. 같은 큰 테이블끼리 반복적으로 같은 key로 join하면 사전 repartition 또는 bucket/partition 설계를 검토한다.
5. 스큐 key를 따로 처리한다. 전체 파티션 수를 늘리는 것만으로 hot key 하나가 해결되지는 않는다.

from pyspark.sql.functions import broadcast

orders_small = orders.select("order_id", "customer_id", "amount")
customers_small = customers.select("customer_id", "segment")

# customers가 충분히 작다면 큰 orders를 customer_id 기준으로 셔플하지 않고 join할 수 있다.
result = orders_small.join(broadcast(customers_small), "customer_id")

Broadcast join도 비용이 있다. 작은 테이블을 모든 Executor에 복사하므로 driver 메모리, executor 메모리, 네트워크를 쓴다. “작으면 무조건 broadcast”가 아니라 실제 크기와 클러스터 메모리를 확인해야 한다.

Output file 수는 마지막 stage에서 결정된다

셔플 파티션 튜닝은 write 결과의 파일 수와도 연결된다. df.write.parquet(...)를 호출하면 보통 마지막 stage의 파티션 수만큼 파일이 생긴다. 그래서 spark.sql.shuffle.partitions를 크게 잡으면 output file도 많아질 수 있다.

마지막 write 직전에 파일 수를 조정하는 패턴은 흔하다.

# 계산 중에는 충분한 병렬성 유지
result = heavy_transform(input_df)

# 최종 결과가 작다면 write 전에 파일 수만 줄인다.
result.coalesce(32).write.mode("overwrite").parquet("s3://bucket/path")

다만 partitioned table에 쓸 때는 조심해야 한다. coalesce(32)는 전체 output task 수를 32개로 줄일 뿐, 각 partition column 값별 파일 수를 균등하게 보장하지 않는다. 날짜별·국가별 데이터량 편차가 크면 어떤 파티션은 여전히 작은 파일이 많거나, 반대로 한 파일이 너무 커질 수 있다. 테이블 포맷이 Delta/Iceberg/Hudi라면 compaction, optimize write, clustering 같은 저장 계층 기능도 함께 봐야 한다.

Spark UI에서 확인할 것

파티셔닝과 셔플 최적화는 코드만 보고 끝낼 수 없다. 반드시 실행 후 Spark UI로 검증해야 한다.

특히 SQL tab에서 Exchange 노드는 셔플 경계를 의미한다. BroadcastExchange는 broadcast join 준비이고, ShuffleExchange는 데이터를 key 또는 partitioning 기준으로 재분배한다. AQE가 켜져 있으면 최종 계획(final plan)과 초기 계획(initial plan)이 다를 수 있으므로 최종 계획을 봐야 한다.

운영 체크리스트

파티션과 셔플 문제를 만났을 때는 다음 순서로 본다.

1. Spark UI에서 가장 오래 걸린 stage를 찾는다.
2. 그 stage의 Shuffle Read, Shuffle Write, Spill, task duration max/median을 본다.
3. SQL physical plan에서 어떤 Exchange가 해당 stage를 만들었는지 찾는다.
4. 불필요한 repartition, distinct, 넓은 groupBy가 있는지 제거한다.
5. join 전 filter/projection을 앞당긴다.
6. 작은 테이블 broadcast 가능성을 검토한다.
7. spark.sql.shuffle.partitions와 AQE 설정을 데이터 크기에 맞게 조정한다.
8. 스큐가 보이면 hot key를 식별하고 salting, 별도 처리, AQE skew join을 검토한다.
9. write 직전 output file 수를 조절하되, 테이블 partition 구조와 compaction 전략까지 함께 확인한다.

한 줄 정리

Spark의 파티션은 병렬 처리 단위이고, 셔플은 데이터를 다시 나누는 비싼 재배치다. 좋은 튜닝은 무작정 파티션 수를 늘리는 것이 아니라 Spark UI의 shuffle 크기, spill, task duration 분포를 보고 불필요한 셔플을 줄이고, 필요한 셔플은 AQE와 적정 partition size로 제어하며, skew key를 따로 다루는 것이다.

References

• https://spark.apache.org/docs/latest/sql-performance-tuning.html
• https://spark.apache.org/docs/latest/rdd-programming-guide.html
• https://spark.apache.org/docs/latest/tuning.html
• https://www.databricks.com/discover/pages/optimize-data-workloads-guide
• https://github.com/apache/spark/blob/master/docs/rdd-programming-guide.md