Joins and aggregations at scale

6 h
requires SPK-05

From your world. Redshift DISTSTYLE ALL ≈ broadcast; distribution keys and skew

Fix a skewed aggregation with salting

One company owns most of the order-item rows. So a plain groupBy("company_id").sum("line_amount") sends all rows of that company to a single shuffle partition. You get one straggler task while the other tasks are idle. The fix for a skewed aggregation that preserves the result is salting. Split the rows of each key across N sub-groups. Aggregate partially. Then combine the partial results into one row for each key.

The purpose of salting is that it must not change the answer. A salt that drops rows, duplicates rows, or does not re-aggregate is a bug. The test compares your salted total against a plain, unsalted DuckDB oracle.

Task

Implement salted_totals(order_items, num_salts). It returns the total line_amount for each company. Use a two-phase salted aggregation:

Add a salt in [0, num_salts) to each row.
Partial-aggregate by (company_id, salt).
Final-aggregate by company_id to one row for each company.

Return the columns company_id, total.

Constraints

The salt can be deterministic (pmod(hash(col), num_salts)) or random. But the final total must equal the unsalted sum for every company.
Return exactly one output row for each company. num_salts >= 1.

How it is graded

Data (DuckDB oracle): SELECT company_id, sum(line_amount) FROM order_items GROUP BY company_id. Your salted result must match exactly (with a float tolerance).
Invariant: the test calls your function with num_salts=1 and num_salts=8. The totals for each company must be identical (salting does not change the result).
Structure: exactly one row for each company. The salt is not an output column.

Starter code

solution.py

"""SPK-06 / 02-salt-skewed-aggregation — starter. See README.md."""


def salted_totals(order_items, num_salts):
    # TODO: two-phase salted aggregation.
    #   1) add salt in [0, num_salts)
    #   2) partial sum by (company_id, salt)
    #   3) final sum by company_id -> one row per company, columns company_id, total
    raise NotImplementedError

Hints

Hint 1

Two-phase salted aggregation: add a salt column salt = pmod(hash(...), N) (or a random int 0..N-1), group by (key, salt) and partial-sum, then group by key and sum the partial sums. The two phases must give the SAME total as a plain groupBy.

Hint 2

The salt must spread the hot key across partitions but NOT change the result: never drop or duplicate rows, and always re-aggregate the partials back to one row per key.

Hint 3

Return columns company_id, total (sum of line_amount per company). The salt is an internal detail; the output has one row per company. Use num_salts salts (parameter, default given).

Run it locally

uv run de start SPK-06/02-salt-skewed-aggregation
# edit workspace/SPK-06/02-salt-skewed-aggregation/
uv run de test