Partitioning: Intermediate

Both repartition(n) and coalesce(n) change how many partitions a DataFrame has, and they get there by opposite routes. repartition is a wide transformation: the plan gains an Exchange, every row passes through shuffle files, and a new stage starts. coalesce is narrow: each output partition reads several existing partitions where they already sit, with no Exchange and no new stage.

Below, 9 orders arrive in 2 input partitions and repartition(4) spreads them over 4. Each task of the first stage writes its rows to shuffle files, then the 4 tasks of the second stage each fetch a share. With no column given, Spark deals rows out round-robin, so the 4 outputs end up close to equal in size.

The physical plan shows this as Exchange RoundRobinPartitioning(4). The price is that of any shuffle: every row is serialized, written to local disk and fetched over the network by the next stage. In return, repartition can raise the count as well as lower it, and the pieces come out balanced however lopsided the inputs were.

Here the same 9 orders start in 4 partitions and coalesce(2) brings them to 2. No row touches a shuffle file. Each of the 2 tasks reads 2 of the original partitions back to back, inside the same stage that read them from storage.

The plan shows a Coalesce 2 node with no Exchange between it and the scan. That has 2 consequences. coalesce only lowers the count: coalesce(10) on a 4-partition DataFrame leaves 4. And each merged partition is the sum of the inputs it absorbed, so lopsided inputs stay lopsided.

repartition("region") hashes on region, so every row with the same region lands in the same partition. Without a count it uses spark.sql.shuffle.partitions, and repartition(3, "region") sets the count directly. The plan shows Exchange hashpartitioning(region, 3). A groupBy("region") right after needs exactly that arrangement, so the planner adds no second Exchange. The run below has 1 shuffle, and the aggregation happens inside the tasks that received the rows.

Since coalesce adds no stage, it shrinks the stage it lands in, and that stage includes every narrow step above it. Take a read of 400 partitions, a filter that keeps 2% of rows, and coalesce(1) so the output is a single file. The whole stage now runs as 1 task. That task reads all 400 input partitions one after another and filters every row while the rest of the cluster sits idle.

repartition(1) pays for a shuffle and still wins here. The read and filter keep all 400 tasks, and only the 2% of rows that survive cross the Exchange to the single writing task. The rule I use: coalesce when the work sharing its stage is light, repartition when that work is heavy, when the count has to go up, or when the partitions need evening out.

The light case looks like this: a small table, a filter, and an output that needs only a couple of partitions. The merged stage reads few rows, so running it on 2 tasks costs little, and coalesce skips the shuffle that repartition would add. Check both counts with df.rdd.getNumPartitions(), which reports how many partitions the DataFrame has. Then look for Coalesce 2 directly above the filter in explain(), with no Exchange between them.

partitionBy on the writer turns a column's values into directories. Each distinct value gets a folder named column=value, the rows carrying that value go into files inside it, and the column is dropped from the files themselves because the path already records it.

Reading /lake/clicks back, Spark walks the directories, parses each name, and adds click_date as a column again. It infers the type from the names: dates come back as dates and digit strings as integers, so a store code like 00123 written as a partition value reads back as 123. Passing a schema on read, or setting spark.sql.sources.partitionColumnTypeInference.enabled to false, keeps it a string.

Each write task opens 1 file for every distinct partition value among the rows it holds. Upstream partitioning decided which rows each task holds without looking at click_date, so a task holding rows from 30 dates writes 30 files. The ceiling is the task count times the value count.

Take a backfill of 30 days of clicks, 30GB in all, written straight after a groupBy that left 200 partitions:

The cost is 1 shuffle, plus a risk on busy days: a date with 10 times the usual traffic becomes 1 large file written by 1 task, and the stage waits for it. Adding .option("maxRecordsPerFile", 5000000) to the writer makes each task start a new file every 5 million rows, so a heavy date gets several bounded files.

Before choosing a column, look at how rows spread across its values, because each value's rows become 1 directory. Run this on store_sales_raw to see the rows each sale_date directory would hold, biggest day first, and compare the top of the list with the bottom.

(store_sales_raw
    .filter(F.col("sale_date").isNotNull())
    .groupBy("sale_date")
    .count()
    .orderBy(F.col("count").desc(), "sale_date")
    .collect())

Every distinct value costs at least 1 directory and 1 file, however few rows it has. partitionBy("user_id") on a table with 2 million users creates 2 million directories. A read has to list 2 million paths before it scans a row, and the files average a few kilobytes, so opening each one takes longer than reading it.

To bring those 500 files down to 10, change which rows each task holds. repartition("country") hashes on the partition column, so all of a country's rows reach 1 task, and that task writes 1 file into its directory. The shuffle runs before the write, and explain() shows Exchange hashpartitioning(country, 200) above the scan. At most 10 of those 200 tasks hold rows, and the rest write nothing.

When a filter references only partition columns, Spark tests it against directory names before it opens a single file. On events, written with partitionBy("event_date") and holding 730 days, a query for 1 day rules out 729 directories from their names alone. Their files are never opened or scanned. That is partition pruning, the return on the layout work of the previous section.

explain() prints a scan node that sorts the filters by where they apply. These are the parts of it to read:

The scan also prints a PushedFilters list. It belongs to a separate mechanism inside the file format and does not affect which directories are read.

The OR row catches teams out. A row with country = 'US' could sit under any date, so Spark cannot rule a directory out. If the date branch needs to be fast, run it as a query of its own.

Filtering on the column the partition was derived from skips pruning. If event_date was computed from event_ts at write time, a dashboard filtering event_ts to a 6-hour window still reads all 730 directories, because Spark has no way to know the 2 columns agree. Adding the matching event_date predicate leaves the result unchanged and cuts the read to 1 directory:

Pointing the read straight at a directory is the other way to lose it. spark.read.parquet("/lake/events/event_date=2024-03-15") reads the right files but treats that directory as the root, so event_date disappears from the schema. Adding .option("basePath", "/lake/events") keeps the column. Reading the root and filtering reaches the same files and keeps it too.

Collect the WHERE clauses of the queries that read the table, the scheduled jobs as well as the dashboards, and count how often each column appears. For events, a typical dashboard read looks like this:

Distinct values over the table's whole retention become directories, so measure them before writing. Run this against user_events to get the row count beside the distinct count of each candidate. Rows divided by a column's count gives the average rows of the directory each value would get, and a count close to the row count marks a column of near-unique values.

user_events.agg(
    F.count("*").alias("rows"),
    F.countDistinct("event_date").alias("event_date_values"),
    F.countDistinct("country").alias("country_values"),
    F.countDistinct("event_type").alias("event_type_values"),
    F.countDistinct("user_id").alias("user_id_values")
).collect()

Bytes decide whether a directory can hold right-sized files. events lands 40GB a day and keeps 2 years, 29.2TB in total. Divide by each candidate's distinct values and compare the result against a common target of files between 128MB and 1GB:

event_date wins on every test. It serves 85% of reads, and each of its 730 directories has room for about 300 files at 128MB. The 30% of reads that add country still prune to 1 day and filter the country inside it. A second level would split each day into 190 slices, and for most countries the slice is too thin to fill a single file.

Event data usually carries a timestamp, far too fine to partition by. Derive a date at write time and partition by that, then point every reader at the derived column, as the pruning section showed with event_ts.

Every Exchange that groupBy, distinct, orderBy or a count-less repartition(col) adds splits its output into spark.sql.shuffle.partitions pieces, 200 unless you change it. It is 1 value for the whole session, so the same 200 applies to a 2GB job and a 2TB job. Each piece becomes 1 task in the next stage, and when that stage ends in a write, each of those tasks that holds rows writes at least 1 file.

Set the count for the job that needs it and leave other jobs at theirs. At launch, --conf spark.sql.shuffle.partitions=500 on spark-submit scopes it to that application. Inside 1 session that runs several queries, call spark.conf.set before each action, since the plan reads the setting when the action triggers planning:

Where 1 step needs a specific count, repartition(n) or repartition(n, col) states it in the code at that step, where a teammate reading the job will see it.

That runtime merging is covered in a later lesson. It only combines partitions, so a setting that is too small stays too small, and a count passed to repartition(n) is kept as written.

A write leaves 1 file per final-stage task that holds rows, multiplied by the partition values each task holds when partitionBy is on. A small aggregate written at the default count leaves up to 200 files of a few kilobytes each. These settings change that:

coalesce(16) after a groupBy merges the reduce side, so the aggregation itself then runs on 16 tasks. That suits a light aggregate and slows a heavy one, where repartition(16) keeps the aggregation at full width and pays 1 extra shuffle for 16 even files:

In code, the count is set on the session before the action that runs the aggregation. With spark.sql.shuffle.partitions at 8, explain() on a groupBy shows Exchange hashpartitioning(key, 8) between the partial and final aggregates, and the reduce stage is planned as 8 tasks. Set it after the action and that job keeps whatever count was in force when it planned.