Grokking the Data Engineering Interview

0% completed

Python for Data Problems

On This Page

The Shipments

Looking Up by Key

Counting

Membership and Differences

Sorting With a Key

What Each Choice Costs

Using This in an Interview

Key Takeaways

Practice Questions

The short script below answers three questions about ten warehouse shipments: how many shipments each carrier handled, which orders have not shipped yet, and which shipments are the heaviest. You can run it in the player.

Python3

. . . .

Output:

per carrier: [('Parcelway', 4), ('Northline', 3), ('Coastal', 3)]
not shipped: ['O-104', 'O-107']
heaviest: [('SH-03', 5100), ('SH-08', 4200), ('SH-05', 3000)]

Each answer took one or two lines, and each used a different built-in structure. The counts came from a counting dictionary, the missing orders from a set difference, and the heaviest shipments from a sort with a key. Picking the structure that matches the question keeps the code short, and it also decides how fast the code runs and how much memory it uses.

Every script in this lesson uses only the Python standard library. Each one ran on Python 3.12, and the output shown is the real output.

The Shipments

A warehouse system exports one record per shipped box. weight_g is the box weight in grams, and day_no is the day it left the warehouse. Each script stores the records as a named tuple, a tuple whose fields can also be read by name, such as s.carrier.

shipment_id	order_id	warehouse	carrier	weight_g	day_no
SH-01	O-101	RNO	Parcelway	2400	1
SH-02	O-102	RNO	Northline	800	1
SH-03	O-103	ATL	Parcelway	5100	1
SH-04	O-101	RNO	Parcelway	1200	2
SH-05	O-105	ATL	Coastal	3000	2
SH-06	O-106	ATL	Northline	800	2
SH-07	O-102	RNO	Northline	650	3
SH-08	O-108	DAL	Coastal	4200	3
SH-09	O-109	DAL	Parcelway	1800	3
SH-10	O-106	DAL	Coastal	1200	4

A second list holds the eight orders customers placed, O-101 to O-108. Three orders shipped in two boxes, two orders have no shipment, and SH-09 belongs to an order that is not on the list.

Looking Up by Key

A dict (dictionary) stores each key with one value, and it finds a key from the key's hash code instead of reading the other entries. A hash code is a number Python computes from the key. The Python FAQ explains that CPython dictionaries are hash tables, so retrieving a key takes "constant time", assuming the keys have different hash values. A dict of three carriers and a dict of three million answer a lookup in about the same time.

The script below prices each shipment from a rate table in cents per kilogram. It also shows what happens when a carrier is missing from the table, and it keeps the last shipment seen for each order.

Python3

. . . .

Output:

SH-01 Parcelway 2400 g -> 216 cents
SH-02 Northline 800 g -> 96 cents
SH-03 Parcelway 5100 g -> 459 cents
get Skyhop: None
KeyError: 'Skyhop'
10 shipments, 7 orders
O-101 -> SH-04 day 2
key order: ['O-101', 'O-102', 'O-103']

SH-01 weighs 2,400 grams at 90 cents per kilogram, so it costs 2400 * 90 // 1000 = 216 cents. A carrier that is not in the table, Skyhop, shows the two ways to read a dict, and each has a cost:

Brackets, RATES[key]: a missing key raises KeyError and stops the job. This is the better choice for a rate table, because a new carrier means the table is out of date, and a guessed price would be billed.
RATES.get(key): a missing key returns None, or a default you pass. This is the better choice when a missing value is normal, such as an optional field. The cost is that a real gap passes without an error.

The last part of the script writes each shipment into latest under its order ID, so a later shipment replaces an earlier one. Ten shipments become 7 keys, and O-101 keeps SH-04 from day 2. This result is correct only because the rows arrive in day order, which is the same "keep the latest row per key" rule that Removing Duplicates applies with ROW_NUMBER. The Python manual also notes that updating a key does not change its position, so O-101 stays first.

Counting

A Counter is a dict that maps each item to the number of times it appeared, and it returns 0 for an item it has never seen. It answers "how many of each" in one pass over the records.

Python3

. . . .

Output:

most_common(2): [('Parcelway', 4), ('Northline', 3)]
Skyhop: 0
ranked: [('Parcelway', 4), ('Coastal', 3), ('Northline', 3)]
top 2 with ties: ['Parcelway', 'Coastal', 'Northline']

The first line is wrong, and nothing reports it. Northline and Coastal both handled 3 shipments, but most_common(2) gave second place to Northline and left Coastal out. The documentation states the rule: "Elements with equal counts are ordered in the order first encountered." Northline appears first in the file (SH-02), Coastal later (SH-05), so a file in another order would give another answer.

The ranked line makes the order independent of the file. Sorting by the pair (-count, name) puts higher counts first and breaks ties by name, so Coastal comes before Northline every time. The last line keeps every carrier whose count is at least the second-highest count, 3, so a tie returns three carriers instead of hiding one.

Membership and Differences

A set is an unordered collection of distinct items, and checking whether an item is in it uses the item's hash code, as a dict does. The manual lists membership testing, removing duplicates, and differences between collections as its common uses. All three appear below.

Python3

. . . .

Output:

10 shipments, 7 distinct orders
ordered, not shipped: ['O-104', 'O-107']
shipped, not ordered: ['O-109']
in both: 6
False True

A set difference keeps the items of the left set that are not in the right set, so its direction matters. ordered - shipped returns O-104 and O-107, which are orders waiting for a box. shipped - ordered returns O-109, which is a shipment with no order, a data problem to report. The 10 shipments cover 7 distinct orders, and 6 of them are on the order list.

Two set differences answer two questions

A set has no order, so the scripts always sort before printing. The order a set prints in can change from one run to the next, because the FAQ notes that a string's hash code depends on "a per-process seed". A test that compares printed sets can therefore fail at random, while a test that compares sorted lists cannot.

Sorting With a Key

sorted() returns a new list ordered by a value computed from each item, called the key, and Python calls the key function once per item. A tuple as the key sorts by its first element, then by the second element among ties, and so on.

Python3

. . . .

Output:

lightest: [('SH-07', 650), ('SH-02', 800), ('SH-06', 800)]
ATL SH-03 5100
ATL SH-05 3000
ATL SH-06 800
DAL SH-08 4200
DAL SH-09 1800
DAL SH-10 1200
RNO SH-01 2400
RNO SH-04 1200
RNO SH-02 800
RNO SH-07 650

SH-02 and SH-06 both weigh 800 grams, and SH-02 comes first because it came first in the input. Python's sort is stable, meaning that items with equal keys keep their input order. Stability is a guarantee, but it repeats whatever order the file had, so a result that must not depend on the file needs a tiebreaker in the key, such as (s.weight_g, s.shipment_id).

The second sort uses the key (s.warehouse, -s.weight_g). Warehouses go from A to Z, and the minus sign puts the heaviest box first inside each warehouse. A minus sign works only for numbers. To sort text in reverse inside an ascending sort, sort twice: first by the inner key in reverse, then by the outer key, and stability keeps the inner order.

What Each Choice Costs

The last script measures the difference between finding 2,000 IDs in a list and in a set of 20,000 IDs. It also times a sort and reports the size of each container. Each time is the fastest of five runs, which reduces the effect of other programs running on the same machine.

Python3

. . . .

Output:

2,000 lookups in 20,000 ids, fastest of 5 runs each
list scan:    217.51 ms
set lookup:     0.16 ms (building the set: 1.31 ms)
sort:           4.01 ms
list object: 173,016 bytes
set object:  2,097,368 bytes

Your times will differ, because they depend on the machine and on what else is running. Checking w in ids on a list compares w with each item from the start until it finds a match. The IDs sit every 10 positions, so the 2,000 checks made about 20 million comparisons. In this run the set answered about 1,400 times faster.

A list compares each item, a set checks one slot

The faster structure has two costs:

Building the set takes time. The 1.31 ms to build it equals about 12 list lookups, so a list is fine when you check only a few items.
The set needs more memory. The set object took about 12 times the bytes of the list object. sys.getsizeof counts only the container itself, not the 20,000 strings that both share.

Question	Structure	What it costs
What is the value for this key?	dict	Memory for the hash table; keys must be hashable (their hash code never changes)
How many of each item?	Counter	Ties come back in input order
Is this item present? Which items are missing?	set	Memory; no order, so sort before printing
In what order? Which are the top N?	`sorted()` with a key	Reads and orders every item, and returns a new list
Read every record once	list	Each membership check scans the list

Using This in an Interview

Meta's published prep guide for its Data Engineer, Product Analytics screen gives Python a 25-minute part of the hour. The guide names "lists, dictionaries, sets, strings, and objects" as the data structures to know, and it says advanced techniques such as dynamic programming are not required.

While you code, say which question each structure answers: "I need to check 2 million shipments against the order IDs, so I build a set of order IDs once." Two follow-up questions are worth preparing:

"Two carriers tie. Which one comes first?" With most_common, the one seen first in the input. Sort by (-count, name) for a fixed order, or keep every item at the cutoff count.
"What if the order list has 500 million IDs?" At the rate measured above, about 105 bytes per ID, 500 million IDs would need about 52 GB before counting the strings. Sort both inputs and compare them in one pass, or load both into the warehouse and find the missing orders with a LEFT JOIN that keeps rows whose match is NULL.

Key Takeaways

Match the structure to the question: a dict for "the value for this key", a Counter for "how many", a set for "is it there" and "what is missing", and sorted() with a key for "in what order".
A dict or set finds a key from its hash code, while a list compares items one by one. On 20,000 IDs, the set answered about 1,400 times faster in one run.
most_common returns tied items in the order they were first seen. Sort by (-count, name) or keep every item at the cutoff count.
A set difference has a direction: ordered - shipped finds waiting orders, and shipped - ordered finds shipments with no order.
Python's sort is stable, so add a tiebreaker to the key when the result must not depend on input order.

The shipments in this lesson arrived as clean Python tuples. Real exports arrive as files with quoted commas, missing fields, and rows that cannot be read at all. The next lesson, Parsing Messy Files, reads such a file and decides what to do with each bad row.

Practice Questions

1. A job checks 2 million shipment rows against a list of 50,000 rejected order IDs. Should the rejected IDs be a list or a set, and why?