Optional Extension: Athena

GSB 5544 · Not part of the required first session

Everything required in Session 1 runs on SageMaker alone. This notebook is for later, or for the curious. Read the Athena extension guide on the course site first. It explains what Athena is, the one-time permission step this notebook needs, what it costs, and how to clean up.

Athena is a query service. You send it SQL; AWS’s own machines scan Parquet files in S3 and write the answer to a bucket you own. So this notebook, unlike the others, creates something in your account (a results bucket) and can incur a small per-query charge.

Run it in JupyterLab inside your SageMaker space. Cells with ____ have a blank to fill before you run them. If a cell stops with NameError: name '____' is not defined, you ran it before filling in its blank; fill the blank and run it again.

# RUNS ON: SageMaker (remote)
%pip install -q s3fs pyarrow psutil      # harmless if already present

import os
import json
import time
import socket
import psutil
import boto3
import pandas as pd
from botocore import UNSIGNED
from botocore.config import Config

BUCKET = "noaa-ghcn-pds"         # NOAA's public bucket
REGION = "us-east-1"             # the region the bucket is in; your SageMaker space should be in it too

# unsigned = anonymous: this client attaches no credentials, because the bucket is public
s3 = boto3.client("s3", region_name=REGION, config=Config(signature_version=UNSIGNED))


def human(n, base=1000):
    """Turn a byte count into a readable string. base=1000 gives KB/MB/GB; base=1024 gives KiB/MiB/GiB."""
    units = ["B", "KB", "MB", "GB", "TB", "PB"] if base == 1000 else ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
    for unit in units:
        if n < base:
            return f"{n:,.1f} {unit}"
        n /= base
    return f"{n:,.1f} (huge)"


def where_am_i():
    """Report which computer this kernel is really on. A RUNS ON comment cannot do that."""
    on_sagemaker = (os.path.exists("/opt/ml/metadata/resource-metadata.json")
                    or os.path.expanduser("~") == "/home/sagemaker-user")
    place = "SageMaker (remote)" if on_sagemaker else "not SageMaker (probably your laptop)"
    print("Computer name   :", socket.gethostname())
    print("This looks like :", place)
    return place


def machine_report(work_dir="."):
    """Measure the computer running this cell. Disk numbers are for the filesystem that holds work_dir."""
    vm = psutil.virtual_memory()
    du = psutil.disk_usage(work_dir)
    return {
        "computer name":                    socket.gethostname(),
        "physical CPU cores":               psutil.cpu_count(logical=False),   # may be None on some virtual machines
        "logical CPUs (vCPUs in the cloud)": psutil.cpu_count(logical=True),
        "total RAM (GB)":                   round(vm.total / 1e9, 1),
        "available RAM right now (GB)":     round(vm.available / 1e9, 1),
        "work disk, total (GB)":            round(du.total / 1e9, 1),
        "work disk, free (GB)":             round(du.free / 1e9, 1),
    }


def slice_prefix(year, element):
    """The key prefix of one year-and-element slice of the Parquet copy of GHCN-Daily."""
    return f"parquet/by_year/YEAR={year}/ELEMENT={element}/"


def prefix_size(prefix):
    """Count the objects under a prefix and add up their sizes, without reading any of them."""
    objects, total = 0, 0
    for page in s3.get_paginator("list_objects_v2").paginate(Bucket=BUCKET, Prefix=prefix):
        for obj in page.get("Contents", []):
            objects += 1
            total += obj["Size"]
    return objects, total


if where_am_i() != "SageMaker (remote)":
    print()
    print("NOTE: this notebook is written to run in your SageMaker space, and this kernel is somewhere else.")
    print("      The RUNS ON comments describe the plan. They do not move the code. Open this file in JupyterLab on SageMaker.")
# RUNS ON: SageMaker (remote)
import io

STATIONS = {"San Luis Obispo": "USC00047851", "Phoenix": "USW00023183"}
STATION  = STATIONS["San Luis Obispo"]
BY_YEAR_2024_BYTES = s3.head_object(Bucket=BUCKET, Key="csv/by_year/2024.csv")["ContentLength"]
print("The 2024 by-year CSV is", human(BY_YEAR_2024_BYTES))

1. A results bucket you own

Athena writes every result to S3, so you need a bucket of your own. These clients are signed: they act as you, because this bucket is yours, not public. If this cell stops with AccessDenied, the permission step in the extension reading has not been done.

Fill in the bucket name in the create call.

# RUNS ON: SageMaker (remote)
sts     = boto3.client("sts",    region_name=REGION)
athena  = boto3.client("athena", region_name=REGION)
s3_mine = boto3.client("s3",     region_name=REGION)

ACCOUNT        = sts.get_caller_identity()["Account"]
RESULTS_BUCKET = f"gsb5544-athena-{ACCOUNT}"     # bucket names are global, so include your account id
RESULTS        = f"s3://{RESULTS_BUCKET}/athena/"

s3_mine.create_bucket(Bucket=____)   # in us-east-1 no LocationConstraint is needed
print("Results will land in", RESULTS)

2. Helpers and a table definition

run_athena starts a query, checks once a second until it finishes, and returns the query id plus Athena’s statistics, which include the bytes scanned. The table definition tells Athena where the Parquet files are and how the YEAR=/ELEMENT= key layout maps to columns (this is called partition projection). Nothing is copied, and these definition statements are free.

Fill in the boto3 call that starts a query.

# RUNS ON: SageMaker (remote)
ATHENA_BYTES = 0                                  # running total of everything scanned today

def run_athena(sql, database=None):
    """Start a query, wait for it, return (query id, statistics)."""
    global ATHENA_BYTES
    kwargs = {"QueryString": sql, "ResultConfiguration": {"OutputLocation": RESULTS}}
    if database:
        kwargs["QueryExecutionContext"] = {"Database": database}
    qid = athena.____(**kwargs)["QueryExecutionId"]
    while True:
        q = athena.get_query_execution(QueryExecutionId=qid)["QueryExecution"]
        state = q["Status"]["State"]
        if state in ("SUCCEEDED", "FAILED", "CANCELLED"):
            break
        time.sleep(1)
    if state != "SUCCEEDED":
        raise RuntimeError(q["Status"].get("StateChangeReason", state))
    ATHENA_BYTES += q["Statistics"].get("DataScannedInBytes", 0)
    return qid, q["Statistics"]


def athena_df(qid):
    """Read a finished query's result CSV from the results bucket."""
    obj = s3_mine.get_object(Bucket=RESULTS_BUCKET, Key=f"athena/{qid}.csv")
    return pd.read_csv(io.BytesIO(obj["Body"].read()))
# RUNS ON: Athena (AWS-managed)
run_athena("CREATE DATABASE IF NOT EXISTS gsb5544")

DDL = '''
CREATE EXTERNAL TABLE IF NOT EXISTS gsb5544.ghcn (
    id         string,
    `date`     string,
    data_value bigint,
    m_flag     string,
    q_flag     string,
    s_flag     string,
    obs_time   string
)
PARTITIONED BY (year int, element string)
STORED AS PARQUET
LOCATION 's3://noaa-ghcn-pds/parquet/by_year/'
TBLPROPERTIES (
    'projection.enabled'        = 'true',
    'projection.year.type'      = 'integer',
    'projection.year.range'     = '1750,2030',
    'projection.element.type'   = 'enum',
    'projection.element.values' = 'TMAX,TMIN,PRCP,SNOW,SNWD,TAVG',
    'storage.location.template' = 's3://noaa-ghcn-pds/parquet/by_year/YEAR=${year}/ELEMENT=${element}/'
)
'''
run_athena(DDL, database="gsb5544")
print("Table gsb5544.ghcn is defined. Bytes scanned so far:", human(ATHENA_BYTES))

3. One station, one year

The WHERE clause is the extraction. Fill in the year and the name of the statistic that reports bytes scanned.

# RUNS ON: Athena (AWS-managed)
SQL = f'''
SELECT id, "date", data_value
FROM   gsb5544.ghcn
WHERE  year = ____
  AND  element IN ('TMAX', 'PRCP')
  AND  q_flag IS NULL
  AND  id = '{STATION}'
'''

qid, stats = run_athena(SQL, database="gsb5544")
scanned = stats["____"]
station_2024 = athena_df(qid)

print(f"Rows returned : {len(station_2024):,}")
print(f"Bytes scanned : {human(scanned)}")
print(f"Query time    : {stats['TotalExecutionTimeInMillis'] / 1000:.1f} s")
print(f"Cost          : ${max(scanned, 10_000_000) / 1e12 * 5:.4f}   (about $5 per TB scanned, 10 MB minimum; check the Athena pricing page)")
print(f"2024 CSV is   : {human(BY_YEAR_2024_BYTES)}")
station_2024.head()

Athena scanned far less than the by-year CSV. Which feature of the bucket’s layout let it skip every other year and element before reading a row? One sentence, as a comment.

# RUNS ON: SageMaker (remote)
# ____

4. Hot days, computed entirely by Athena

Write one query that returns the number of days with TMAX at or above 35 °C (data_value >= 350) per station per year from 2015 through 2024, for both lab stations, ignoring quality-flagged values. Pivot the result to years as rows and markets as columns. It should match the hot_days table from the lab.

# RUNS ON: Athena (AWS-managed)
# your code here

A station-year with no row in this result had no day at or above 35 °C among the days it reported. The fillna(0) in the pivot is correct for a count of hot days, but it cannot tell you whether the station reported all year. The lab’s days_reported column exists for that reason.

5. Clean up: empty the results bucket

Unlike the required session, this notebook created something in S3 that stays after the space stops. Empty it. Fill in the call that lists the bucket and the call that deletes a batch of objects.

# RUNS ON: SageMaker (remote)
listing = s3_mine.____(Bucket=RESULTS_BUCKET)
objects = listing.get("Contents", [])
print(f"{len(objects)} result objects, {human(sum(o['Size'] for o in objects))}")

if objects:
    s3_mine.____(Bucket=RESULTS_BUCKET, Delete={"Objects": [{"Key": o["Key"]} for o in objects]})

print("Bucket empty :", "Contents" not in s3_mine.list_objects_v2(Bucket=RESULTS_BUCKET))
print("Athena today :", human(ATHENA_BYTES), f"scanned, about ${max(ATHENA_BYTES, 10_000_000) / 1e12 * 5:.4f}")

Then stop the space, as always.