Salt the Hot Merchant

A hard Spark interview practice problem on DataDriven. Write the Spark code and run it against the tables, with instant feedback.

Domain
Spark
Difficulty
hard
Seniority
L4
Asked in
Uber

Problem

transactions holds one payment per row keyed by user_id, and a single internal payout account owns 38% of them, so matching them to users on that key parks the whole account on one task for three days while the other 199 finish in seconds. Without broadcasting users (a downstream 2 GB enrichment step rules that out), spread that account's rows across many tasks and return every account in users that has payments, with its payment count and total amount.

One merchant owns 38% of your rows. Salt or suffer.

Expected output

First 5 of 15 rows
user_idtxn_counttotal_amount
100128605.98
197148814.54
294149003.12
1167127959.42
1264128121.06

Worked solution and explanation

What this really is

This is a hash-partitioning problem dressed up as a payments incident. Every row with the same user_id hashes to the same shuffle partition, so one account holding 38% of the ledger means one task holding 38% of the work, no matter how many partitions you configure. The fix is to invent new distinct key values for that account. Anyone can add a salt to transactions. What separates candidates is remembering the other side: if users is not replicated across the exact same salt range, salted payments find no match and the job 'succeeds' while silently dropping rows for the most important account on the platform.

Step 1: Salt the heavy side

Give every payment a bucket with F.col("transaction_id") % SALT_BUCKETS. Payment ids are unique and roughly sequential, so the hot account's 456M rows land evenly on values 0 to N-1 and hash to about N different partitions, roughly 1/N of the load each.

Step 2: Replicate the light side over the same range

Each row of users must exist once per salt value, or a payment in bucket 5 has nothing to meet. F.explode over an array of 0 to N-1 produces exactly N copies, so every (user_id, salt) pair on the fact side has exactly one partner.

Step 3: Match on the composite key, then fold the salt away

Match on [user_id, salt], then group by user_id alone. Each payment matched exactly one replica, so COUNT(*) and SUM(total_amount) are the same numbers the unsalted query would produce, just computed by N tasks instead of one.

Two-sided salting
SALT_BUCKETS = 8

transactions = spark.table("transactions")
users = spark.table("users")

# Heavy side: every payment gets a bucket in [0, SALT_BUCKETS), derived from its own id
salted_txns = transactions.withColumn("salt", F.col("transaction_id") % SALT_BUCKETS)

# Light side: one copy of each account per bucket, so every salted payment finds its match
salt_values = F.array(*[F.lit(i) for i in range(SALT_BUCKETS)])
users_replicated = users.select("user_id").withColumn("salt", F.explode(salt_values))

matched = salted_txns.join(users_replicated, ["user_id", "salt"], "inner")

result = matched.groupBy("user_id").agg(
    F.count("*").alias("txn_count"),
    F.sum("total_amount").alias("total_amount"),
)

result

The salt changes how rows are routed, never which rows match: the totals equal the unsalted answer.

One-sided salt
transactions gets a salt, users does not. The match on (user_id, salt) finds nothing, or someone concatenates the salt into the key and only bucket 0 lines up. Rows vanish, disproportionately for the hot account.
Two-sided salt
users is exploded to N copies, one per salt value. Every payment meets exactly one replica, so nothing is dropped and nothing is double counted.

Common follow-up questions