A Straggler Stage in a CDE PySpark Job

AQE skew handling splits an oversized shuffle partition into sub-partitions by key range and, when useful, replicates the matching build-side rows. That fixes partitions that contain many moderately large keys. It cannot fix single-key skew: if one customer_id owns 40M rows, every sub-partition that contains that key still contains all 40M rows. The condition thresholds (advisory size ~256 MB, skew factor 5) are met and irrelevant - splitting has nothing to split on.

There is a second, simpler reason for stragglers: 380 GB across 200 shuffle partitions is ~1.9 GB per partition. AQE's coalescing only merges small partitions; it does not increase partition count. Most tasks are already working on partitions 7× above the 128–256 MB target, so the baseline is wrong before skew is considered.

Changes, in order:

1. Set spark.sql.shuffle.partitions to about 1500–2000 (380 GB ÷ ~256 MB) and let AQE coalesce downward.
2. Salt the hot key: add pmod(hash(customer_id), 16) to the big side and explode the dimension side 16×, then re-aggregate. Alternatively branch the top few customers into a separate job and union the results - simpler to reason about and easier to validate.
3. Confirm skew in the Spark UI's stage task-duration distribution rather than guessing; the max/median ratio should be near 1 when fixed.

As-is: 2M rows × 1.2 KB ≈ 2.4 GB. With 60 executors that is 2.4 GB of heap per executor plus 2.4 GB × 60 ≈ 144 GB of network transfer, and it exceeds the 8 GB broadcast block limit. That is a direct route to the container-kill errors. So: not as-is.

But the dimension is almost certainly being read wider than needed. If the job only needs product_id, category, and brand, a 2M-row projection at ~100 bytes is ~200 MB - small enough to broadcast safely. Conditions for broadcasting:

• After column pruning and filtering (drop discontinued products), the build side is under ~10% of executor memory, and under the 8 GB block limit.
• Raise spark.sql.autoBroadcastJoinThreshold to ~256 MB rather than relying on the 10 MB default.
• The probe side has no pathological single-key skew that broadcasting would merely relocate.
• The dimension is stable across the run, so a stale copy is acceptable.

If those do not hold, keep the sort-merge join and fix partition sizing instead. Broadcasting a 2.4 GB dimension to fix a skew problem would trade a straggler for OOM kills - the worst of both outcomes. Verify with EXPLAIN that BroadcastHashJoin is chosen and that the build side byte count matches the projection estimate.

Sequence by cost and by what actually causes the failure:

1. Column pruning and dimension filtering, then broadcast if it fits (about 200 MB after projection). Removes a 380 GB-class shuffle entirely if it qualifies.
2. Shuffle partition count to ~1500–2000 so the baseline partitions sit near 256 MB.
3. Salt or split the 40M-row customer - the only change that addresses the 2.5–4 hour tasks.
4. Memory overhead: spark.executor.memoryOverhead at 2–3 GB and remove any Python UDFs, which are a common hidden cause of YARN overhead kills.
5. Only then add executors or cores. More executors do not split one hot key and mostly raise cost.

Validation: run against a single day's partition into a scratch Kudu table, then compare COUNT(*), SUM(revenue), and COUNT(DISTINCT customer_id) against the current output - salting must not change results, and a mismatch means the explode/re-aggregate is wrong. Check the Spark UI for max/median task duration below ~3×, spill near zero, and no executor losses. Only after a clean comparison, promote to the scheduled run, and keep the previous Kudu snapshot for one cycle so the dashboard has a fallback if the 06:00 deadline is missed.