Pandas API on Spark: Six Groups, One Model, and a Very Long Night

Parallelism: at most 6, effectively 1.

groupby(...).applyInPandas creates one task per group. Six regions = six concurrent tasks, regardless of how many cores the cluster has. The 62% region is a single task on a single executor scanning ~560 GB - the job is a serial job with five small companions. No amount of cluster size changes this; it is a property of the grouping key.

Problem 1 (scheduling/skew): grouping by a 6-value key concentrates work. The long pole is 6x larger than the average task and ~10x larger than the smallest.

Problem 2 (Python worker memory):

• The 400 MB joblib model is loaded per Python worker, at module import, inside every task. Python workers are separate processes with their own memory, and the JVM sees only memoryOverhead (2 GB) beyond the heap - an OOM in the worker kills the executor, and it is frequently misreported as a Spark memory error.
• On top of the model, each worker holds an Arrow record batch, the deserialised pandas frame, and the output frame simultaneously. With 10,000 rows of a wide schema, those intermediate copies can dwarf the batch itself.
• Net: worker RSS = model + input batch + intermediates, competing against heap + 2 GB overhead.

So the crash is not 'not enough executor heap' - raising heap alone will not fix a Python-worker OOM.

Because batch size controls a fixed cost, not a linear one.

Each Arrow batch pays a roughly constant overhead: serialisation across the JVM/Python boundary, worker context setup, pandas object construction, and per-call model use. Going from 10,000 to 100 rows multiplies the number of batches by ~100 while the per-batch cost stays flat - so the 14-hour runtime is dominated by overhead, not by the actual model inference. The crash disappeared because peak worker memory fell, but the throughput collapsed. That is a symptom swap, not a fix.

How to set it properly:

• Size batches by bytes, not rows: a target of ~32-64 MB per batch keeps Arrow conversion efficient while bounding peak memory; derive the row count from measured row width.
• Shrink the per-worker constant instead of the batch: load the model lazily once per worker (a module-level singleton guarded by a flag, populated on first use) rather than per task/per batch.
• Fix the parallelism, which is the actual cause of the 14 hours: applyInPandas on 6 groups is 6 tasks. Repartition by region plus a salt (e.g. pmod(xxhash64(id), 8)) and use mapInPandas over the salted key, so the 62% region is processed by many tasks instead of one.
• Move to mapInPandas over an iterator of pandas/Arrow batches so memory is streamed per batch and the UDF never holds a whole group.
• If worker OOMs persist, raise memoryOverhead - the worker lives outside the JVM heap, so heap increases do nothing for it.

Rule of thumb: keep batches large enough that fixed costs amortise, and create parallelism through partitioning, not through tiny batches.

Reject it. 900 GB compressed is not the relevant number: .toPandas() materialises an in-memory pandas frame on the driver - uncompressed, object-dtype strings inflated, Python objects with per-value overhead. A 64 GB driver will OOM well before finishing, and if it somehow succeeded, you would have moved an embarrassingly parallel transformation onto a single core. The proposal confuses 'fits on cheap storage' with 'fits in one process heap'.

The narrow, legitimate version of the same idea: .toPandas() is fine for a small aggregate (a few hundred thousand rows) after reducing with Spark; it is never fine for the base table.

Safe path:

1. Keep the work distributed: salt the region key so no task owns 62% of the data; use mapInPandas with a batch iterator.
2. Load the 400 MB model once per worker, not per task - a lazy module-level singleton; in classic mode SparkContext.addFile() (or a cluster-wide mount) makes the file available to workers, so use a cluster-visible path rather than a client-local one.
3. Set ps.options.compute.default_index_type = 'distributed'. The default, distributed-sequence, requires a global ordering to produce sequential indices, which costs an extra job and shuffle when converting to a pandas-on-Spark frame; if you never use the index, that is pure overhead. Choose distributed-sequence only when you genuinely need contiguous indices.
4. Only convert to Pandas at the boundary, and only for bounded results - assert a row count before any .toPandas().
5. Verify Arrow settings globally (spark.sql.execution.arrow.pyspark.enabled) and confirm that failures in the Python worker are surfaced as worker OOMs, not heap OOMs, so the right knob gets turned.