GCP Interview Questions for Data Engineers, Part 2: The Apache Beam Coding Round

import csv, glob, os, tempfile, logging
from datetime import datetime, timezone
import apache_beam as beam
from apache_beam.transforms import window
logging.getLogger().setLevel(logging.ERROR)

OUT = tempfile.mkdtemp()
def collect(name):                                  # read back what a pipeline wrote, sorted for stable output
    rows = []
    for f in glob.glob(os.path.join(OUT, name + "*")):
        rows += [line.rstrip("\n") for line in open(f, encoding="utf-8")]
    return sorted(rows)

def bookings():
    with open("hotels.csv", newline="", encoding="utf-8") as f:
        yield from csv.DictReader(f)

with beam.Pipeline() as p:                          # no runner given: Beam uses its local runner (Prism here)
    rows = p | "read" >> beam.Create(list(bookings()))
    kept = rows | "not cancelled" >> beam.Filter(lambda r: r["is_canceled"] == "0")

    # Task 1: bookings per market segment (count per key)
    (kept | "segment" >> beam.Map(lambda r: r["market_segment"])
          | "count" >> beam.combiners.Count.PerElement()
          | "fmt1" >> beam.MapTuple(lambda k, n: f"{k}: {n}")
          | "w1" >> beam.io.WriteToText(os.path.join(OUT, "t1")))

    # Task 2: average nightly rate per hotel, with a combiner rather than GroupByKey
    (kept | "rate" >> beam.Map(lambda r: (r["hotel"], float(r["adr"])))
          | "mean" >> beam.combiners.Mean.PerKey()
          | "fmt2" >> beam.MapTuple(lambda k, v: f"{k}: {v:.2f}")
          | "w2" >> beam.io.WriteToText(os.path.join(OUT, "t2")))

    # Task 3: top 3 guest countries per hotel
    (kept | "country" >> beam.Map(lambda r: ((r["hotel"], r["country"]), 1))
          | "per country" >> beam.CombinePerKey(sum)
          | "rekey" >> beam.MapTuple(lambda k, n: (k[0], (n, k[1])))
          | "top3" >> beam.combiners.Top.PerKey(3)
          | "fmt3" >> beam.MapTuple(lambda h, top: f"{h}: " + ", ".join(f"{c} {n}" for n, c in top))
          | "w3" >> beam.io.WriteToText(os.path.join(OUT, "t3")))

    # Task 4: enrich with a small lookup table passed as a side input
    hotel_city = p | "lookup" >> beam.Create([("City Hotel", "Lisbon"), ("Resort Hotel", "Algarve")])
    (kept | "hotel key" >> beam.Map(lambda r: r["hotel"])
          | "enrich" >> beam.Map(lambda h, cities: cities.get(h, "unknown"), cities=beam.pvalue.AsDict(hotel_city))
          | "count city" >> beam.combiners.Count.PerElement()
          | "fmt4" >> beam.MapTuple(lambda k, n: f"{k}: {n}")
          | "w4" >> beam.io.WriteToText(os.path.join(OUT, "t4")))

class LabelWithDay(beam.DoFn):                     # attach the window's start date to each daily total
    def process(self, total, win=beam.DoFn.WindowParam):
        yield f"{datetime.fromtimestamp(win.start.micros / 1e6, timezone.utc):%Y-%m-%d}: {total}"

with beam.Pipeline() as p:
    # Task 5: daily totals in event time, from hourly bike rentals
    def stamped(r):
        # the file gives local hours; they are used as-is, so windows follow the local calendar day
        ts = datetime.strptime(f"{r['dteday']} {int(r['hr']):02d}", "%Y-%m-%d %H").replace(tzinfo=timezone.utc)
        return window.TimestampedValue(int(r["cnt"]), ts.timestamp())
    with open("hour.csv", newline="", encoding="utf-8") as f:
        hours = [r for r in csv.DictReader(f) if r["dteday"].startswith("2012-07-0")]
    (p | "hours" >> beam.Create(hours)
       | "event time" >> beam.Map(stamped)
       | "daily" >> beam.WindowInto(window.FixedWindows(24 * 3600))
       | "sum" >> beam.CombineGlobally(sum).without_defaults()
       | "label" >> beam.ParDo(LabelWithDay())
       | "w5" >> beam.io.WriteToText(os.path.join(OUT, "t5")))

for task, title in [("t1", "T1 non-cancelled bookings per market segment"), ("t2", "T2 mean nightly rate per hotel"),
                    ("t3", "T3 top 3 guest countries per hotel"), ("t4", "T4 bookings by city (side-input lookup)"),
                    ("t5", "T5 daily bike rentals, 1-9 July 2012 (fixed event-time windows)")]:
    print(title)
    for line in collect(task):
        print("  ", line)

with open("hour.csv", newline="", encoding="utf-8") as f:     # cross-check task 5 without Beam
    check = sum(int(r["cnt"]) for r in csv.DictReader(f) if r["dteday"] == "2012-07-04")
print("check: 2012-07-04 summed directly from the file:", check)

T1 non-cancelled bookings per market segment
   Aviation: 185
   Complementary: 646
   Corporate: 4303
   Direct: 10672
   Groups: 7714
   Offline TA/TO: 15908
   Online TA: 35738
T2 mean nightly rate per hotel
   City Hotel: 105.75
   Resort Hotel: 90.79
T3 top 3 guest countries per hotel
   City Hotel: PRT 10879, FRA 7081, DEU 5012
   Resort Hotel: PRT 10192, GBR 5923, ESP 3106
T4 bookings by city (side-input lookup)
   Algarve: 28938
   Lisbon: 46228
T5 daily bike rentals, 1-9 July 2012 (fixed event-time windows)
   2012-07-01: 5531
   2012-07-02: 6227
   2012-07-03: 6660
   2012-07-04: 7403
   2012-07-05: 6241
   2012-07-06: 6207
   2012-07-07: 4840
   2012-07-08: 4672
   2012-07-09: 6569
check: 2012-07-04 summed directly from the file: 7403