[Python] How to best partition my data with a 32 core EMR instance and make sure I max out the...

I’m optimizing a PySpark pipeline that processes records with a heavily skewed categorical column (category). The data has:

• A few high-frequency categories (e.g., 90% of records fall into 2-3 categories).
• Many low-frequency categories (long tail).

Current approach: Group records into 10 semantic category_groups (predefined business logic). Sub-partition each group into equal-sized chunks (target=1000 rows/sub-partition). Repartition using (category_group, sub_partition) as keys. Current Code:

# Define a target threshold for each group:
target = 1000
# 1) Categorical grouping (e.g., maps "A","B","C" → "group_1")
group_mapping = {"A": "group_1", "B": "group_2", ...}
df = df.withColumn("category_group", F.create_map([F.lit(x) for x in group_mapping.items()])[col("category")])

# 2) Sub-partitioning
w = Window.partitionBy("category_group").orderBy("timestamp")
df = df.withColumn("row_in_group", F.row_number().over(w))
df = df.withColumn("sub_partition", F.floor((col("row_in_group") - 1) / target))

# 3) Repartition
n_partitions = min(math.ceil(df.count() / target), 200)
df = df.repartition(n_partitions, "category_group", "sub_partition")


Problem:

• Works for small datasets (e.g., 10K rows → 10 partitions and 100k rows -> 110 partitions).
• Fails at scale (2M+ rows):
• Fixed target=1000 creates too many partitions (2000+).
• High-frequency categories still cause stragglers.

I would really appreciate if someone can guide me how to best partition my dataset df and how to best do sub-partitions. I think current sub-partition approach looks better where I can salt my groups. But maybe it can be better, especially with an auto-scaling 'target'.

Continue reading...