diff --git a/.dockerignore b/.dockerignore index c5ca45a..2bba668 100644 --- a/.dockerignore +++ b/.dockerignore @@ -11,6 +11,10 @@ data/raw/generated # The external dataset cache is 44 MB and downloads on demand. data/raw/external +# The pipeline's lake: a 96 MB raw CSV and Spark's output. The pipeline +# service mounts ./data instead, and baking it in copied it into every image. +data/lake + # Build outputs, not inputs. reports **/*.xlsx diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0453875..b63a7dc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,6 +19,26 @@ jobs: # the floor was the only version anyone had ever run. python-version: ["3.12", "3.13"] + # The warehouse tests compare every query on PostgreSQL with SQLite. Without + # a server they skip, so CI provides one rather than letting them pass by + # not running. + services: + postgres: + image: postgres:16 + env: + POSTGRES_USER: rootsignal + POSTGRES_PASSWORD: rootsignal + POSTGRES_DB: rootsignal + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U rootsignal" + --health-interval 5s + --health-retries 10 + + env: + ROOTSIGNAL_WAREHOUSE_URL: postgresql://rootsignal:rootsignal@localhost:5432/rootsignal + steps: - name: Checkout repository uses: actions/checkout@v5 @@ -29,14 +49,21 @@ jobs: python-version: ${{ matrix.python-version }} cache: pip + # Spark runs on the JVM; the conform job's tests need one. + - name: Set up Java + uses: actions/setup-java@v5 + with: + distribution: temurin + java-version: "17" + - name: Install project and development dependencies - run: python -m pip install -e ".[dev]" + run: python -m pip install -e ".[dev,pipeline]" - name: Generate controlled sample dataset run: python scripts/generate_sample_data.py --output-dir data/raw/generated - name: Run test suite - run: pytest + run: pytest -rs - name: Run Ruff checks run: ruff check . diff --git a/.gitignore b/.gitignore index be3f4e1..5c2dc9b 100644 --- a/.gitignore +++ b/.gitignore @@ -40,3 +40,6 @@ mutants/ .vscode/ .idea/ .DS_Store + +# The pipeline's local lake (rebuilt by scripts/run_pipeline.py) +data/lake/ diff --git a/Dockerfile b/Dockerfile index c1440b5..62ce430 100644 --- a/Dockerfile +++ b/Dockerfile @@ -22,13 +22,19 @@ ENV PYTHONDONTWRITEBYTECODE=1 \ WORKDIR /app +# Spark runs on the JVM. The conform job is tested in here, on Linux, which is +# also what Glue and EMR run it on. +RUN apt-get update \ + && apt-get install -y --no-install-recommends openjdk-17-jre-headless \ + && rm -rf /var/lib/apt/lists/* + # Dependencies are installed from pyproject.toml alone, against a placeholder # package, so that editing source code does not invalidate the layer holding # pandas, numpy and streamlit. COPY pyproject.toml README.md ./ RUN mkdir -p src/rootsignal \ && touch src/rootsignal/__init__.py \ - && pip install --no-cache-dir ".[dev,dashboard,ai]" \ + && pip install --no-cache-dir ".[dev,dashboard,ai,pipeline]" \ && rm -rf src COPY . . diff --git a/README.md b/README.md index dfb4281..225b133 100644 --- a/README.md +++ b/README.md @@ -93,8 +93,10 @@ Reproduced by the committed test suite and evaluation scripts. | **Data quality** | **8 errors** detected across 10 raw tables; cleaning imputes 3 fields, drops 2 exact duplicates, quarantines 1 impossible line; **zero errors** after | | **Plan vs actual** | Attainment spans **87.4% to 125.6%** across 64 region/category/channel segments, evenly split 32 above plan and 32 below. KAM quota attainment runs **89.8% to 106.0%** across four managers, three ahead and one behind | | **Reproducibility** | Identical results on **numpy 1.26 under Windows and numpy 2.5 under Linux** — same signal, same impact to the cent, same forecast improvement. The container resolves the top of the declared dependency range rather than a lockfile, which is how the numpy 2.x break was found | -| **Tests** | **316 automated tests**, passing on both dependency sets. One asserts that no output ever claims causation | -| **Mutation testing** | `scripts/mutation_test.py` applies 37 operator and constant mutations to the KPI layer; the suite kills **34 (91.9%)**. The 3 survivors are equivalent mutants (`> 0` vs `>= 0` on a count that is never zero, or on a ratio where 0/0 is already NaN) | +| **Data pipeline** | The same 1,067,371 lines landed in S3, conformed by a **PySpark** job, cleaned, and loaded into **PostgreSQL** (996,528 sales lines) and **Redshift Serverless** (996,531, loaded before stock codes were made case-insensitive). Spark's output is **identical to the pandas adapter's in all six tables**, compared column by column on the full dataset | +| **Warehouse parity** | Every staging view and all seven analytical queries return the **same rows on PostgreSQL and Redshift as on SQLite**. The one tolerance is `aov`, which may differ by a cent on an exact half cent because the engines round ties differently | +| **Tests** | **427 automated tests**, passing on both dependency sets. One asserts that no output ever claims causation | +| **Mutation testing** | `scripts/mutation_test.py` applies 37 operator and constant mutations to the KPI layer; the suite kills **34 (91.9%)**. The 3 survivors are equivalent mutants (`> 0` vs `>= 0` on a count that is never zero, or on a ratio where 0/0 is already NaN). On the pipeline, `scripts/mutation_test_pipeline.py` applies **28 realistic bugs** (Spark's rounding, a lost CSV option, a dropped tie-breaker, a spending cap that only logs) and the suite catches **all 28**, up from 22 before the gaps it found were tested. Operator mutants on the pipeline code: **93 of 108** caught; the 15 survivors are wait times, output file counts and print formatting | **On that false-alarm figure.** The 12 counts every signal that matched nothing planted, wherever it was raised: 9 on the two control windows and 3 alongside @@ -119,6 +121,13 @@ Or don't, and check the claims instead: docker build -t rootsignal . && docker run --rm rootsignal pytest ``` +Or run the whole pipeline on the real dataset, from download to warehouse: + +```bash +docker compose up -d warehouse +docker compose run --rm pipeline # land, conform in Spark, load, publish +``` + Then, in rough order of what is worth seeing first: ```bash @@ -182,9 +191,35 @@ Evidence package — evidence, pattern, impact, confidence criteria, alternati Recommended investigation ``` -**Stack:** Python, pandas, NumPy, SQL (SQLite/PostgreSQL), Streamlit, Plotly, -pytest, GitHub Actions, and an optional OpenAI-compatible LLM that the system -works fully without. Power BI is deliberately out of scope. +**Stack:** Python, pandas, NumPy, PySpark, SQL (SQLite, PostgreSQL, Amazon +Redshift), AWS S3, Streamlit, Plotly, Power BI, pytest, Docker, GitHub Actions, +and an optional OpenAI-compatible LLM that the system works fully without. + +### Data pipeline + +```text +UCI workbook -> S3 raw zone -> Spark conform job -> S3 curated zone (Parquet) + -> cleaning + validation (nothing loads while an error remains) + -> PostgreSQL warehouse + S3 clean zone -> Redshift Serverless (COPY) + -> report tables (rate/mix split, drivers, forecast accuracy) -> Power BI +``` + +`jobs/conform_online_retail.py` is a PySpark version of the Online Retail +adapter with no imports from the package, so the same file runs locally, in the +container and on a Spark cluster. `scripts/compare_spark_to_pandas.py` checks it +against the pandas adapter on the full dataset. Cleaning stays in pandas: it +owns the audit trail and the quarantine, and a second implementation would mean +a second definition of every rule. + +Locally the lake is a folder and the warehouse a PostgreSQL container +(`docker compose run --rm pipeline`). `scripts/run_pipeline_aws.py` runs the same +stages with the lake in S3 and a Redshift Serverless warehouse, capped at 8 RPU +with a usage limit that switches it off, and a teardown step that removes +everything that bills. Glue and EMR Serverless subcommands exist for the Spark +step but have not been run: the account used is not yet permitted Glue jobs, +and its EMR Serverless quota had not taken effect. The Spark job itself has only +run locally so far. [power_bi.md](docs/power_bi.md) describes the report on the +warehouse. ## What's built @@ -204,6 +239,8 @@ works fully without. Power BI is deliberately out of scope. | RootSignal engine | Evidence, pattern, impact, confidence and recommended investigation, ranked | [signals.md](docs/signals.md) | | Scenario evaluation | Measures whether the engine tells three planted situations apart and stays silent on two controls | [signals.md](docs/signals.md) | | SQL layer | Staging views, commercial mart and seven business queries, all executed by tests | [sql.md](docs/sql.md) | +| Warehouse | The same SQL on PostgreSQL and Redshift, translated from one source and held row for row to SQLite | [power_bi.md](docs/power_bi.md) | +| Pipeline | Raw, curated and clean zones on disk or in S3, a PySpark conform job, and report tables for BI | [power_bi.md](docs/power_bi.md) | | Excel reporting | Six operational workbooks, each opening with what its figures mean | [reporting.md](docs/reporting.md) | | Dashboard | Five Streamlit pages over a Streamlit-free, tested data layer | [dashboard.md](docs/dashboard.md) | | Presentation | One map from schema keys to language a business reader already has | [dashboard.md](docs/dashboard.md#the-words-on-the-screen) | @@ -251,8 +288,10 @@ against six named criteria rather than a tuned score. It is validated against five planted scenarios — classifying three of three correctly at a medium confidence floor, with zero false alarms at a high one — and runs unchanged on a million real invoice lines from a public dataset, where it beats a naive -forecast baseline by 38.8% across 49 backtest folds. Python, pandas, SQL, -Streamlit; 308 tests, including one asserting that no output ever claims -causation. +forecast baseline by 38.8% across 49 backtest folds. That dataset moves through +a pipeline from S3 through a PySpark job into PostgreSQL and Redshift, with +Spark's output checked against the pandas implementation on every row. Python, +pandas, PySpark, SQL, AWS, Streamlit; 427 tests, including one asserting that no +output ever claims causation. diff --git a/docker-compose.yml b/docker-compose.yml index 3bbedb9..bbe36a8 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -24,3 +24,38 @@ services: image: rootsignal profiles: ["tools"] command: pytest + + # Source to warehouse: land the published data, conform it in Spark, load + # it into the warehouse below. The lake is ./data/lake on the host. + pipeline: + build: . + image: rootsignal + profiles: ["tools"] + depends_on: + warehouse: + condition: service_healthy + environment: + - ROOTSIGNAL_WAREHOUSE_URL=postgresql://rootsignal:rootsignal@warehouse:5432/rootsignal + volumes: + - ./data:/app/data + command: python scripts/run_pipeline.py + + # The warehouse the pipeline loads into. Credentials are for a local + # container only; a deployed warehouse takes its URL from the environment. + warehouse: + image: postgres:16 + environment: + - POSTGRES_USER=rootsignal + - POSTGRES_PASSWORD=rootsignal + - POSTGRES_DB=rootsignal + ports: + - "5433:5432" + volumes: + - warehouse-data:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U rootsignal"] + interval: 5s + retries: 10 + +volumes: + warehouse-data: diff --git a/docs/RootSignal.pbit b/docs/RootSignal.pbit new file mode 100644 index 0000000..e6f78b2 Binary files /dev/null and b/docs/RootSignal.pbit differ diff --git a/docs/docker.md b/docs/docker.md index 28bd74f..4dd55d3 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -19,7 +19,7 @@ docker run --rm rootsignal pytest ```bash docker run --rm -p 8501:8501 rootsignal # the dashboard -docker run --rm rootsignal pytest # all 308 tests +docker run --rm rootsignal pytest # the full suite, Spark tests included docker run --rm rootsignal python scripts/detect_signals.py ``` diff --git a/docs/power_bi.md b/docs/power_bi.md new file mode 100644 index 0000000..1c9b320 --- /dev/null +++ b/docs/power_bi.md @@ -0,0 +1,87 @@ +# Power BI on the RootSignal warehouse + +The pipeline loads the cleaned model and the analysis results into PostgreSQL, +and this report reads them from there. It computes nothing the project has not +already computed and tested: revenue comes from the fact table, and the country +contributions, the rate and mix split, month-on-month change and forecast +accuracy come from report tables the pipeline writes. + +| Overview | Drivers | +| --- | --- | +| ![Overview](screenshots/powerbi_overview.png) | ![Drivers](screenshots/powerbi_drivers.png) | +| **Returns** | **Forecast** | +| ![Returns](screenshots/powerbi_returns.png) | ![Forecast](screenshots/powerbi_forecast.png) | + +## Open it + +1. Fill the warehouse: + ```bash + docker compose up -d warehouse + docker compose run --rm pipeline + ``` +2. Open `docs/RootSignal.pbit` in Power BI Desktop. It is a template: the + report and model without data, so it asks for the connection and loads. +3. Credentials: **Database**, user `rootsignal`, password `rootsignal`. These + belong to the local container in `docker-compose.yml`. +4. The container has no TLS certificate, so Power BI asks to connect without + encryption: **OK**. + +Loading a million sales lines takes about a minute. On a machine with 8 GB of +RAM, run the report and the Spark pipeline at different times; together they +exhaust Docker's memory. + +## What is in it + +**Overview.** Net sales, invoices, average order value and units as cards, with +a country slicer that filters them. Net sales by month, and a table of each +month against the one before, taken from `rpt_monthly_trend`. + +**Drivers.** How much each country moved net sales between October and +November 2011, from `rpt_country_contribution`. + +**Returns.** Why the return rate moved over the same two months, split into a +rate effect (countries returning a different share) and a mix effect (demand +moving toward countries that always return more), from +`rpt_return_rate_components` and `rpt_return_rate_split`. The rate effect +carries almost all of the movement. + +**Forecast.** Backtested error of each daily revenue model against the naive +baseline, from `rpt_forecast_accuracy`. `seasonal_mean_7` is 38.8% better than +naive over 49 folds, the figure the README reports. + +## Building it yourself + +The template is the reliable route. To rebuild the report from an empty file: + +- **Tables.** Get data → PostgreSQL database, server `localhost:5433`, database + `rootsignal`, Import. Tables appear with the schema as a prefix, as + `rootsignal fact_sales` and so on. +- **Relationships.** `fact_sales` to `dim_date` on `date`, to `dim_customer` on + `customer_id` and to `dim_region` on `region_code`, each many to one, single + direction. Leave the `rpt_` tables unconnected; they are finished results. +- **Products.** `fact_sales` to `dim_sku` on `sku_id` works too. Building the + report found 172 stock codes written in two cases or with a trailing space + (`15056BL` and `15056bl`); Power BI compares text without case and saw + duplicates. The pipeline now trims and upper-cases codes, so the relationship + is many to one like the others. +- **Measures.** Power BI names ignore case, so a measure cannot be called + `Units` next to a column called `units`: + ```DAX + Net Sales = SUM ( 'rootsignal fact_sales'[net_sales] ) + Total Units = SUM ( 'rootsignal fact_sales'[units] ) + Invoices = DISTINCTCOUNT ( 'rootsignal fact_sales'[order_id] ) + AOV = DIVIDE ( [Net Sales], [Invoices] ) + ``` + `Invoices` is a distinct count because one invoice lists several products. + Expected totals: £20,972,594 net sales, 40,077 invoices, 11,438,460 units, + £523.31 average order value. +- **Month-on-month change** comes from `rpt_monthly_trend` rather than + `DATEADD`. `dim_date` holds only the days the retailer traded, and Power BI's + time-intelligence functions assume an unbroken calendar. +- **Partial months** are already excluded from the monthly report tables; the + data stops on 9 December 2011 and the pipeline drops that month. +- **Totals.** Turn off the total row on tables of rates and changes. Adding + percentages across rows produces a number that means nothing. + +A return rate is not a fill rate. Returns are customers sending goods back, not +a warehouse failing to ship. diff --git a/docs/screenshots/powerbi_drivers.png b/docs/screenshots/powerbi_drivers.png new file mode 100644 index 0000000..3193681 Binary files /dev/null and b/docs/screenshots/powerbi_drivers.png differ diff --git a/docs/screenshots/powerbi_forecast.png b/docs/screenshots/powerbi_forecast.png new file mode 100644 index 0000000..8f0a2f6 Binary files /dev/null and b/docs/screenshots/powerbi_forecast.png differ diff --git a/docs/screenshots/powerbi_overview.png b/docs/screenshots/powerbi_overview.png new file mode 100644 index 0000000..6109031 Binary files /dev/null and b/docs/screenshots/powerbi_overview.png differ diff --git a/docs/screenshots/powerbi_returns.png b/docs/screenshots/powerbi_returns.png new file mode 100644 index 0000000..c7f3e5b Binary files /dev/null and b/docs/screenshots/powerbi_returns.png differ diff --git a/jobs/conform_online_retail.py b/jobs/conform_online_retail.py new file mode 100644 index 0000000..64ba272 --- /dev/null +++ b/jobs/conform_online_retail.py @@ -0,0 +1,273 @@ +"""Conform raw Online Retail II invoice lines into the RootSignal model, in Spark. + +This is the Spark form of rootsignal.adapters.online_retail.adapt. It reads the +raw zone, applies the same mapping, and writes one Parquet dataset per table to +the curated zone. Cleaning, validation and the warehouse load run after it, on +its output, exactly as they run on the pandas adapter's. + +It is deliberately a single file that imports nothing from the rootsignal +package, so the same script runs under spark-submit on a laptop, as an AWS Glue +job and on EMR Serverless without packaging the project. It also keeps to +Python 3.9 syntax, which is what EMR Serverless runs. + +Every rule matches the pandas adapter, down to how numbers are rounded: +pandas rounds by scaling, rounding half to even and scaling back, so that is +what rint(x * 100) / 100 does here. Spark's own round() rounds the decimal +value half up, and gives a different cent on an exact half cent. On Online +Retail II that never happens: prices carry at most two decimals and quantities +are whole, and scripts/compare_spark_to_pandas.py counts 0 of 1,044,420 sales +lines affected. The job rounds the pandas way regardless, so the two stay +equal on data where it does happen; the test fixture includes such a line. +tests/test_spark_conform.py asserts the two produce the same tables. + + spark-submit jobs/conform_online_retail.py --input --output +""" + +from __future__ import annotations + +import argparse +import json +import sys + +from pyspark.sql import DataFrame, SparkSession +from pyspark.sql import functions as F +from pyspark.sql import types as T + +RAW_SCHEMA = T.StructType( + [ + T.StructField("Invoice", T.StringType()), + T.StructField("StockCode", T.StringType()), + T.StructField("Description", T.StringType()), + T.StructField("Quantity", T.StringType()), + T.StructField("InvoiceDate", T.StringType()), + T.StructField("Price", T.StringType()), + T.StructField("Customer ID", T.StringType()), + T.StructField("Country", T.StringType()), + ] +) + +FACT_SALES_COLUMNS = [ + "order_id", "date", "customer_id", "region_code", "channel", "kam_id", + "sku_id", "category", "sales_type", "units", "unit_price", "discount_pct", "net_sales", +] +SALES_KEY = ["order_id", "sku_id", "sales_type"] + +# Files written per table. Left alone, Spark writes one file per shuffle +# partition, 200 by default: 200 files of about 100 KB for a million sales lines, +# which every reader, Redshift's COPY included, then has to open one by one. +OUTPUT_FILES = {"fact_sales": 4} + + +def round_half_even(column, places: int): + """NumPy's rounding: scale, round half to even, scale back.""" + scale = float(10**places) + return F.rint(column * F.lit(scale)) / F.lit(scale) + + +def first_in_file_order(column: str): + """pandas groupby 'first': the first non-null value, in source row order.""" + return F.min_by(F.col(column), F.when(F.col(column).isNotNull(), F.col("_row"))) + + +def read_raw(spark: SparkSession, path: str) -> DataFrame: + """Invoice lines as published, every value read as text and typed explicitly. + + multiLine keeps a quoted description containing a line break as one row, and + reads the file in one pass, so _row follows the order of lines in the file. + """ + raw = ( + spark.read.option("header", True) + .option("multiLine", True) + .option("escape", '"') + .schema(RAW_SCHEMA) + .csv(path) + .withColumn("_row", F.monotonically_increasing_id()) + ) + return raw.select( + F.col("_row"), + F.col("Invoice").alias("order_id"), + # The same product appears as 15056BL and 15056bl, and as 47503J with + # and without a trailing space. Trimmed as Python's str.strip() trims, + # any Unicode whitespace, then upper-cased. + F.upper(F.regexp_replace("StockCode", r"(?U)^\s+|\s+$", "")).alias("sku_id"), + F.col("Description").alias("sku_name"), + F.col("Quantity").cast("long").alias("units"), + F.to_timestamp("InvoiceDate").alias("timestamp"), + F.col("Price").cast("double").alias("unit_price"), + F.col("Customer ID").alias("customer_id_raw"), + F.col("Country").alias("region_code"), + ) + + +def conform_lines(raw: DataFrame) -> DataFrame: + """Derived columns on every line, before sales and returns are separated.""" + first_word = F.regexp_extract(F.upper(F.coalesce(F.col("sku_name"), F.lit("UNKNOWN"))), r"(?U)^\s*(\S+)", 1) + return ( + raw.withColumn("date", F.to_date("timestamp")) + .withColumn("category", F.when(first_word == "", "UNKNOWN").otherwise(first_word)) + # A fifth of lines have no customer. They are named rather than dropped. + .withColumn( + "customer_id", + F.coalesce(F.regexp_replace("customer_id_raw", r"\.0$", ""), F.lit("UNKNOWN")), + ) + .withColumn("channel", F.lit("Online")) + .withColumn("sales_type", F.lit("PRIMARY")) + .withColumn("kam_id", F.lit("UNASSIGNED")) + .withColumn( + "is_return", + (F.col("units") < 0) | F.upper(F.col("order_id")).startswith("C"), + ) + ) + + +def consolidate_repeated_lines(sales: DataFrame) -> tuple[DataFrame, int]: + """Combine lines repeating a product on the same invoice. + + Units and revenue are added and the price follows from them, so volume and + revenue are both unchanged. As in the pandas adapter, the price is + recomputed on every line whenever any line was combined. + """ + total = sales.count() + distinct = sales.select(*SALES_KEY).distinct().count() + repeated = total - distinct + if repeated == 0: + return sales.select(*FACT_SALES_COLUMNS), 0 + + grouped = sales.groupBy(*SALES_KEY).agg( + *(first_in_file_order(name).alias(name) + for name in ("date", "customer_id", "region_code", "channel", "kam_id", "category", "discount_pct")), + F.sum("units").alias("units"), + F.sum("net_sales").alias("net_sales"), + ) + grouped = grouped.withColumn( + "unit_price", + round_half_even( + F.col("net_sales") / F.when(F.col("units") != 0, F.col("units").cast("double")), 4 + ), + ) + return grouped.select(*FACT_SALES_COLUMNS), repeated + + +def build_dim_date(fact_sales: DataFrame) -> DataFrame: + return ( + fact_sales.select("date").where(F.col("date").isNotNull()).distinct() + .select( + "date", + F.weekofyear("date").alias("week"), + F.month("date").alias("month"), + F.concat(F.lit("Q"), F.quarter("date").cast("string")).alias("quarter"), + F.year("date").alias("year"), + ) + ) + + +def build_dim_sku(sales: DataFrame) -> DataFrame: + # Price comes from lines with a real one; samples and adjustments sit at or + # below zero and would drag a median negative. + firsts = sales.groupBy("sku_id").agg( + first_in_file_order("sku_name").alias("sku_name"), + first_in_file_order("category").alias("category"), + ) + medians = ( + sales.where(F.col("unit_price") > 0) + .groupBy("sku_id") + .agg(F.expr("percentile(unit_price, 0.5)").alias("list_price")) + ) + sku = firsts.join(medians, "sku_id", "left").fillna({"list_price": 0.0}) + return sku.select( + "sku_id", + F.coalesce("sku_name", F.lit("Not recorded")).alias("sku_name"), + "category", + F.lit("Standard").alias("sub_category"), + F.lit(1.0).alias("pack_size_kg"), + round_half_even(F.col("list_price") * F.lit(0.7), 4).alias("unit_cost"), + "list_price", + F.lit(True).alias("active_flag"), + ) + + +def build_dim_customer(sales: DataFrame) -> DataFrame: + customer = sales.groupBy("customer_id").agg(first_in_file_order("region_code").alias("region_code")) + return customer.select( + "customer_id", + F.when(F.col("customer_id") == "UNKNOWN", F.lit("Unidentified customer")) + .otherwise(F.concat(F.lit("Customer "), F.col("customer_id"))) + .alias("customer_name"), + F.lit("Retail").alias("customer_type"), + F.lit("Online").alias("channel"), + "region_code", + F.col("region_code").alias("city"), + F.col("region_code").alias("zone"), + F.lit("UNASSIGNED").alias("kam_id"), + ) + + +def build_dim_region(sales: DataFrame) -> DataFrame: + regions = sales.select("region_code").where(F.col("region_code").isNotNull()).distinct() + return regions.select("region_code", F.col("region_code").alias("city"), F.col("region_code").alias("zone")) + + +def conform(spark: SparkSession, input_path: str) -> tuple[dict, dict]: + """Every curated table, plus a summary of what the run did.""" + raw = read_raw(spark, input_path) + missing_keys = raw.where(F.col("order_id").isNull() | F.col("sku_id").isNull()).count() + if missing_keys: + raise ValueError(f"{missing_keys} raw line(s) have no invoice or stock code.") + + lines = conform_lines(raw).cache() + sales = ( + lines.where(~F.col("is_return")) + .withColumn("discount_pct", F.lit(0.0)) + .withColumn("net_sales", round_half_even(F.col("units") * F.col("unit_price"), 2)) + .cache() + ) + returns = lines.where(F.col("is_return")).select( + "order_id", "sku_id", "region_code", "customer_id", "date", "units", + F.abs("units").alias("returned_units"), + ) + + fact_sales, repeated = consolidate_repeated_lines(sales) + fact_sales = fact_sales.cache() + + tables = { + "fact_sales": fact_sales, + "dim_date": build_dim_date(fact_sales), + "dim_sku": build_dim_sku(sales), + "dim_customer": build_dim_customer(sales), + "dim_region": build_dim_region(sales), + "dim_kam": spark.createDataFrame([("UNASSIGNED", "Not recorded")], "kam_id string, kam_name string"), + "returns": returns, + } + summary = { + "source_rows": lines.count(), + "returned_lines": returns.count(), + "consolidated_lines": repeated, + "tables": {name: frame.count() for name, frame in tables.items()}, + } + return tables, summary + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--input", required=True, help="Raw invoice CSV file or prefix.") + parser.add_argument("--output", required=True, help="Curated prefix; one Parquet dataset per table.") + # Glue passes its own arguments alongside these. + args, _ = parser.parse_known_args(argv) + + spark = ( + SparkSession.builder.appName("rootsignal-conform-online-retail") + .config("spark.sql.session.timeZone", "UTC") + .getOrCreate() + ) + tables, summary = conform(spark, args.input) + output = args.output.rstrip("/") + for name, frame in tables.items(): + frame.coalesce(OUTPUT_FILES.get(name, 1)).write.mode("overwrite").parquet(f"{output}/{name}") + spark.createDataFrame([(json.dumps(summary),)], "summary string").coalesce(1).write.mode( + "overwrite" + ).text(f"{output}/_summary") + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/pyproject.toml b/pyproject.toml index 5bbd337..285bcce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ dev = [ "pytest>=8,<9", "ruff>=0.6,<1", + "moto[s3]>=5,<6", ] dashboard = [ @@ -30,6 +31,20 @@ ai = [ "openai>=1.50,<2", ] +# The warehouse: PostgreSQL through psycopg. +warehouse = [ + "psycopg[binary]>=3.2,<4", +] + +# The data pipeline: Spark for the conform step, Parquet for the curated zone, +# boto3 for S3, Glue and Redshift. Spark needs a Java runtime as well. +pipeline = [ + "pyspark>=3.5,<4", + "pyarrow>=15", + "boto3>=1.34,<2", + "psycopg[binary]>=3.2,<4", +] + [tool.pytest.ini_options] pythonpath = ["src"] testpaths = ["tests"] diff --git a/scripts/compare_spark_to_pandas.py b/scripts/compare_spark_to_pandas.py new file mode 100644 index 0000000..09c23f8 --- /dev/null +++ b/scripts/compare_spark_to_pandas.py @@ -0,0 +1,126 @@ +"""Compare the Spark conform job with the pandas adapter on the whole dataset. + +tests/test_spark_conform.py holds the two to each other on a planted fixture. +This does the same on all 1,067,371 published lines, table by table, and +reports every column that differs. Run it after the conform stage: + + docker compose run --rm pipeline python scripts/run_pipeline.py --stages land,conform + docker compose run --rm pipeline python scripts/compare_spark_to_pandas.py + +It also counts the lines where Spark's built-in round() would have disagreed +with the pandas-style rounding the job uses. On this dataset the count is 0. +""" + +from __future__ import annotations + +import argparse +import time + +import numpy as np +import pandas as pd + +from rootsignal.adapters.online_retail import adapt, read_raw +from rootsignal.console import use_utf8_output +from rootsignal.pipeline import RAW, Lake, read_curated +from rootsignal.pipeline.stages import RAW_CSV + +KEYS = { + "fact_sales": ["order_id", "sku_id", "sales_type"], + "dim_date": ["date"], + "dim_sku": ["sku_id"], + "dim_customer": ["customer_id"], + "dim_region": ["region_code"], + "dim_kam": ["kam_id"], +} + + +def compare(spark_table: pd.DataFrame, pandas_table: pd.DataFrame, key: list[str]) -> list[str]: + """Every difference between two tables, as readable lines. Empty when equal.""" + problems = [] + if sorted(spark_table.columns) != sorted(pandas_table.columns): + return [f"columns differ: {sorted(spark_table.columns)} vs {sorted(pandas_table.columns)}"] + if len(spark_table) != len(pandas_table): + problems.append(f"row count {len(spark_table):,} vs {len(pandas_table):,}") + + merged = spark_table.merge( + pandas_table, on=key, how="outer", suffixes=("_spark", "_pandas"), indicator=True + ) + unmatched = merged["_merge"] != "both" + if unmatched.any(): + problems.append(f"{int(unmatched.sum()):,} key(s) present on one side only") + both = merged.loc[~unmatched] + + for column in spark_table.columns: + if column in key: + continue + a, b = both[f"{column}_spark"], both[f"{column}_pandas"] + if pd.api.types.is_numeric_dtype(a) and pd.api.types.is_numeric_dtype(b) and a.dtype != bool: + x, y = a.astype(float).to_numpy(), b.astype(float).to_numpy() + gap = np.abs(np.where(np.isnan(x) & np.isnan(y), 0.0, x - y)) + differing = int((gap > 1e-9).sum() + (np.isnan(x) != np.isnan(y)).sum()) + if differing: + problems.append(f"{column}: {differing:,} row(s) differ, largest gap {np.nanmax(gap):.3g}") + else: + differing = int((a.astype(str).to_numpy() != b.astype(str).to_numpy()).sum()) + if differing: + problems.append(f"{column}: {differing:,} row(s) differ") + return problems + + +def spark_round_disagreements(raw_csv: str) -> tuple[int, int]: + """Sales lines where Spark's own round() would have given different revenue. + + Counted in Spark, on the job's own reading of the raw zone: round() against + the half-to-even rounding the job uses to match pandas. + """ + import importlib.util + + from pyspark.sql import SparkSession + from pyspark.sql import functions as F + + from rootsignal.pipeline.stages import JOB_PATH + + spec = importlib.util.spec_from_file_location("conform_online_retail", JOB_PATH) + job = importlib.util.module_from_spec(spec) + spec.loader.exec_module(job) + + spark = SparkSession.builder.appName("round-check").getOrCreate() + sales = job.conform_lines(job.read_raw(spark, raw_csv)).where(~F.col("is_return")) + revenue = F.col("units") * F.col("unit_price") + total = sales.count() + differing = sales.where(F.round(revenue, 2) != job.round_half_even(revenue, 2)).count() + spark.stop() + return differing, total + + +def main() -> None: + use_utf8_output() + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--lake", default="data/lake") + args = parser.parse_args() + + started = time.time() + spark_dataset, _ = read_curated(Lake(args.lake)) + raw = read_raw() + pandas_dataset = adapt(raw=raw) + print(f"{len(raw):,} published lines; both implementations read in {time.time() - started:.0f}s\n") + + failures = 0 + for name, key in KEYS.items(): + problems = compare(spark_dataset.tables[name], pandas_dataset.tables[name], key) + rows = len(spark_dataset.tables[name]) + print(f"{name:14s} {rows:>10,} rows {'identical' if not problems else 'DIFFERS'}") + for problem in problems: + print(f" {problem}") + failures += bool(problems) + + disagree, lines = spark_round_disagreements(Lake(args.lake).uri(RAW, RAW_CSV)) + print( + f"\nSpark's round() would have changed the revenue on {disagree:,} of {lines:,} sales lines " + f"({disagree / lines:.2%})." + ) + raise SystemExit(1 if failures else 0) + + +if __name__ == "__main__": + main() diff --git a/scripts/mutation_test_pipeline.py b/scripts/mutation_test_pipeline.py new file mode 100644 index 0000000..b4dc696 --- /dev/null +++ b/scripts/mutation_test_pipeline.py @@ -0,0 +1,164 @@ +"""Mutation testing for the pipeline with bugs it could realistically ship. + +scripts/mutation_test.py swaps operators and nudges constants, which suits the +KPI layer. Most of the pipeline's mistakes would not look like that: a +rounding rule swapped for Spark's own, a NULL ordering dropped, a CSV option +lost, a spending cap that only logs. Each mutant below is one such change, +applied alone; the suite that should notice is run, and the file restored. + +Run it where Spark and PostgreSQL are both available: + + docker compose up -d warehouse + docker compose run --rm pipeline python scripts/mutation_test_pipeline.py + +A surviving mutant is either a gap in the tests or a change no test here can +see. The second kind is listed with its reason, not hidden. +""" + +from __future__ import annotations + +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +from rootsignal.console import use_utf8_output + +ROOT = Path(__file__).resolve().parents[1] +JOB = "jobs/conform_online_retail.py" +SPARK_TESTS = ("tests/test_spark_conform.py",) +WAREHOUSE_TESTS = ("tests/test_warehouse.py",) + + +@dataclass(frozen=True) +class Mutant: + name: str + path: str + old: str + new: str + tests: tuple[str, ...] + + +MUTANTS = [ + # The Spark job + Mutant("Spark's own round() instead of pandas rounding", JOB, + "return F.rint(column * F.lit(scale)) / F.lit(scale)", "return F.round(column, places)", SPARK_TESTS), + Mutant("A C-prefixed invoice is no longer a return", JOB, + '(F.col("units") < 0) | F.upper(F.col("order_id")).startswith("C")', '(F.col("units") < 0)', SPARK_TESTS), + Mutant("A negative quantity is no longer a return", JOB, + '(F.col("units") < 0) | F.upper(F.col("order_id")).startswith("C")', + 'F.upper(F.col("order_id")).startswith("C")', SPARK_TESTS), + Mutant("Customer IDs keep the trailing .0", JOB, + 'F.regexp_replace("customer_id_raw", r"\\.0$", "")', 'F.col("customer_id_raw")', SPARK_TESTS), + Mutant("List price is the mean, not the median", JOB, + 'F.expr("percentile(unit_price, 0.5)")', 'F.avg("unit_price")', SPARK_TESTS), + Mutant("Zero prices count toward the list price", JOB, + 'sales.where(F.col("unit_price") > 0)', "sales", SPARK_TESTS), + Mutant("Prices are recomputed even with nothing to combine", JOB, + "if repeated == 0:", "if repeated < 0:", SPARK_TESTS), + Mutant("A blank description is not UNKNOWN", JOB, + 'F.when(first_word == "", "UNKNOWN").otherwise(first_word)', "first_word", SPARK_TESTS), + Mutant("CSV read line by line", JOB, '.option("multiLine", True)', '.option("multiLine", False)', SPARK_TESTS), + Mutant("CSV quotes escaped the Spark default way", JOB, """.option("escape", '"')""", "", SPARK_TESTS), + Mutant("Whitespace matched as ASCII only", JOB, 'r"(?U)^\\s*(\\S+)"', 'r"^\\s*(\\S+)"', SPARK_TESTS), + Mutant("First value taken in any order", JOB, + 'return F.min_by(F.col(column), F.when(F.col(column).isNotNull(), F.col("_row")))', + "return F.first(F.col(column), ignorenulls=True)", SPARK_TESTS), + # The SQL translation + Mutant("Week starts at the month", "src/rootsignal/sql/dialect.py", + "CAST(DATE_TRUNC('week', \\1) AS DATE)", "CAST(DATE_TRUNC('month', \\1) AS DATE)", WAREHOUSE_TESTS), + Mutant("REAL left as PostgreSQL's 4-byte float", "src/rootsignal/sql/dialect.py", + 'out = _REAL.sub("DOUBLE PRECISION", out)', "out = out", WAREHOUSE_TESTS), + Mutant("Untranslated SQLite passed through silently", "src/rootsignal/sql/dialect.py", + " if match:\n raise ValueError(", " if False:\n raise ValueError(", + WAREHOUSE_TESTS), + Mutant("CHECK constraints sent to Redshift", "src/rootsignal/sql/dialect.py", + " out = _strip_checks(out)\n", "", ("tests/test_aws.py",)), + Mutant("The warehouse truncates instead of rounding", "src/rootsignal/sql/dialect.py", + "AS $$ SELECT round(value::numeric, places)::double precision $$;", + "AS $$ SELECT trunc(value::numeric, places)::double precision $$;", + WAREHOUSE_TESTS), + # The queries + Mutant("Target variance sorted without a tie-breaker", "sql/analytics/target_variance.sql", + "ORDER BY variance_pct NULLS FIRST, p.region_code, p.category, p.channel;", + "ORDER BY variance_pct;", WAREHOUSE_TESTS), + Mutant("Weekly movers sorted without a tie-breaker", "sql/analytics/weekly_movers.sql", + "ORDER BY c.week_start, c.contribution, c.region_code, c.category;", + "ORDER BY c.week_start, c.contribution;", WAREHOUSE_TESTS), + Mutant("Stockout rate averaged over integers", "sql/analytics/supply_watchlist.sql", + "ROUND(AVG(CAST(i.stockout_flag AS REAL)), 4)", "ROUND(AVG(i.stockout_flag), 4)", WAREHOUSE_TESTS), + # Loading + Mutant("The old warehouse is dropped outside the load's transaction", "src/rootsignal/sql/warehouse.py", + " cursor.execute(f'CREATE SCHEMA \"{schema}\"')", + " conn.commit()\n cursor.execute(f'CREATE SCHEMA \"{schema}\"')", WAREHOUSE_TESTS), + Mutant("Clean zone integers written as 64-bit", "src/rootsignal/pipeline/stages.py", + '"INTEGER": pa.int32()', '"INTEGER": pa.int64()', ("tests/test_pipeline_stages.py",)), + Mutant("Report tables keep pandas Periods", "src/rootsignal/pipeline/publish.py", + " out[column] = series.dt.start_time.dt.date\n", " pass\n", + ("tests/test_pipeline_stages.py",)), + # The lake and AWS + Mutant("A folder picks up its siblings", "src/rootsignal/pipeline/lake.py", + 'if item["Key"] == base or item["Key"].startswith(base + "/"):', "if True:", ("tests/test_lake.py",)), + Mutant("Teardown reports a Glue job it did not remove", "src/rootsignal/pipeline/aws.py", + " glue.get_job(JobName=GLUE_JOB)\n glue.delete_job(JobName=GLUE_JOB)", + " glue.delete_job(JobName=GLUE_JOB)", ("tests/test_aws.py",)), + Mutant("Redshift's spending cap only logs", "src/rootsignal/pipeline/aws.py", + 'breachAction="deactivate"', 'breachAction="log"', ("tests/test_aws.py",)), + Mutant("Redshift may scale past its base size", "src/rootsignal/pipeline/aws.py", + "maxCapacity=base_capacity,", "maxCapacity=base_capacity * 4,", ("tests/test_aws.py",)), + Mutant("SET and SELECT sent as separate calls", "src/rootsignal/pipeline/aws.py", + 'statements = [f"SET search_path TO {schema}", sql] if schema else [sql]', + "statements = [sql]", ("tests/test_aws.py",)), +] + +# Mutants no test in this repository can catch, and why. They are still run, so +# the day a test starts catching one, this list is visibly out of date. +EXPECTED_SURVIVORS: dict[str, str] = {} + + +def run(mutant: Mutant) -> bool: + """Apply one mutant, run its tests, restore the file. True if killed.""" + path = ROOT / mutant.path + original = path.read_text(encoding="utf-8") + count = original.count(mutant.old) + if count != 1: + raise SystemExit(f"{mutant.name}: expected the original text once in {mutant.path}, found {count}") + try: + path.write_text(original.replace(mutant.old, mutant.new), encoding="utf-8") + result = subprocess.run( + [sys.executable, "-m", "pytest", "-x", "-q", "-p", "no:cacheprovider", "-o", "addopts=", *mutant.tests], + cwd=ROOT, + capture_output=True, + text=True, + ) + finally: + path.write_text(original, encoding="utf-8") + # Only a failing test is a kill. pytest also exits non-zero when a file is + # missing or nothing was collected, and counting that would inflate the score. + if result.returncode not in (0, 1): + raise SystemExit(f"{mutant.name}: pytest exited {result.returncode}\n{result.stdout[-2000:]}") + if result.returncode == 0 and " passed" not in result.stdout: + raise SystemExit(f"{mutant.name}: its tests all skipped, so the result would mean nothing") + return result.returncode == 1 + + +def main() -> int: + use_utf8_output() + killed, survived = [], [] + for mutant in MUTANTS: + (killed if run(mutant) else survived).append(mutant) + print(f"{'killed ' if mutant in killed else 'SURVIVED'} {mutant.name}", flush=True) + + print(f"\n{len(killed)}/{len(MUTANTS)} killed") + unexplained = [mutant for mutant in survived if mutant.name not in EXPECTED_SURVIVORS] + for mutant in survived: + reason = EXPECTED_SURVIVORS.get(mutant.name, "no test catches this") + print(f" survived: {mutant.name}. {reason}") + stale = [name for name in EXPECTED_SURVIVORS if name in {mutant.name for mutant in killed}] + for name in stale: + print(f" now killed, remove from EXPECTED_SURVIVORS: {name}") + return 1 if unexplained or stale else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_external_dataset.py b/scripts/run_external_dataset.py index db26ab9..55bc6eb 100644 --- a/scripts/run_external_dataset.py +++ b/scripts/run_external_dataset.py @@ -33,6 +33,7 @@ rolling_origin_evaluate, summarise_backtest, ) +from rootsignal.pipeline.publish import BACKTEST from rootsignal.validation import DatasetValidator @@ -148,7 +149,7 @@ def main() -> None: series = extract_metric(daily, "net_sales") print(f"{len(series)} days of history.") accuracy = compare_against_baseline( - summarise_backtest(rolling_origin_evaluate(series, horizon=7, initial_train=56, step=14)) + summarise_backtest(rolling_origin_evaluate(series, **BACKTEST)) ) print(accuracy[["model", "folds", "n_observations", "mae", "wape", "wape_improvement_vs_baseline"]].to_string(index=False)) diff --git a/scripts/run_pipeline.py b/scripts/run_pipeline.py new file mode 100644 index 0000000..6e3be8f --- /dev/null +++ b/scripts/run_pipeline.py @@ -0,0 +1,74 @@ +"""Run the Online Retail II pipeline: land, conform in Spark, load, publish. + +Locally, the lake is a folder and Spark runs on this machine: + + docker compose up -d warehouse + docker compose run --rm pipeline + +On AWS the lake is an S3 bucket and the conform step runs as a Glue job; see +scripts/run_pipeline_aws.py. +""" + +from __future__ import annotations + +import argparse +import time + +from rootsignal.console import use_utf8_output +from rootsignal.pipeline import Lake, conform_locally, land, load, publish, read_curated + +STAGES = ("land", "conform", "load", "publish") + + +def main() -> None: + use_utf8_output() + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--lake", default="data/lake", help="Local folder for the raw and curated zones.") + parser.add_argument("--warehouse-url", default=None, help="Defaults to ROOTSIGNAL_WAREHOUSE_URL.") + parser.add_argument("--schema", default="rootsignal") + parser.add_argument("--stages", default=",".join(STAGES), help="Comma-separated subset of land,conform,load,publish.") + args = parser.parse_args() + + stages = [stage.strip() for stage in args.stages.split(",") if stage.strip()] + unknown = sorted(set(stages) - set(STAGES)) + if unknown: + parser.error(f"Unknown stage(s): {unknown}") + + lake = Lake(args.lake) + started = time.time() + + if "land" in stages: + landed = land(lake) + print(f"land {landed['raw_rows']:,} raw lines to {landed['raw_csv']}") + + if "conform" in stages: + summary = conform_locally(lake) + print( + f"conform {summary['source_rows']:,} lines in, {summary['returned_lines']:,} returns held apart, " + f"{summary['consolidated_lines']:,} repeated lines combined" + ) + for name, rows in summary["tables"].items(): + print(f" {name:14s} {rows:>10,}") + + if "load" in stages or "publish" in stages: + dataset, returns = read_curated(lake) + report = load(dataset, url=args.warehouse_url, schema=args.schema, lake=lake) + print( + f"load {report.errors_before_cleaning} validation error(s) before cleaning, " + f"{report.quarantined:,} row(s) quarantined, 0 after" + ) + for name, rows in report.loaded.items(): + if rows: + print(f" {name:14s} {rows:>10,}") + + if "publish" in stages: + published = publish(report.tables, returns, url=args.warehouse_url, schema=args.schema) + print("publish report tables") + for name, rows in published.items(): + print(f" {name:28s} {rows:>6,}") + + print(f"\nDone in {time.time() - started:.0f}s.") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_pipeline_aws.py b/scripts/run_pipeline_aws.py new file mode 100644 index 0000000..4614e86 --- /dev/null +++ b/scripts/run_pipeline_aws.py @@ -0,0 +1,187 @@ +"""Run the pipeline on AWS: S3, Glue, EMR Serverless and Redshift Serverless. + +Each step is a subcommand, so a run can be stopped, inspected and resumed, and +nothing that bills is created without being asked for by name: + + python scripts/run_pipeline_aws.py setup bucket and IAM roles (free) + python scripts/run_pipeline_aws.py land raw zone to S3 (free tier) + python scripts/run_pipeline_aws.py glue conform in Glue (bills) + python scripts/run_pipeline_aws.py emr conform on EMR Serverless (bills) + python scripts/run_pipeline_aws.py load clean, into PostgreSQL and the S3 clean zone + python scripts/run_pipeline_aws.py redshift load Redshift and check it (bills) + python scripts/run_pipeline_aws.py status everything that can bill + python scripts/run_pipeline_aws.py teardown remove everything that can bill + +Credentials come from the AWS CLI configuration, as for any boto3 program. +""" + +from __future__ import annotations + +import argparse +import tempfile +import time + +import numpy as np +import pandas as pd + +from rootsignal.console import use_utf8_output +from rootsignal.pipeline import Lake, land, load, publish, read_curated +from rootsignal.pipeline import aws + +EMR_OUTPUT = "curated_emr/online_retail" + + +def lake_for(bucket: str | None) -> Lake: + return Lake(f"s3://{bucket or aws.default_bucket()}") + + +def roles(bucket: str) -> dict[str, str]: + return { + "glue": aws.ensure_role( + aws.GLUE_ROLE, + ["glue.amazonaws.com"], + bucket, + write=True, + managed=("arn:aws:iam::aws:policy/service-role/AWSGlueServiceRole",), + ), + "emr": aws.ensure_role(aws.EMR_ROLE, ["emr-serverless.amazonaws.com"], bucket, write=True), + "redshift": aws.ensure_role( + aws.REDSHIFT_ROLE, + ["redshift.amazonaws.com", "redshift-serverless.amazonaws.com"], + bucket, + write=False, + ), + } + + +def show_run(run: aws.SparkRun) -> None: + print(f"{run.engine}: {run.state} in {run.seconds:.0f}s, {run.usage}") + print(f" output {run.output}") + if run.error: + print(f" error {run.error}") + + +def compare_frames(left: pd.DataFrame, right: pd.DataFrame, tolerance: dict[str, float]) -> list[str]: + """Columns that differ between two results of the same query.""" + if list(left.columns) != list(right.columns): + return [f"columns {list(left.columns)} vs {list(right.columns)}"] + if len(left) != len(right): + return [f"{len(left)} rows vs {len(right)}"] + problems = [] + left, right = left.reset_index(drop=True), right.reset_index(drop=True) + for column in left.columns: + a, b = left[column], right[column] + numeric_a = pd.to_numeric(a, errors="coerce") + numeric_b = pd.to_numeric(b, errors="coerce") + if numeric_a.notna().sum() == a.notna().sum() and numeric_b.notna().sum() == b.notna().sum(): + x, y = numeric_a.astype(float).to_numpy(), numeric_b.astype(float).to_numpy() + gap = np.abs(np.where(np.isnan(x) & np.isnan(y), 0.0, x - y)) + bad = int((gap > tolerance.get(column, 1e-9)).sum() + (np.isnan(x) != np.isnan(y)).sum()) + else: + bad = int((a.astype(str).to_numpy() != b.astype(str).to_numpy()).sum()) + if bad: + problems.append(f"{column}: {bad} row(s)") + return problems + + +def check_redshift_against_sqlite(schema: str) -> int: + """Every query on the generated sample, Redshift against SQLite.""" + from rootsignal.dataset import load_tables_for_analysis + from rootsignal.sql import available_queries, build_database, run_query_file + + sqlite = build_database(load_tables_for_analysis("data/raw/generated")) + failures = 0 + for name in available_queries(): + problems = compare_frames(aws.query_redshift(schema, name), run_query_file(sqlite, name), {}) + print(f" {name:24s} {'identical' if not problems else 'DIFFERS: ' + '; '.join(problems)}") + failures += bool(problems) + return failures + + +def main() -> None: + use_utf8_output() + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("step", choices=["setup", "land", "glue", "emr", "load", "redshift", "status", "teardown"]) + parser.add_argument("--bucket", default=None, help="Defaults to rootsignal-lake--.") + parser.add_argument("--warehouse-url", default=None, help="PostgreSQL; defaults to ROOTSIGNAL_WAREHOUSE_URL.") + args = parser.parse_args() + started = time.time() + + if args.step == "status": + found = aws.billable_resources() + print("\n".join(found) if found else "Nothing that can bill.") + return + if args.step == "teardown": + removed = aws.teardown() + print("\n".join(f"removed {item}" for item in removed) if removed else "Nothing to remove.") + left = aws.billable_resources() + print("Still present: " + ", ".join(left) if left else "Nothing that can bill is left.") + return + + bucket = args.bucket or aws.default_bucket() + lake = lake_for(bucket) + + if args.step == "setup": + aws.ensure_bucket(bucket) + for name, arn in roles(bucket).items(): + print(f"role {name:8s} {arn}") + print(f"lake {lake.root}") + + elif args.step == "land": + landed = land(lake) + print(f"{landed['raw_rows']:,} raw lines to {landed['raw_csv']}") + + elif args.step == "glue": + show_run(aws.run_glue(lake, roles(bucket)["glue"])) + + elif args.step == "emr": + show_run(aws.run_emr(lake, roles(bucket)["emr"], lake.uri(EMR_OUTPUT))) + # The same script on a second engine should write the same tables. + with tempfile.TemporaryDirectory() as scratch: + glue_dataset, _ = read_curated(lake, into=f"{scratch}/glue") + emr_dataset, _ = read_curated(lake, into=f"{scratch}/emr", zone=EMR_OUTPUT) + for name, frame in glue_dataset.tables.items(): + if frame.empty: + continue + key = list(frame.columns[:1]) if name != "fact_sales" else ["order_id", "sku_id", "sales_type"] + problems = compare_frames( + frame.sort_values(key).reset_index(drop=True), + emr_dataset.tables[name].sort_values(key).reset_index(drop=True), + {"net_sales": 1e-6, "unit_price": 1e-6}, + ) + print(f" EMR vs Glue {name:14s} {'identical' if not problems else 'DIFFERS: ' + '; '.join(problems)}") + + elif args.step == "load": + dataset, returns = read_curated(lake) + report = load(dataset, url=args.warehouse_url, lake=lake) + published = publish(report.tables, returns, url=args.warehouse_url) + print( + f"{report.errors_before_cleaning} validation errors before cleaning, " + f"{report.quarantined:,} rows quarantined, 0 after" + ) + print(f"PostgreSQL: {report.loaded['fact_sales']:,} sales lines, {len(published)} report tables") + print(f"clean zone: {lake.uri('clean/online_retail')}") + + elif args.step == "redshift": + role = roles(bucket)["redshift"] + aws.ensure_redshift(role) + counts = aws.load_redshift(lake, "online_retail") + print("Redshift online_retail: " + ", ".join(f"{k} {v:,}" for k, v in counts.items() if v)) + + # The generated sample fills every fact, so every query has rows to compare. + from rootsignal.dataset import load_tables_for_analysis + from rootsignal.pipeline import write_clean_zone + + sample = load_tables_for_analysis("data/raw/generated") + sample_lake = Lake(f"s3://{bucket}/sample") + write_clean_zone(sample_lake, sample) + aws.load_redshift(sample_lake, "sample") + print("Redshift against SQLite, every analytical query on the generated sample:") + failures = check_redshift_against_sqlite("sample") + print("all identical" if not failures else f"{failures} quer(ies) differ") + + print(f"\n{time.time() - started:.0f}s") + + +if __name__ == "__main__": + main() diff --git a/sql/analytics/daily_sales_tracker.sql b/sql/analytics/daily_sales_tracker.sql index ed384ce..8ce6acf 100644 --- a/sql/analytics/daily_sales_tracker.sql +++ b/sql/analytics/daily_sales_tracker.sql @@ -32,7 +32,7 @@ SELECT d.date, d.net_sales, d.units, - do.orders, + od.orders, CASE WHEN d.ordered_units <> 0 THEN ROUND(CAST(d.fulfilled_units AS REAL) / d.ordered_units, 4) END AS fill_rate, p.sales_target, @@ -42,6 +42,6 @@ SELECT LAG(d.net_sales) OVER (ORDER BY d.date) AS prior_day_net_sales, ROUND(d.net_sales - LAG(d.net_sales) OVER (ORDER BY d.date), 2) AS day_over_day_change FROM daily AS d -JOIN daily_orders AS do ON do.date = d.date +JOIN daily_orders AS od ON od.date = d.date LEFT JOIN plan AS p ON p.date = d.date ORDER BY d.date; diff --git a/sql/analytics/kam_scorecard.sql b/sql/analytics/kam_scorecard.sql index 150db3d..5dc9979 100644 --- a/sql/analytics/kam_scorecard.sql +++ b/sql/analytics/kam_scorecard.sql @@ -44,4 +44,7 @@ FROM dim_kam AS k LEFT JOIN sales AS s ON s.kam_id = k.kam_id LEFT JOIN orders AS o ON o.kam_id = k.kam_id LEFT JOIN quota AS q ON q.kam_id = k.kam_id -ORDER BY sales_attainment; +-- Ties are broken on the segment key, and NULLs placed explicitly: SQLite +-- sorts them first and PostgreSQL last, so leaving either implicit gives the +-- two engines different row orders for the same result. +ORDER BY sales_attainment NULLS FIRST, k.kam_id; diff --git a/sql/analytics/supply_watchlist.sql b/sql/analytics/supply_watchlist.sql index bd23253..2acd8d0 100644 --- a/sql/analytics/supply_watchlist.sql +++ b/sql/analytics/supply_watchlist.sql @@ -23,12 +23,14 @@ WITH weekly AS ( stock AS ( -- Stock is a level and is averaged; stockout_flag is a rate over SKU-days. -- Summing a level would report a week's stock as seven times its own size. + -- Both are cast before averaging: Redshift's AVG of an integer column is an + -- integer, which would report every stockout rate below 1 as 0. SELECT - DATE(i.date, 'weekday 0', '-6 days') AS week_start, + DATE(i.date, 'weekday 0', '-6 days') AS week_start, i.region_code, i.category, - ROUND(AVG(i.available_stock), 2) AS available_stock, - ROUND(AVG(i.stockout_flag), 4) AS stockout_rate + ROUND(AVG(CAST(i.available_stock AS REAL)), 2) AS available_stock, + ROUND(AVG(CAST(i.stockout_flag AS REAL)), 4) AS stockout_rate FROM stg_inventory AS i GROUP BY week_start, i.region_code, i.category ), @@ -51,15 +53,16 @@ joined AS ( AND s.region_code = w.region_code AND s.category = w.category ), +-- The window is written out on each LAG because Redshift has no named WINDOW +-- clause. movement AS ( SELECT *, - LAG(fill_rate) OVER segment AS prior_fill_rate, - LAG(ordered_units) OVER segment AS prior_ordered_units, - LAG(available_stock) OVER segment AS prior_available_stock, - LAG(stockout_rate) OVER segment AS prior_stockout_rate + LAG(fill_rate) OVER (PARTITION BY region_code, category ORDER BY week_start) AS prior_fill_rate, + LAG(ordered_units) OVER (PARTITION BY region_code, category ORDER BY week_start) AS prior_ordered_units, + LAG(available_stock) OVER (PARTITION BY region_code, category ORDER BY week_start) AS prior_available_stock, + LAG(stockout_rate) OVER (PARTITION BY region_code, category ORDER BY week_start) AS prior_stockout_rate FROM joined - WINDOW segment AS (PARTITION BY region_code, category ORDER BY week_start) ) SELECT week_start, @@ -82,4 +85,4 @@ WHERE prior_fill_rate IS NOT NULL AND fill_rate - prior_fill_rate <= -0.05 -- ...while demand did not, so this is not simply a shrinking segment. AND CAST(ordered_units AS REAL) / prior_ordered_units >= 0.95 -ORDER BY fill_rate_change; +ORDER BY fill_rate_change, week_start, region_code, category; diff --git a/sql/analytics/target_variance.sql b/sql/analytics/target_variance.sql index f500840..94fd515 100644 --- a/sql/analytics/target_variance.sql +++ b/sql/analytics/target_variance.sql @@ -38,4 +38,7 @@ LEFT JOIN actual AS a ON a.region_code = p.region_code AND a.category = p.category AND a.channel = p.channel -ORDER BY variance_pct; +-- Ties are broken on the segment key, and NULLs placed explicitly: SQLite +-- sorts them first and PostgreSQL last, so leaving either implicit gives the +-- two engines different row orders for the same result. +ORDER BY variance_pct NULLS FIRST, p.region_code, p.category, p.channel; diff --git a/sql/analytics/weekly_movers.sql b/sql/analytics/weekly_movers.sql index 6a8591d..522c7ed 100644 --- a/sql/analytics/weekly_movers.sql +++ b/sql/analytics/weekly_movers.sql @@ -52,4 +52,4 @@ SELECT AS share_of_absolute_movement, c.days_observed FROM contribution AS c -ORDER BY c.week_start, c.contribution; +ORDER BY c.week_start, c.contribution, c.region_code, c.category; diff --git a/src/rootsignal/adapters/online_retail.py b/src/rootsignal/adapters/online_retail.py index 7b654f0..8a4eeca 100644 --- a/src/rootsignal/adapters/online_retail.py +++ b/src/rootsignal/adapters/online_retail.py @@ -62,6 +62,9 @@ "Roughly a fifth of lines have no customer identifier. Those rows are " "kept with an explicit UNKNOWN customer rather than dropped, because " "discarding a fifth of revenue would misstate the business.", + "Stock codes are upper-cased and trimmed. 172 codes appear in two forms, " + "such as 15056BL and 15056bl, or 47503J with a trailing space, for the " + "same product; kept apart, one product's sales would split across two.", ), ) @@ -137,7 +140,10 @@ def adapt( frame["date"] = pd.to_datetime(frame["timestamp"], errors="coerce").dt.date frame["order_id"] = frame["order_id"].astype(str) - frame["sku_id"] = frame["sku_id"].astype(str) + # The same product appears as 15056BL and 15056bl, and as 47503J with and + # without a trailing space. The code identifies the product; its case and + # padding do not. + frame["sku_id"] = frame["sku_id"].astype(str).str.strip().str.upper() frame["category"] = _derive_category(frame["sku_name"]) # A fifth of lines have no customer. Dropping them would remove a fifth of diff --git a/src/rootsignal/pipeline/__init__.py b/src/rootsignal/pipeline/__init__.py new file mode 100644 index 0000000..92b9448 --- /dev/null +++ b/src/rootsignal/pipeline/__init__.py @@ -0,0 +1,35 @@ +"""From the published source to the warehouse: land, conform in Spark, load, publish.""" + +from .lake import CLEAN, CURATED, JOBS, RAW, Lake +from .stages import ( + LoadReport, + clean, + conform_locally, + land, + load, + publish, + publish_job, + read_curated, + read_summary, + write_clean_zone, + write_raw_csv, +) + +__all__ = [ + "CLEAN", + "CURATED", + "JOBS", + "RAW", + "Lake", + "LoadReport", + "clean", + "conform_locally", + "land", + "load", + "publish", + "publish_job", + "read_curated", + "read_summary", + "write_clean_zone", + "write_raw_csv", +] diff --git a/src/rootsignal/pipeline/aws.py b/src/rootsignal/pipeline/aws.py new file mode 100644 index 0000000..7cdd9ef --- /dev/null +++ b/src/rootsignal/pipeline/aws.py @@ -0,0 +1,464 @@ +"""The same pipeline on AWS: S3 for the lake, Glue and EMR Serverless for Spark, +Redshift Serverless as the warehouse. + +Nothing here changes what the pipeline computes. The Spark script is the one +tests/test_spark_conform.py holds to the pandas adapter, uploaded unchanged; +the SQL is the repository's, translated by rootsignal.sql.dialect.to_redshift. + +Everything that bills is created with a ceiling and removed by teardown(): + + Glue 2 workers, a 20-minute timeout + EMR Serverless capped at 8 vCPU, stops itself after 2 idle minutes + Redshift 8 RPU, with a usage limit that switches the workgroup off + after 4 RPU-hours in a day + +The bucket and the IAM roles cost nothing to keep and are left in place. +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass + +import pandas as pd + +from ..sql.database import ANALYTICS_DIR, LOAD_ORDER, MARTS_DIR, SCHEMA_PATH, STAGING_DIR, _sql_files +from ..sql.dialect import _strip_comments, to_redshift +from .lake import CLEAN, CURATED, RAW, Lake +from .stages import RAW_CSV, publish_job + +PREFIX = "rootsignal" +GLUE_JOB = f"{PREFIX}-conform" +EMR_APP = f"{PREFIX}-conform" +NAMESPACE = f"{PREFIX}-ns" +WORKGROUP = f"{PREFIX}-wg" +DATABASE = "dev" + +GLUE_ROLE = f"{PREFIX}-glue" +EMR_ROLE = f"{PREFIX}-emr-serverless" +REDSHIFT_ROLE = f"{PREFIX}-redshift" + + +def _client(service: str): + import boto3 + + return boto3.client(service) + + +def account_id() -> str: + return _client("sts").get_caller_identity()["Account"] + + +def region() -> str: + import boto3 + + return boto3.session.Session().region_name + + +def default_bucket() -> str: + return f"{PREFIX}-lake-{account_id()}-{region()}" + + +# -------------------------------------------------------------------------- +# Storage and access +# -------------------------------------------------------------------------- + + +def ensure_bucket(name: str) -> str: + """A private bucket for the lake. Public access is blocked outright.""" + s3 = _client("s3") + try: + s3.head_bucket(Bucket=name) + except s3.exceptions.ClientError: + s3.create_bucket(Bucket=name, CreateBucketConfiguration={"LocationConstraint": region()}) + s3.put_public_access_block( + Bucket=name, + PublicAccessBlockConfiguration={ + "BlockPublicAcls": True, + "IgnorePublicAcls": True, + "BlockPublicPolicy": True, + "RestrictPublicBuckets": True, + }, + ) + return name + + +def _lake_policy(bucket: str, write: bool) -> dict: + actions = ["s3:GetObject", "s3:ListBucket", "s3:GetBucketLocation"] + if write: + actions += ["s3:PutObject", "s3:DeleteObject"] + return { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": actions, + "Resource": [f"arn:aws:s3:::{bucket}", f"arn:aws:s3:::{bucket}/*"], + } + ], + } + + +def ensure_role( + name: str, services: list[str], bucket: str, write: bool, managed: tuple[str, ...] = () +) -> str: + """A role the services can assume, allowed into this bucket and nowhere else.""" + iam = _client("iam") + trust = { + "Version": "2012-10-17", + "Statement": [{"Effect": "Allow", "Principal": {"Service": services}, "Action": "sts:AssumeRole"}], + } + try: + arn = iam.get_role(RoleName=name)["Role"]["Arn"] + created = False + except iam.exceptions.NoSuchEntityException: + arn = iam.create_role(RoleName=name, AssumeRolePolicyDocument=json.dumps(trust))["Role"]["Arn"] + created = True + iam.put_role_policy( + RoleName=name, PolicyName=f"{PREFIX}-lake", PolicyDocument=json.dumps(_lake_policy(bucket, write)) + ) + for policy in managed: + iam.attach_role_policy(RoleName=name, PolicyArn=policy) + if created: + # A new role takes a few seconds to become assumable everywhere. + time.sleep(15) + return arn + + +# -------------------------------------------------------------------------- +# Spark: Glue and EMR Serverless +# -------------------------------------------------------------------------- + + +@dataclass +class SparkRun: + engine: str + state: str + seconds: float + usage: str + output: str + error: str = "" + + +def run_glue(lake: Lake, role_arn: str, output: str | None = None, workers: int = 2) -> SparkRun: + """Run the conform script as a Glue Spark job and wait for it.""" + glue = _client("glue") + output = output or lake.uri(CURATED) + job = { + "Role": role_arn, + "Command": {"Name": "glueetl", "ScriptLocation": publish_job(lake), "PythonVersion": "3"}, + "GlueVersion": "5.0", + "WorkerType": "G.1X", + "NumberOfWorkers": workers, + "Timeout": 20, + "MaxRetries": 0, + "DefaultArguments": {"--input": lake.uri(RAW, RAW_CSV), "--output": output}, + } + try: + glue.get_job(JobName=GLUE_JOB) + glue.update_job(JobName=GLUE_JOB, JobUpdate=job) + except glue.exceptions.EntityNotFoundException: + glue.create_job(Name=GLUE_JOB, **job) + + run_id = glue.start_job_run(JobName=GLUE_JOB)["JobRunId"] + while True: + run = glue.get_job_run(JobName=GLUE_JOB, RunId=run_id)["JobRun"] + if run["JobRunState"] in {"SUCCEEDED", "FAILED", "ERROR", "TIMEOUT", "STOPPED"}: + break + time.sleep(15) + seconds = float(run.get("ExecutionTime", 0)) + # Glue bills per second with a one-minute minimum, per worker (1 DPU each on G.1X). + dpu_hours = max(seconds, 60) / 3600 * workers + return SparkRun( + engine="Glue 5.0", + state=run["JobRunState"], + seconds=seconds, + usage=f"{dpu_hours:.3f} DPU-hours", + output=output, + error=run.get("ErrorMessage", ""), + ) + + +def ensure_emr_application() -> str: + """An EMR Serverless Spark application with a hard capacity ceiling.""" + emr = _client("emr-serverless") + for app in emr.list_applications()["applications"]: + if app["name"] == EMR_APP and app["state"] not in {"TERMINATED"}: + return app["id"] + return emr.create_application( + name=EMR_APP, + releaseLabel="emr-7.5.0", + type="SPARK", + maximumCapacity={"cpu": "8 vCPU", "memory": "32 GB", "disk": "100 GB"}, + autoStartConfiguration={"enabled": True}, + autoStopConfiguration={"enabled": True, "idleTimeoutMinutes": 2}, + )["applicationId"] + + +def run_emr(lake: Lake, role_arn: str, output: str) -> SparkRun: + """Run the same conform script on EMR Serverless and wait for it.""" + emr = _client("emr-serverless") + app_id = ensure_emr_application() + submit = " ".join( + [ + "--conf spark.dynamicAllocation.enabled=false", + "--conf spark.executor.instances=2", + "--conf spark.executor.cores=1", + "--conf spark.executor.memory=2g", + "--conf spark.driver.cores=1", + "--conf spark.driver.memory=2g", + ] + ) + job_id = emr.start_job_run( + applicationId=app_id, + executionRoleArn=role_arn, + name="conform-online-retail", + jobDriver={ + "sparkSubmit": { + "entryPoint": publish_job(lake), + "entryPointArguments": ["--input", lake.uri(RAW, RAW_CSV), "--output", output], + "sparkSubmitParameters": submit, + } + }, + configurationOverrides={ + "monitoringConfiguration": {"s3MonitoringConfiguration": {"logUri": lake.uri("logs", "emr")}} + }, + executionTimeoutMinutes=20, + )["jobRunId"] + while True: + run = emr.get_job_run(applicationId=app_id, jobRunId=job_id)["jobRun"] + if run["state"] in {"SUCCESS", "FAILED", "CANCELLED"}: + break + time.sleep(15) + used = run.get("totalResourceUtilization", {}) + return SparkRun( + engine=f"EMR Serverless {run.get('releaseLabel', 'emr-7.5.0')}", + state=run["state"], + seconds=float(run.get("totalExecutionDurationSeconds", 0)), + usage=f"{used.get('vCPUHour', 0):.3f} vCPU-hours, {used.get('memoryGBHour', 0):.3f} GB-hours", + output=output, + error=run.get("stateDetails", ""), + ) + + +# -------------------------------------------------------------------------- +# Redshift Serverless +# -------------------------------------------------------------------------- + + +def _default_network() -> tuple[list[str], list[str]]: + ec2 = _client("ec2") + vpc = ec2.describe_vpcs(Filters=[{"Name": "isDefault", "Values": ["true"]}])["Vpcs"][0]["VpcId"] + subnets = [ + subnet["SubnetId"] + for subnet in ec2.describe_subnets(Filters=[{"Name": "vpc-id", "Values": [vpc]}])["Subnets"] + ] + group = ec2.describe_security_groups( + Filters=[{"Name": "vpc-id", "Values": [vpc]}, {"Name": "group-name", "Values": ["default"]}] + )["SecurityGroups"][0]["GroupId"] + return subnets, [group] + + +def ensure_redshift(role_arn: str, base_capacity: int = 8, daily_rpu_hours: int = 4) -> str: + """A Redshift Serverless workgroup at the smallest size, with a spending cap. + + The usage limit deactivates the workgroup once it has used daily_rpu_hours + of compute in a day, so a forgotten workgroup cannot run up a bill. + """ + rs = _client("redshift-serverless") + try: + rs.get_namespace(namespaceName=NAMESPACE) + except rs.exceptions.ResourceNotFoundException: + rs.create_namespace( + namespaceName=NAMESPACE, dbName=DATABASE, iamRoles=[role_arn], defaultIamRoleArn=role_arn + ) + try: + workgroup = rs.get_workgroup(workgroupName=WORKGROUP)["workgroup"] + except rs.exceptions.ResourceNotFoundException: + subnets, groups = _default_network() + workgroup = rs.create_workgroup( + workgroupName=WORKGROUP, + namespaceName=NAMESPACE, + baseCapacity=base_capacity, + maxCapacity=base_capacity, + publiclyAccessible=False, + subnetIds=subnets, + securityGroupIds=groups, + )["workgroup"] + while workgroup["status"] != "AVAILABLE": + time.sleep(15) + workgroup = rs.get_workgroup(workgroupName=WORKGROUP)["workgroup"] + + limits = rs.list_usage_limits(resourceArn=workgroup["workgroupArn"]).get("usageLimits", []) + if not limits: + rs.create_usage_limit( + resourceArn=workgroup["workgroupArn"], + usageType="serverless-compute", + amount=daily_rpu_hours, + period="daily", + breachAction="deactivate", + ) + return workgroup["workgroupArn"] + + +def execute(statements: list[str], timeout: int = 900) -> list[str]: + """Run statements in one transaction through the Redshift Data API.""" + data = _client("redshift-data") + statements = [sql for sql in statements if sql.strip()] + if len(statements) == 1: + response = data.execute_statement(WorkgroupName=WORKGROUP, Database=DATABASE, Sql=statements[0]) + else: + response = data.batch_execute_statement(WorkgroupName=WORKGROUP, Database=DATABASE, Sqls=statements) + started = time.time() + while True: + described = data.describe_statement(Id=response["Id"]) + if described["Status"] in {"FINISHED", "FAILED", "ABORTED"}: + break + if time.time() - started > timeout: + data.cancel_statement(Id=response["Id"]) + raise TimeoutError(f"Redshift statement did not finish in {timeout}s") + time.sleep(2) + if described["Status"] != "FINISHED": + raise RuntimeError(f"Redshift: {described.get('Error', described['Status'])}") + if len(statements) == 1: + return [response["Id"]] + return [part["Id"] for part in described["SubStatements"]] + + +def fetch(sql: str, schema: str | None = None) -> pd.DataFrame: + """A SELECT through the Data API, as a frame. + + Each Data API call is its own session, so a search_path set in one call is + gone by the next. With a schema, the SET and the SELECT go in one batch and + the SELECT's result is read. + """ + data = _client("redshift-data") + statements = [f"SET search_path TO {schema}", sql] if schema else [sql] + statement_id = execute(statements)[-1] + columns, rows, token = None, [], None + while True: + kwargs = {"Id": statement_id, **({"NextToken": token} if token else {})} + page = data.get_statement_result(**kwargs) + columns = columns or [column["name"] for column in page["ColumnMetadata"]] + for record in page["Records"]: + rows.append([None if field.get("isNull") else next(iter(field.values())) for field in record]) + token = page.get("NextToken") + if not token: + break + return pd.DataFrame(rows, columns=columns) + + +def statements_in(sql: str) -> list[str]: + """The statements in a SQL file, comments removed first. + + Splitting on semicolons before removing comments would cut a statement at a + semicolon inside a comment, and schema.sql's second line has one. + """ + return [part.strip() for part in _strip_comments(sql).split(";") if part.strip()] + + +def load_redshift(lake: Lake, schema: str) -> dict[str, int]: + """Create the schema, COPY the clean zone into it, and register the views.""" + statements = [ + f"DROP SCHEMA IF EXISTS {schema} CASCADE", + f"CREATE SCHEMA {schema}", + f"SET search_path TO {schema}", + *statements_in(to_redshift(SCHEMA_PATH.read_text(encoding="utf-8"))), + ] + for name in LOAD_ORDER: + statements.append(f"COPY {name} FROM '{lake.uri(CLEAN, name)}/' IAM_ROLE default FORMAT AS PARQUET") + execute(statements) + + views = [f"SET search_path TO {schema}"] + for directory in (STAGING_DIR, MARTS_DIR): + for path in _sql_files(directory): + views += statements_in(to_redshift(path.read_text(encoding="utf-8"))) + execute(views) + + counts = fetch(" UNION ALL ".join(f"SELECT '{name}' AS name, COUNT(*) AS n FROM {name}" for name in LOAD_ORDER), schema) + return dict(zip(counts["name"], counts["n"].astype(int), strict=True)) + + +def query_redshift(schema: str, name: str) -> pd.DataFrame: + """One of the analytical queries, run inside `schema`.""" + (sql,) = statements_in(to_redshift((ANALYTICS_DIR / f"{name}.sql").read_text(encoding="utf-8"))) + return fetch(sql, schema) + + +# -------------------------------------------------------------------------- +# What is running, and removing it +# -------------------------------------------------------------------------- + + +def billable_resources() -> list[str]: + """Everything this module creates that can bill while it exists.""" + found = [] + glue = _client("glue") + try: + glue.get_job(JobName=GLUE_JOB) + found.append(f"Glue job {GLUE_JOB} (bills only while running)") + except glue.exceptions.EntityNotFoundException: + pass + for app in _client("emr-serverless").list_applications()["applications"]: + if app["name"] == EMR_APP and app["state"] != "TERMINATED": + found.append(f"EMR Serverless application {app['id']} ({app['state']})") + rs = _client("redshift-serverless") + for workgroup in rs.list_workgroups()["workgroups"]: + if workgroup["workgroupName"] == WORKGROUP: + found.append(f"Redshift Serverless workgroup {WORKGROUP} ({workgroup['status']})") + for namespace in rs.list_namespaces()["namespaces"]: + if namespace["namespaceName"] == NAMESPACE: + found.append(f"Redshift Serverless namespace {NAMESPACE} (storage)") + return found + + +def teardown() -> list[str]: + """Remove every billable resource. The bucket and roles are kept.""" + removed = [] + glue = _client("glue") + # DeleteJob succeeds whether or not the job exists, so asking first is the + # only way to report a removal that happened rather than one that did not. + try: + glue.get_job(JobName=GLUE_JOB) + glue.delete_job(JobName=GLUE_JOB) + removed.append(f"Glue job {GLUE_JOB}") + except glue.exceptions.EntityNotFoundException: + pass + + emr = _client("emr-serverless") + for app in emr.list_applications()["applications"]: + if app["name"] != EMR_APP or app["state"] == "TERMINATED": + continue + if app["state"] not in {"STOPPED", "CREATED"}: + emr.stop_application(applicationId=app["id"]) + while emr.get_application(applicationId=app["id"])["application"]["state"] not in {"STOPPED"}: + time.sleep(10) + emr.delete_application(applicationId=app["id"]) + removed.append(f"EMR Serverless application {app['id']}") + + rs = _client("redshift-serverless") + try: + rs.get_workgroup(workgroupName=WORKGROUP) + rs.delete_workgroup(workgroupName=WORKGROUP) + while True: + try: + rs.get_workgroup(workgroupName=WORKGROUP) + time.sleep(15) + except rs.exceptions.ResourceNotFoundException: + break + removed.append(f"Redshift Serverless workgroup {WORKGROUP}") + except rs.exceptions.ResourceNotFoundException: + pass + try: + rs.delete_namespace(namespaceName=NAMESPACE) + while True: + try: + rs.get_namespace(namespaceName=NAMESPACE) + time.sleep(15) + except rs.exceptions.ResourceNotFoundException: + break + removed.append(f"Redshift Serverless namespace {NAMESPACE}") + except rs.exceptions.ResourceNotFoundException: + pass + return removed diff --git a/src/rootsignal/pipeline/lake.py b/src/rootsignal/pipeline/lake.py new file mode 100644 index 0000000..fd20bcc --- /dev/null +++ b/src/rootsignal/pipeline/lake.py @@ -0,0 +1,103 @@ +"""Where the pipeline keeps its files: a local folder, or an S3 bucket. + +Both are laid out the same way, one prefix per zone: + + raw/online_retail/ the source as published, plus a CSV Spark can read + curated/online_retail/ one Parquet dataset per conformed table + clean/online_retail/ the cleaned model, typed as schema.sql, for Redshift + jobs/ the Spark script, for Glue or EMR to fetch + +The rest of the pipeline only asks for a zone's URI, puts a file into one, or +fetches one to local disk, so a run on a laptop and a run against S3 differ in +the root they were given and nothing else. +""" + +from __future__ import annotations + +import shutil +from dataclasses import dataclass +from pathlib import Path + +RAW = "raw/online_retail" +CURATED = "curated/online_retail" +CLEAN = "clean/online_retail" +JOBS = "jobs" + + +@dataclass(frozen=True) +class Lake: + root: str + + @property + def is_s3(self) -> bool: + return self.root.startswith("s3://") + + @property + def bucket(self) -> str: + return self.root.removeprefix("s3://").split("/", 1)[0] + + @property + def prefix(self) -> str: + parts = self.root.removeprefix("s3://").split("/", 1) + return parts[1].strip("/") if len(parts) > 1 else "" + + def key(self, *parts: str) -> str: + return "/".join(part.strip("/") for part in (self.prefix, *parts) if part) + + def uri(self, *parts: str) -> str: + """The location Spark reads or writes: s3:// on AWS, a path locally.""" + if self.is_s3: + return f"s3://{self.bucket}/{self.key(*parts)}" + return str(Path(self.root, *parts)) + + def _client(self): + import boto3 + + return boto3.client("s3") + + def put(self, local: str | Path, *parts: str) -> str: + """Copy a local file into the lake and return its URI.""" + if self.is_s3: + self._client().upload_file(str(local), self.bucket, self.key(*parts)) + else: + target = Path(self.root, *parts) + target.parent.mkdir(parents=True, exist_ok=True) + if Path(local).resolve() != target.resolve(): + shutil.copyfile(local, target) + return self.uri(*parts) + + def _objects(self, *parts: str): + """Keys at exactly this path, or under it as a folder. + + A bare prefix would also match siblings: fact_sales would pick up + fact_sales_backup. + """ + base = self.key(*parts) + paginator = self._client().get_paginator("list_objects_v2") + for page in paginator.paginate(Bucket=self.bucket, Prefix=base): + for item in page.get("Contents", []): + if item["Key"] == base or item["Key"].startswith(base + "/"): + yield item["Key"] + + def exists(self, *parts: str) -> bool: + if not self.is_s3: + return Path(self.root, *parts).exists() + return next(self._objects(*parts), None) is not None + + def fetch(self, *parts: str, into: str | Path) -> Path: + """A local copy of a file or folder. Local lakes are read in place.""" + if not self.is_s3: + return Path(self.root, *parts) + client = self._client() + base = self.key(*parts) + target = Path(into, *parts) + found = False + for key in self._objects(*parts): + relative = key[len(base):].lstrip("/") + destination = target / relative if relative else target + destination.parent.mkdir(parents=True, exist_ok=True) + client.download_file(self.bucket, key, str(destination)) + found = True + if not found: + raise FileNotFoundError(f"Nothing in the lake at {self.uri(*parts)}") + return target diff --git a/src/rootsignal/pipeline/publish.py b/src/rootsignal/pipeline/publish.py new file mode 100644 index 0000000..7aa5326 --- /dev/null +++ b/src/rootsignal/pipeline/publish.py @@ -0,0 +1,87 @@ +"""The analyses the warehouse serves alongside the model. + +The fact and dimension tables let a BI tool slice revenue, but the findings this +project exists for, which segments moved a metric and whether a rate moved +because segments changed or because the mix did, come from the Python layers. +These are computed here once, with the same calls and parameters as +scripts/run_external_dataset.py, and written to the warehouse as report tables, +so a dashboard reads the project's results rather than re-deriving them. +""" + +from __future__ import annotations + +import pandas as pd + +from ..adapters.online_retail import returns_by_period +from ..analysis import calculate_trend, summarise_by_period +from ..decomposition import component_coherence, decompose_additive, decompose_movement +from ..forecasting import ( + build_daily_series, + compare_against_baseline, + extract_metric, + rolling_origin_evaluate, + summarise_backtest, +) + + +# The backtest behind the README's forecast figure: a week ahead, eight weeks of +# history before the first fold, a new fold every fortnight. On Online Retail II +# that is 49 folds. scripts/run_external_dataset.py uses the same settings, so +# the report table and the published number cannot drift apart. +BACKTEST = {"horizon": 7, "initial_train": 56, "step": 14} + + +def analysis_tables(tables: dict[str, pd.DataFrame], returns: pd.DataFrame) -> dict[str, pd.DataFrame]: + """Every report table, keyed by the name it is published under.""" + sales = tables["fact_sales"] + + monthly = summarise_by_period(tables, period="month") + trend = calculate_trend(monthly, "net_sales") + + by_country = summarise_by_period(tables, period="month", group_by=["region_code"]) + contribution = decompose_additive(by_country, "net_sales", ["region_code"]) + + # Countries with almost no volume swing between extremes; min_share folds + # them into one bucket rather than dropping volume, as the script does. + returns_monthly = returns_by_period(sales, returns, period="month", group_by=["region_code"]) + return_split = decompose_movement(returns_monthly, "return_rate", ["region_code"], min_share=0.01) + + daily = build_daily_series(sales, metrics=["net_sales", "units"]) + accuracy = compare_against_baseline( + summarise_backtest( + rolling_origin_evaluate(extract_metric(daily, "net_sales"), **BACKTEST) + ) + ) + + return { + "rpt_monthly_trade": monthly, + "rpt_monthly_trend": trend, + "rpt_country_contribution": contribution, + "rpt_returns_monthly": returns_monthly, + "rpt_return_rate_split": return_split, + "rpt_return_rate_components": component_coherence(return_split), + "rpt_forecast_accuracy": accuracy, + } + + +def flatten(frame: pd.DataFrame) -> pd.DataFrame: + """A frame a SQL table can hold: plain columns, no index, no Python objects. + + Periods and timestamps become dates, and anything else that is not a + number, a boolean or a date is written as text. + """ + out = frame.reset_index(drop=True).copy() + out.columns = [str(column) for column in out.columns] + for column in out.columns: + series = out[column] + if isinstance(series.dtype, pd.PeriodDtype): + out[column] = series.dt.start_time.dt.date + elif pd.api.types.is_datetime64_any_dtype(series): + out[column] = series.dt.date + elif pd.api.types.is_bool_dtype(series) or pd.api.types.is_numeric_dtype(series): + continue + elif series.map(lambda value: value is None or hasattr(value, "isoformat")).all(): + continue + else: + out[column] = series.map(lambda value: None if value is None or value != value else str(value)) + return out diff --git a/src/rootsignal/pipeline/stages.py b/src/rootsignal/pipeline/stages.py new file mode 100644 index 0000000..3550992 --- /dev/null +++ b/src/rootsignal/pipeline/stages.py @@ -0,0 +1,238 @@ +"""The stages that take Online Retail II from its source to the warehouse. + + land the published workbook into the raw zone, with a CSV rendering of + it, because Spark has no Excel reader + conform jobs/conform_online_retail.py, in Spark: raw zone to curated zone + load curated Parquet through the existing cleaning and validation, into + PostgreSQL and the clean zone; nothing is loaded if validation + still finds an error + publish the analysis results, written to the warehouse as report tables + +Cleaning stays in pandas on purpose. It is the part of the project with the +audit trail and the quarantine, it runs in seconds on the conformed tables, and +rewriting it for Spark would create a second definition of every rule. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +import tempfile +from dataclasses import dataclass, field +from pathlib import Path + +import pandas as pd + +from ..adapters.base import AdaptedDataset, DatasetCapabilities, empty_fact +from ..adapters.online_retail import ARCHIVE_NAME, CAPABILITIES, DEFAULT_CACHE, download, read_raw +from .lake import CLEAN, CURATED, JOBS, RAW, Lake + +JOB_PATH = Path(__file__).resolve().parents[3] / "jobs" / "conform_online_retail.py" +RAW_CSV = "invoices.csv" +CURATED_TABLES = ("fact_sales", "dim_date", "dim_sku", "dim_customer", "dim_region", "dim_kam") +EMPTY_FACTS = ("fact_orders", "fact_inventory", "fact_targets", "fact_kam_targets") +DATE_COLUMNS = {"fact_sales": ["date"], "dim_date": ["date"], "returns": ["date"]} + + +def write_raw_csv(raw: pd.DataFrame, path: str | Path) -> Path: + """The workbook's rows as CSV, values unchanged. + + Numbers are written at full precision and a missing value as an empty + field, which the Spark job reads back as null. Nothing is cleaned here. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + raw.to_csv(path, index=False) + return path + + +def land(lake: Lake, cache_dir: str | Path = DEFAULT_CACHE) -> dict: + """Put the source into the raw zone: the archive as published, and a CSV.""" + archive = download(cache_dir) + raw = read_raw(cache_dir) + lake.put(archive, RAW, ARCHIVE_NAME) + if lake.is_s3: + with tempfile.TemporaryDirectory() as scratch: + lake.put(write_raw_csv(raw, Path(scratch, RAW_CSV)), RAW, RAW_CSV) + else: + write_raw_csv(raw, Path(lake.uri(RAW, RAW_CSV))) + return {"raw_rows": len(raw), "raw_csv": lake.uri(RAW, RAW_CSV)} + + +def publish_job(lake: Lake) -> str: + """Put the Spark script where Glue or EMR can fetch it.""" + return lake.put(JOB_PATH, JOBS, JOB_PATH.name) + + +def conform_locally(lake: Lake) -> dict: + """Run the Spark job on this machine against a local lake.""" + if lake.is_s3: + raise ValueError("Run the conform stage on AWS for an S3 lake; see rootsignal.pipeline.aws.") + completed = subprocess.run( + [sys.executable, str(JOB_PATH), "--input", lake.uri(RAW, RAW_CSV), "--output", lake.uri(CURATED)], + capture_output=True, + text=True, + ) + if completed.returncode != 0: + raise RuntimeError(f"The Spark job failed:\n{completed.stderr[-4000:]}") + return read_summary(lake) + + +def read_summary(lake: Lake, into: str | Path | None = None) -> dict: + """The job's own record of what it read and wrote.""" + with tempfile.TemporaryDirectory() as scratch: + folder = lake.fetch(CURATED, "_summary", into=into or scratch) + text = "".join(path.read_text(encoding="utf-8") for path in sorted(folder.glob("part-*"))) + return json.loads(text) + + +def _read_table(folder: Path, name: str) -> pd.DataFrame: + frame = pd.read_parquet(folder / name) + for column in DATE_COLUMNS.get(name, []): + frame[column] = pd.to_datetime(frame[column]).dt.date + return frame + + +def read_curated( + lake: Lake, into: str | Path | None = None, zone: str = CURATED +) -> tuple[AdaptedDataset, pd.DataFrame]: + """The conformed tables as the pandas adapter would return them, and the returns. + + The facts this dataset cannot supply are added empty, as the adapter adds + them, so cleaning and validation see the same shape either way. `zone` reads + a Spark run written somewhere other than the curated zone. + """ + scratch = tempfile.TemporaryDirectory() if into is None else None + try: + folder = lake.fetch(zone, into=into or scratch.name) + tables = {name: _read_table(folder, name) for name in CURATED_TABLES} + returns = _read_table(folder, "returns") + summary = json.loads( + "".join(path.read_text(encoding="utf-8") for path in sorted((folder / "_summary").glob("part-*"))) + ) + finally: + if scratch is not None: + scratch.cleanup() + + for name in EMPTY_FACTS: + tables[name] = empty_fact(name) + + notes = CAPABILITIES.notes + ( + f"{summary['returned_lines']:,} of {summary['source_rows']:,} lines are returns and are " + "held separately from sales.", + f"{summary['consolidated_lines']:,} invoice lines repeat a product already on the same " + "invoice and are consolidated into one, adding units and weighting the price by quantity.", + "Conformed in Spark by jobs/conform_online_retail.py.", + ) + dataset = AdaptedDataset( + tables=tables, + capabilities=DatasetCapabilities( + name=CAPABILITIES.name, available_facts=CAPABILITIES.available_facts, notes=notes + ), + source_rows=summary["source_rows"], + ) + return dataset, returns + + +@dataclass +class LoadReport: + loaded: dict[str, int] + quarantined: int + errors_before_cleaning: int + audit: pd.DataFrame = field(repr=False) + tables: dict[str, pd.DataFrame] = field(repr=False, default_factory=dict) + + +def clean(dataset: AdaptedDataset): + """Clean and validate. Any error left after cleaning stops the pipeline.""" + from ..cleaning import clean_dataset + from ..validation import DatasetValidator + + before = DatasetValidator().validate(dataset.tables) + result = clean_dataset(dataset.tables) + after = DatasetValidator().validate(result.tables) + if after.errors(): + problems = "; ".join(f"{issue.table}.{issue.check}: {issue.message}" for issue in after.errors()) + raise ValueError(f"Not loading: validation still fails after cleaning. {problems}") + return result, len(before.errors()) + + +def write_clean_zone(lake: Lake, tables: dict[str, pd.DataFrame]) -> dict[str, str]: + """The cleaned model as Parquet, typed exactly as schema.sql declares it. + + This is what Redshift loads with COPY, so the types have to match the + tables: 32-bit integers for INTEGER, dates for DATE, 0/1 for flags. + """ + import pyarrow as pa + import pyarrow.parquet as pq + + from ..sql.database import LOAD_ORDER, SCHEMA_PATH + from ..sql.dialect import schema_columns + + arrow_types = {"TEXT": pa.string(), "INTEGER": pa.int32(), "REAL": pa.float64(), "DATE": pa.date32()} + declared = schema_columns(SCHEMA_PATH.read_text(encoding="utf-8")) + written = {} + with tempfile.TemporaryDirectory() as scratch: + for name in LOAD_ORDER: + frame = tables[name] + columns = declared[name] + arrays = [] + for column, kind in columns: + values = frame[column] + if kind == "INTEGER": + values = values.astype("int64") + arrays.append(pa.array(values.tolist(), type=arrow_types[kind])) + table = pa.Table.from_arrays(arrays, names=[column for column, _ in columns]) + local = Path(scratch, f"{name}.parquet") + pq.write_table(table, local) + written[name] = lake.put(local, CLEAN, name, "part-00000.parquet") + return written + + +def load( + dataset: AdaptedDataset, + url: str | None = None, + schema: str = "rootsignal", + lake: Lake | None = None, +) -> LoadReport: + """Clean, validate, and replace the warehouse schema with the result. + + Validation runs again after cleaning, and any remaining error stops the load + before the warehouse is touched. With a lake, the cleaned tables are also + written to its clean zone for Redshift. + """ + from ..sql.warehouse import build_warehouse + + result, errors_before = clean(dataset) + connection = build_warehouse(result.tables, url=url, schema=schema) + connection.close() + if lake is not None: + write_clean_zone(lake, result.tables) + return LoadReport( + loaded={name: len(frame) for name, frame in result.tables.items()}, + quarantined=sum(len(frame) for frame in result.quarantined.values()), + errors_before_cleaning=errors_before, + audit=result.audit, + tables=result.tables, + ) + + +def publish( + tables: dict[str, pd.DataFrame], + returns: pd.DataFrame, + url: str | None = None, + schema: str = "rootsignal", +) -> dict[str, int]: + """Write the analysis results into the warehouse as report tables.""" + from ..sql.warehouse import connect, write_report_table + from .publish import analysis_tables, flatten + + connection = connect(url, schema) + try: + return { + name: write_report_table(connection, name, flatten(frame)) + for name, frame in analysis_tables(tables, returns).items() + } + finally: + connection.close() diff --git a/src/rootsignal/sql/dialect.py b/src/rootsignal/sql/dialect.py new file mode 100644 index 0000000..3bcff16 --- /dev/null +++ b/src/rootsignal/sql/dialect.py @@ -0,0 +1,125 @@ +"""Run the one set of SQL files on PostgreSQL as well as SQLite. + +The staging views, the mart and the analytical queries are written once. Where +the two engines disagree, this module rewrites the SQLite form into the +PostgreSQL one, and nothing else: every rule below exists because a specific +construct in sql/ needs it, and a construct with no rule is refused rather than +passed through to fail somewhere less obvious. + +Rounding is the one difference not handled by rewriting. PostgreSQL has no +ROUND for double precision with a number of places, so the warehouse defines +one (see ROUND_FUNCTION) rather than every ROUND call in the repository being +wrapped in a cast. +""" + +from __future__ import annotations + +import re + +# Week starting Monday. SQLite moves forward to the next Sunday (or stays on a +# Sunday) and steps back six days; DATE_TRUNC('week') is the ISO Monday, which +# is the same day. +_WEEK_START = re.compile( + r"DATE\(\s*([\w.]+)\s*,\s*'weekday 0'\s*,\s*'-6 days'\s*\)", re.IGNORECASE +) +_YEAR_MONTH = re.compile(r"STRFTIME\(\s*'%Y-%m'\s*,\s*([\w.]+)\s*\)", re.IGNORECASE) +_PRAGMA = re.compile(r"^\s*PRAGMA\b[^;]*;\s*$", re.IGNORECASE | re.MULTILINE) +# SQLite's REAL is an 8-byte float. PostgreSQL's REAL is 4 bytes, which would +# lose cents on large revenue figures, so it becomes DOUBLE PRECISION. +_REAL = re.compile(r"\bREAL\b") + +# Anything still SQLite-specific after translation. +_UNTRANSLATED = ( + re.compile(r"\bSTRFTIME\s*\(", re.IGNORECASE), + re.compile(r"'weekday\s+\d'", re.IGNORECASE), + re.compile(r"\bPRAGMA\b", re.IGNORECASE), + re.compile(r"\bJULIANDAY\s*\(", re.IGNORECASE), + re.compile(r"\bIFNULL\s*\(", re.IGNORECASE), +) + +ROUND_FUNCTION = """ +CREATE OR REPLACE FUNCTION round(value DOUBLE PRECISION, places INTEGER) +RETURNS DOUBLE PRECISION +LANGUAGE SQL IMMUTABLE STRICT +AS $$ SELECT round(value::numeric, places)::double precision $$; +""" + + +def _strip_comments(sql: str) -> str: + return re.sub(r"--[^\n]*", "", sql) + + +_CREATE_TABLE = re.compile(r"CREATE TABLE\s+(\w+)\s*\((.*?)\n\);", re.IGNORECASE | re.DOTALL) +_COLUMN = re.compile(r"^\s*(\w+)\s+(TEXT|INTEGER|REAL|DATE)\b", re.IGNORECASE) + + +def schema_columns(schema_sql: str) -> dict[str, list[tuple[str, str]]]: + """Each table's columns and declared types, in order, from schema.sql. + + Used to write Parquet whose types match the tables exactly, which Redshift's + COPY requires: it will not narrow a 64-bit integer into an INTEGER column. + """ + tables = {} + for name, body in _CREATE_TABLE.findall(_strip_comments(schema_sql)): + columns = [] + for line in body.splitlines(): + match = _COLUMN.match(line) + if match: + columns.append((match.group(1), match.group(2).upper())) + tables[name] = columns + return tables + + +def _strip_checks(sql: str) -> str: + """Remove CHECK (...) clauses, following parentheses to their close.""" + out, index = [], 0 + pattern = re.compile(r"\s+CHECK\s*\(", re.IGNORECASE) + while True: + match = pattern.search(sql, index) + if not match: + out.append(sql[index:]) + return "".join(out) + out.append(sql[index:match.start()]) + depth, position = 1, match.end() + while depth: + depth += {"(": 1, ")": -1}.get(sql[position], 0) + position += 1 + index = position + + +def to_redshift(sql: str) -> str: + """Rewrite SQL from this repository for Amazon Redshift. + + The PostgreSQL translation applies, plus what Redshift does not support: + CHECK constraints and indexes are removed (validation enforces the checks + before anything is loaded, and Redshift sorts and distributes instead of + indexing), and TEXT becomes VARCHAR(1024), because Redshift's TEXT is + VARCHAR(256) and would truncate a long product description. + """ + out = to_postgres(sql) + out = _strip_checks(out) + out = re.sub(r"^\s*CREATE INDEX\b[^;]*;\s*$", "", out, flags=re.IGNORECASE | re.MULTILINE) + out = re.sub(r"\bTEXT\b", "VARCHAR(1024)", out) + return out + + +def to_postgres(sql: str) -> str: + """Rewrite SQLite SQL from this repository into PostgreSQL. + + Raises ValueError if a SQLite-only construct survives, so a new query using + one fails at translation rather than as a syntax error from the server. + """ + out = _PRAGMA.sub("", sql) + out = _WEEK_START.sub(r"CAST(DATE_TRUNC('week', \1) AS DATE)", out) + out = _YEAR_MONTH.sub(r"TO_CHAR(\1, 'YYYY-MM')", out) + out = _REAL.sub("DOUBLE PRECISION", out) + + code = _strip_comments(out) + for pattern in _UNTRANSLATED: + match = pattern.search(code) + if match: + raise ValueError( + f"No PostgreSQL translation for {match.group(0)!r}. " + "Add a rule to rootsignal.sql.dialect or write it portably." + ) + return out diff --git a/src/rootsignal/sql/warehouse.py b/src/rootsignal/sql/warehouse.py new file mode 100644 index 0000000..d506c52 --- /dev/null +++ b/src/rootsignal/sql/warehouse.py @@ -0,0 +1,195 @@ +"""Load the cleaned business model into PostgreSQL and query it. + +The SQLite database in database.py is what the test suite builds in memory. +This is the same schema, the same views and the same queries on a PostgreSQL +server: the warehouse the pipeline loads into. Nothing is written twice. The +files in sql/ are translated by dialect.to_postgres, and tests assert both +engines return the same rows for every query. + +Each load replaces one schema wholesale, so a run is repeatable and a failed +load never leaves a half-updated warehouse behind: the new tables are built in +a transaction and the old schema is only dropped inside the same one. +""" + +from __future__ import annotations + +import datetime as dt +import io +import os +from collections.abc import Sequence +from decimal import Decimal +from pathlib import Path + +import pandas as pd + +from .database import ( + ANALYTICS_DIR, + LOAD_ORDER, + MARTS_DIR, + SCHEMA_PATH, + STAGING_DIR, + _sql_files, + to_sqlite_types, +) +from .dialect import ROUND_FUNCTION, to_postgres + +URL_VARIABLE = "ROOTSIGNAL_WAREHOUSE_URL" +DEFAULT_SCHEMA = "rootsignal" + + +def warehouse_url(url: str | None = None) -> str: + """The connection string, from the argument or ROOTSIGNAL_WAREHOUSE_URL.""" + resolved = url or os.environ.get(URL_VARIABLE, "") + if not resolved: + raise ValueError( + f"No warehouse configured. Pass a URL or set {URL_VARIABLE}, for example " + "postgresql://rootsignal:rootsignal@localhost:5433/rootsignal " + "(docker compose up -d warehouse)." + ) + return resolved + + +def connect(url: str | None = None, schema: str = DEFAULT_SCHEMA): + """Open a connection whose unqualified names resolve inside `schema`.""" + import psycopg + + conn = psycopg.connect(warehouse_url(url), autocommit=False) + conn.execute(f'SET search_path TO "{schema}"') + conn.commit() + return conn + + +def _copy_frame(cursor, table: str, frame: pd.DataFrame) -> None: + """Bulk-load a frame with COPY, which is what makes a million rows cheap. + + CSV is used as the wire format. An empty unquoted field is NULL in COPY's + CSV mode, which is how pandas writes a missing value. + """ + buffer = io.StringIO() + to_sqlite_types(frame).to_csv(buffer, index=False, header=False, na_rep="") + columns = ", ".join(f'"{column}"' for column in frame.columns) + with cursor.copy(f'COPY "{table}" ({columns}) FROM STDIN WITH (FORMAT csv)') as copy: + copy.write(buffer.getvalue()) + + +def build_warehouse( + tables: dict[str, pd.DataFrame], + url: str | None = None, + schema: str = DEFAULT_SCHEMA, + with_views: bool = True, +): + """Replace `schema` with the cleaned tables, and register the views. + + Foreign keys are enforced by the server as rows arrive, so a load that + violates the declared model fails here, and the transaction leaves the + previous warehouse as it was. + """ + missing = [name for name in LOAD_ORDER if name not in tables] + if missing: + raise ValueError(f"Cannot build the warehouse; missing tables: {missing}") + + import psycopg + + conn = psycopg.connect(warehouse_url(url), autocommit=False) + try: + with conn.cursor() as cursor: + cursor.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE') + cursor.execute(f'CREATE SCHEMA "{schema}"') + cursor.execute(f'SET LOCAL search_path TO "{schema}"') + cursor.execute(ROUND_FUNCTION) + cursor.execute(to_postgres(SCHEMA_PATH.read_text(encoding="utf-8"))) + for name in LOAD_ORDER: + _copy_frame(cursor, name, tables[name]) + if with_views: + for directory in (STAGING_DIR, MARTS_DIR): + for path in _sql_files(directory): + cursor.execute(to_postgres(path.read_text(encoding="utf-8"))) + cursor.execute("ANALYZE") + conn.commit() + except Exception: + conn.rollback() + conn.close() + raise + conn.close() + return connect(url, schema) + + +def _column_type(series: pd.Series) -> str: + if pd.api.types.is_bool_dtype(series): + return "BOOLEAN" + if pd.api.types.is_integer_dtype(series): + return "BIGINT" + if pd.api.types.is_numeric_dtype(series): + return "DOUBLE PRECISION" + if series.map(lambda value: isinstance(value, dt.date)).any(): + return "DATE" + return "TEXT" + + +def write_report_table(conn, name: str, frame: pd.DataFrame) -> int: + """Replace one report table with a frame, typing columns from its dtypes. + + Report tables are outputs of the analysis layer rather than part of the + declared model, so they carry no keys and are rebuilt on every run. + """ + columns = ", ".join(f'"{column}" {_column_type(frame[column])}' for column in frame.columns) + with conn.cursor() as cursor: + cursor.execute(f'DROP TABLE IF EXISTS "{name}"') + cursor.execute(f'CREATE TABLE "{name}" ({columns})') + buffer = io.StringIO() + frame.to_csv(buffer, index=False, header=False, na_rep="") + names = ", ".join(f'"{column}"' for column in frame.columns) + with cursor.copy(f'COPY "{name}" ({names}) FROM STDIN WITH (FORMAT csv)') as copy: + copy.write(buffer.getvalue()) + conn.commit() + return len(frame) + + +def _plain(value): + """Server types as the plain Python values the SQLite path returns.""" + if isinstance(value, Decimal): + return float(value) + return value + + +def query_warehouse(conn, sql: str, params: Sequence | None = None) -> pd.DataFrame: + """Run a SELECT against the warehouse and return the result as a frame. + + NUMERIC results (AVG of an integer column, for instance) arrive as Decimal + and are returned as float, and dates as ISO strings, so a frame from either + engine can be compared with the other directly. + """ + statement = to_postgres(sql) + if params: + statement = statement.replace("?", "%s") + with conn.cursor() as cursor: + cursor.execute(statement, list(params) if params else None) + columns = [column.name for column in cursor.description] + rows = [[_plain(value) for value in row] for row in cursor.fetchall()] + conn.rollback() + frame = pd.DataFrame(rows, columns=columns) + for column in frame.columns: + if frame[column].map(lambda value: isinstance(value, dt.date)).any(): + frame[column] = frame[column].map( + lambda value: value.isoformat() if isinstance(value, dt.date) else value + ) + elif frame[column].dtype == object and frame[column].map( + lambda value: value is None or isinstance(value, float) + ).all(): + frame[column] = frame[column].astype(float) + return frame + + +def run_warehouse_query_file( + conn, + name: str, + directory: Path | None = None, + params: Sequence | None = None, +) -> pd.DataFrame: + """Run a named analytical query from sql/analytics on the warehouse.""" + folder = directory or ANALYTICS_DIR + path = folder / f"{name}.sql" + if not path.exists(): + available = [candidate.stem for candidate in _sql_files(folder)] + raise ValueError(f"No query named '{name}' in {folder}; available: {available}") + return query_warehouse(conn, path.read_text(encoding="utf-8"), params) diff --git a/tests/test_adapters.py b/tests/test_adapters.py index cd7c1d0..7bfb584 100644 --- a/tests/test_adapters.py +++ b/tests/test_adapters.py @@ -134,6 +134,38 @@ def test_a_product_repeated_on_one_invoice_is_consolidated() -> None: assert not sales.duplicated(subset=["order_id", "sku_id", "sales_type"]).any() +def test_a_stock_code_is_one_product_whatever_its_case_or_padding() -> None: + """The published file writes some codes in both cases, and one with a trailing space. + + 172 codes appear twice this way. 163 pairs carry the same description and + the rest the same product worded differently, so the code is the product + and its case is noise. Kept apart, one product's sales split across two + rows, and any tool comparing text without case (Power BI does) sees + duplicate keys. + """ + raw = raw_lines() + extra = pd.DataFrame( + { + "Invoice": ["536370", "536370", "536371"], + "StockCode": ["15056bl", "15056BL", "47503J "], + "Description": ["EDWARDIAN PARASOL BLACK", "EDWARDIAN PARASOL BLACK", "SET/3 FLORAL GARDEN TOOLS IN BAG"], + "Quantity": [2, 3, 1], + "InvoiceDate": pd.to_datetime(["2009-12-04 11:00", "2009-12-04 11:00", "2009-12-04 12:00"]), + "Price": [5.95, 5.95, 3.75], + "Customer ID": [17850.0, 17850.0, 17850.0], + "Country": ["United Kingdom"] * 3, + } + ) + dataset = adapt(raw=pd.concat([raw, extra], ignore_index=True)) + sales, skus = dataset.tables["fact_sales"], dataset.tables["dim_sku"] + + parasol = sales[sales["sku_id"] == "15056BL"] + assert len(parasol) == 1 and parasol["units"].iloc[0] == 5 + assert "15056bl" not in set(sales["sku_id"]) | set(skus["sku_id"]) + assert "47503J" in set(skus["sku_id"]) and "47503J " not in set(skus["sku_id"]) + assert skus["sku_id"].str.upper().is_unique + + def test_missing_customers_are_named_rather_than_dropped() -> None: """A fifth of real lines have no customer. Dropping them would lose a fifth of revenue.""" dataset = adapt(raw=raw_lines()) diff --git a/tests/test_aws.py b/tests/test_aws.py new file mode 100644 index 0000000..efe5f1c --- /dev/null +++ b/tests/test_aws.py @@ -0,0 +1,505 @@ +"""The AWS side of the pipeline, without an AWS account. + +S3 runs on moto. Glue, EMR Serverless, IAM, EC2 and the two Redshift APIs are +small recording fakes: each test states which calls must happen, which must +not, and what the module does when AWS reports a failure. The fakes answer the +way the services were observed to answer on a real run, including the one +that surprised: Glue's DeleteJob succeeds for a job that does not exist. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from rootsignal.pipeline import CLEAN, Lake, aws +from rootsignal.sql.database import LOAD_ORDER + + +class NotFound(Exception): + pass + + +class FakeClient: + """Records every call; answers from per-method handlers.""" + + def __init__(self, **handlers): + self.calls: list[tuple[str, dict]] = [] + self.handlers = handlers + self.exceptions = SimpleNamespace( + EntityNotFoundException=NotFound, + ResourceNotFoundException=NotFound, + NoSuchEntityException=NotFound, + ClientError=NotFound, + ) + + def __getattr__(self, name): + if name.startswith("__"): + raise AttributeError(name) + + def call(**kwargs): + self.calls.append((name, kwargs)) + handler = self.handlers.get(name) + if handler is None: + return {} + return handler(**kwargs) if callable(handler) else handler + + return call + + def called(self, name: str) -> list[dict]: + return [kwargs for method, kwargs in self.calls if method == name] + + +def missing(**_): + raise NotFound() + + +def sequence(*answers): + """Successive answers to repeated calls, the last one repeating.""" + remaining = list(answers) + + def answer(**_): + value = remaining.pop(0) if len(remaining) > 1 else remaining[0] + if isinstance(value, Exception): + raise value + return value + + return answer + + +@pytest.fixture +def clients(monkeypatch): + """Route aws._client to fakes a test registers, and make waiting instant.""" + registry: dict[str, FakeClient] = {} + monkeypatch.setattr(aws, "_client", lambda service: registry[service]) + monkeypatch.setattr(aws.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(aws, "publish_job", lambda lake: lake.uri("jobs", "conform_online_retail.py")) + return registry + + +LAKE = Lake("s3://rootsignal-lake-test") + + +# -------------------------------------------------------------------------- +# Roles and the bucket +# -------------------------------------------------------------------------- + + +def test_a_role_can_only_reach_its_own_bucket(clients) -> None: + iam = clients["iam"] = FakeClient( + get_role=missing, create_role={"Role": {"Arn": "arn:aws:iam::1:role/r"}} + ) + + arn = aws.ensure_role("r", ["glue.amazonaws.com"], "lake-bucket", write=False) + + assert arn == "arn:aws:iam::1:role/r" + trust = json.loads(iam.called("create_role")[0]["AssumeRolePolicyDocument"]) + assert trust["Statement"][0]["Principal"]["Service"] == ["glue.amazonaws.com"] + policy = json.loads(iam.called("put_role_policy")[0]["PolicyDocument"])["Statement"][0] + assert policy["Resource"] == ["arn:aws:s3:::lake-bucket", "arn:aws:s3:::lake-bucket/*"] + assert "s3:PutObject" not in policy["Action"] and "s3:DeleteObject" not in policy["Action"] + + +def test_a_writing_role_may_put_and_delete_and_managed_policies_attach(clients) -> None: + iam = clients["iam"] = FakeClient(get_role={"Role": {"Arn": "arn:existing"}}) + + aws.ensure_role("r", ["glue.amazonaws.com"], "b", write=True, managed=("arn:managed",)) + + assert not iam.called("create_role"), "an existing role is reused, not recreated" + actions = json.loads(iam.called("put_role_policy")[0]["PolicyDocument"])["Statement"][0]["Action"] + assert {"s3:PutObject", "s3:DeleteObject"} <= set(actions) + assert iam.called("attach_role_policy") == [{"RoleName": "r", "PolicyArn": "arn:managed"}] + + +def test_the_bucket_is_created_private(monkeypatch) -> None: + moto = pytest.importorskip("moto") + import boto3 + + for name, value in { + "AWS_ACCESS_KEY_ID": "testing", + "AWS_SECRET_ACCESS_KEY": "testing", + "AWS_DEFAULT_REGION": "ap-south-1", + }.items(): + monkeypatch.setenv(name, value) + with moto.mock_aws(): + aws.ensure_bucket("rootsignal-private") + block = boto3.client("s3").get_public_access_block(Bucket="rootsignal-private") + assert all(block["PublicAccessBlockConfiguration"].values()) + aws.ensure_bucket("rootsignal-private") # a second run changes nothing + + +# -------------------------------------------------------------------------- +# Glue and EMR Serverless +# -------------------------------------------------------------------------- + + +def test_glue_creates_the_job_with_the_lake_paths_and_waits_for_it(clients) -> None: + glue = clients["glue"] = FakeClient( + get_job=missing, + start_job_run={"JobRunId": "jr_1"}, + get_job_run=sequence( + {"JobRun": {"JobRunState": "RUNNING"}}, + {"JobRun": {"JobRunState": "SUCCEEDED", "ExecutionTime": 150}}, + ), + ) + + run = aws.run_glue(LAKE, "arn:role", workers=2) + + job = glue.called("create_job")[0] + assert job["Name"] == aws.GLUE_JOB and job["Role"] == "arn:role" + assert job["DefaultArguments"] == { + "--input": "s3://rootsignal-lake-test/raw/online_retail/invoices.csv", + "--output": "s3://rootsignal-lake-test/curated/online_retail", + } + assert job["NumberOfWorkers"] == 2 and job["Timeout"] == 20 and job["MaxRetries"] == 0 + assert len(glue.called("get_job_run")) == 2 + assert run.state == "SUCCEEDED" + assert run.usage == f"{150 / 3600 * 2:.3f} DPU-hours" + + +def test_glue_updates_an_existing_job_and_bills_at_least_a_minute(clients) -> None: + glue = clients["glue"] = FakeClient( + get_job={"Job": {}}, + start_job_run={"JobRunId": "jr_2"}, + get_job_run={"JobRun": {"JobRunState": "SUCCEEDED", "ExecutionTime": 20}}, + ) + + run = aws.run_glue(LAKE, "arn:role", workers=2) + + assert not glue.called("create_job") and len(glue.called("update_job")) == 1 + assert run.usage == f"{60 / 3600 * 2:.3f} DPU-hours" + + +def test_a_failed_glue_run_is_reported_with_its_error(clients) -> None: + clients["glue"] = FakeClient( + get_job={"Job": {}}, + start_job_run={"JobRunId": "jr_3"}, + get_job_run={"JobRun": {"JobRunState": "FAILED", "ExecutionTime": 70, "ErrorMessage": "boom"}}, + ) + + run = aws.run_glue(LAKE, "arn:role") + + assert (run.state, run.error) == ("FAILED", "boom") + + +def test_emr_reuses_a_live_application_and_ignores_a_terminated_one(clients) -> None: + emr = clients["emr-serverless"] = FakeClient( + list_applications={ + "applications": [ + {"name": aws.EMR_APP, "state": "TERMINATED", "id": "old"}, + {"name": aws.EMR_APP, "state": "STOPPED", "id": "live"}, + ] + } + ) + assert aws.ensure_emr_application() == "live" + assert not emr.called("create_application") + + +def test_a_new_emr_application_is_capped_and_stops_itself(clients) -> None: + emr = clients["emr-serverless"] = FakeClient( + list_applications={"applications": []}, create_application={"applicationId": "new"} + ) + assert aws.ensure_emr_application() == "new" + created = emr.called("create_application")[0] + assert created["maximumCapacity"]["cpu"] == "8 vCPU" + assert created["autoStopConfiguration"] == {"enabled": True, "idleTimeoutMinutes": 2} + + +def test_emr_runs_the_same_script_with_the_same_arguments(clients) -> None: + emr = clients["emr-serverless"] = FakeClient( + list_applications={"applications": [{"name": aws.EMR_APP, "state": "STARTED", "id": "app"}]}, + start_job_run={"jobRunId": "run"}, + get_job_run=sequence( + {"jobRun": {"state": "RUNNING"}}, + {"jobRun": {"state": "SUCCESS", "totalExecutionDurationSeconds": 95, + "totalResourceUtilization": {"vCPUHour": 0.1, "memoryGBHour": 0.4}}}, + ), + ) + + run = aws.run_emr(LAKE, "arn:emr", "s3://rootsignal-lake-test/curated_emr/online_retail") + + driver = emr.called("start_job_run")[0]["jobDriver"]["sparkSubmit"] + assert driver["entryPoint"] == "s3://rootsignal-lake-test/jobs/conform_online_retail.py" + assert driver["entryPointArguments"] == [ + "--input", "s3://rootsignal-lake-test/raw/online_retail/invoices.csv", + "--output", "s3://rootsignal-lake-test/curated_emr/online_retail", + ] + assert emr.called("start_job_run")[0]["executionTimeoutMinutes"] == 20 + assert (run.state, run.seconds) == ("SUCCESS", 95.0) + assert run.usage == "0.100 vCPU-hours, 0.400 GB-hours" + + +# -------------------------------------------------------------------------- +# Redshift Serverless +# -------------------------------------------------------------------------- + + +def network_clients(registry) -> None: + registry["ec2"] = FakeClient( + describe_vpcs={"Vpcs": [{"VpcId": "vpc-1"}]}, + describe_subnets={"Subnets": [{"SubnetId": "a"}, {"SubnetId": "b"}, {"SubnetId": "c"}]}, + describe_security_groups={"SecurityGroups": [{"GroupId": "sg-1"}]}, + ) + + +def test_redshift_is_created_at_one_size_with_a_spending_cap(clients) -> None: + network_clients(clients) + rs = clients["redshift-serverless"] = FakeClient( + get_namespace=missing, + get_workgroup=sequence( + NotFound(), + {"workgroup": {"status": "AVAILABLE", "workgroupArn": "arn:wg"}}, + ), + create_workgroup={"workgroup": {"status": "CREATING", "workgroupArn": "arn:wg"}}, + list_usage_limits={"usageLimits": []}, + ) + + assert aws.ensure_redshift("arn:rs", base_capacity=8, daily_rpu_hours=4) == "arn:wg" + + namespace = rs.called("create_namespace")[0] + assert namespace["defaultIamRoleArn"] == "arn:rs" + workgroup = rs.called("create_workgroup")[0] + assert workgroup["baseCapacity"] == workgroup["maxCapacity"] == 8 + assert workgroup["publiclyAccessible"] is False + assert workgroup["subnetIds"] == ["a", "b", "c"] and workgroup["securityGroupIds"] == ["sg-1"] + assert rs.called("create_usage_limit") == [ + {"resourceArn": "arn:wg", "usageType": "serverless-compute", "amount": 4, + "period": "daily", "breachAction": "deactivate"} + ] + + +def test_an_existing_spending_cap_is_not_added_twice(clients) -> None: + rs = clients["redshift-serverless"] = FakeClient( + get_namespace={"namespace": {}}, + get_workgroup={"workgroup": {"status": "AVAILABLE", "workgroupArn": "arn:wg"}}, + list_usage_limits={"usageLimits": [{"usageLimitId": "u"}]}, + ) + aws.ensure_redshift("arn:rs") + assert not rs.called("create_namespace") and not rs.called("create_workgroup") + assert not rs.called("create_usage_limit") + + +def test_one_statement_is_executed_alone_and_several_as_one_batch(clients) -> None: + data = clients["redshift-data"] = FakeClient( + execute_statement={"Id": "single"}, + batch_execute_statement={"Id": "batch"}, + describe_statement=lambda Id: { + "Status": "FINISHED", "SubStatements": [{"Id": f"{Id}:1"}, {"Id": f"{Id}:2"}] + }, + ) + assert aws.execute(["SELECT 1", " "]) == ["single"] + assert aws.execute(["SET x", "SELECT 1"]) == ["batch:1", "batch:2"] + assert data.called("batch_execute_statement")[0]["Sqls"] == ["SET x", "SELECT 1"] + + +def test_a_failed_statement_raises_with_redshifts_own_message(clients) -> None: + clients["redshift-data"] = FakeClient( + execute_statement={"Id": "s"}, + describe_statement={"Status": "FAILED", "Error": "column \"x\" does not exist"}, + ) + with pytest.raises(RuntimeError, match='column "x" does not exist'): + aws.execute(["SELECT x"]) + + +def test_a_statement_that_never_finishes_is_cancelled(clients, monkeypatch) -> None: + data = clients["redshift-data"] = FakeClient( + execute_statement={"Id": "slow"}, describe_statement={"Status": "STARTED"} + ) + clock = iter([0.0, 5.0, 11.0]) + monkeypatch.setattr(aws.time, "time", lambda: next(clock)) + with pytest.raises(TimeoutError): + aws.execute(["SELECT pg_sleep(100)"], timeout=10) + assert data.called("cancel_statement") == [{"Id": "slow"}] + + +def test_fetch_sets_the_schema_in_the_same_batch_and_reads_every_page(clients) -> None: + pages = { + None: { + "ColumnMetadata": [{"name": "region_code"}, {"name": "net_sales"}], + "Records": [[{"stringValue": "France"}, {"doubleValue": 10.5}]], + "NextToken": "page2", + }, + "page2": { + "ColumnMetadata": [{"name": "region_code"}, {"name": "net_sales"}], + "Records": [[{"stringValue": "Spain"}, {"isNull": True}]], + }, + } + data = clients["redshift-data"] = FakeClient( + batch_execute_statement={"Id": "b"}, + describe_statement={"Status": "FINISHED", "SubStatements": [{"Id": "b:1"}, {"Id": "b:2"}]}, + get_statement_result=lambda Id, NextToken=None: pages[NextToken], + ) + + frame = aws.fetch("SELECT * FROM mart", schema="online_retail") + + assert data.called("batch_execute_statement")[0]["Sqls"] == [ + "SET search_path TO online_retail", "SELECT * FROM mart" + ] + assert {call["Id"] for call in data.called("get_statement_result")} == {"b:2"} + assert frame["region_code"].tolist() == ["France", "Spain"] + assert frame["net_sales"].tolist()[0] == 10.5 and frame["net_sales"].isna().tolist()[1] + + +def test_the_redshift_load_copies_every_table_from_the_clean_zone(clients) -> None: + batches = [] + + def batch(Sqls, **_): + batches.append(Sqls) + return {"Id": f"b{len(batches)}"} + + counts_page = { + "ColumnMetadata": [{"name": "name"}, {"name": "n"}], + "Records": [[{"stringValue": name}, {"longValue": 7}] for name in LOAD_ORDER], + } + clients["redshift-data"] = FakeClient( + batch_execute_statement=batch, + describe_statement=lambda Id: {"Status": "FINISHED", "SubStatements": [{"Id": f"{Id}:1"}, {"Id": f"{Id}:2"}]}, + get_statement_result=counts_page, + ) + + counts = aws.load_redshift(LAKE, "online_retail") + + ddl = batches[0] + assert ddl[:3] == ["DROP SCHEMA IF EXISTS online_retail CASCADE", "CREATE SCHEMA online_retail", + "SET search_path TO online_retail"] + assert sum(statement.startswith("CREATE TABLE") for statement in ddl) == len(LOAD_ORDER) + assert not any("CHECK" in statement or "CREATE INDEX" in statement for statement in ddl) + copies = [statement for statement in ddl if statement.startswith("COPY")] + assert copies == [ + f"COPY {name} FROM '{LAKE.uri(CLEAN, name)}/' IAM_ROLE default FORMAT AS PARQUET" for name in LOAD_ORDER + ] + views = batches[1] + assert views[0] == "SET search_path TO online_retail" + assert any(statement.startswith("CREATE VIEW mart_commercial_daily") for statement in views) + assert counts == {name: 7 for name in LOAD_ORDER} + + +# -------------------------------------------------------------------------- +# Teardown +# -------------------------------------------------------------------------- + + +def test_teardown_reports_only_what_it_actually_removed(clients) -> None: + """DeleteJob succeeds for a job that never existed; that is not a removal.""" + glue = clients["glue"] = FakeClient(get_job=missing, delete_job={"JobName": aws.GLUE_JOB}) + clients["emr-serverless"] = FakeClient(list_applications={"applications": []}) + clients["redshift-serverless"] = FakeClient(get_workgroup=missing, delete_namespace=missing) + + assert aws.teardown() == [] + assert not glue.called("delete_job") + + +def test_teardown_stops_a_running_application_before_deleting_it(clients) -> None: + clients["glue"] = FakeClient(get_job={"Job": {}}) + emr = clients["emr-serverless"] = FakeClient( + list_applications={"applications": [{"name": aws.EMR_APP, "state": "STARTED", "id": "app"}]}, + get_application=sequence({"application": {"state": "STOPPING"}}, {"application": {"state": "STOPPED"}}), + ) + rs = clients["redshift-serverless"] = FakeClient( + get_workgroup=sequence({"workgroup": {}}, {"workgroup": {}}, NotFound()), + get_namespace=sequence({"namespace": {}}, NotFound()), + ) + + removed = aws.teardown() + + order = [method for method, _ in emr.calls] + assert order.index("stop_application") < order.index("delete_application") + assert rs.called("delete_workgroup") and rs.called("delete_namespace") + assert removed == [ + f"Glue job {aws.GLUE_JOB}", + "EMR Serverless application app", + f"Redshift Serverless workgroup {aws.WORKGROUP}", + f"Redshift Serverless namespace {aws.NAMESPACE}", + ] + + +def test_billable_resources_lists_what_exists_and_nothing_else(clients) -> None: + clients["glue"] = FakeClient(get_job=missing) + clients["emr-serverless"] = FakeClient( + list_applications={"applications": [ + {"name": aws.EMR_APP, "state": "TERMINATED", "id": "gone"}, + {"name": "someone-elses-app", "state": "STARTED", "id": "other"}, + ]} + ) + clients["redshift-serverless"] = FakeClient( + list_workgroups={"workgroups": [{"workgroupName": aws.WORKGROUP, "status": "AVAILABLE"}]}, + list_namespaces={"namespaces": [{"namespaceName": "other-ns"}]}, + ) + + assert aws.billable_resources() == [f"Redshift Serverless workgroup {aws.WORKGROUP} (AVAILABLE)"] + + +# -------------------------------------------------------------------------- +# The ceilings, and the edges of the arithmetic +# -------------------------------------------------------------------------- + + +def test_the_spending_ceilings_are_the_ones_the_module_states() -> None: + """Glue at two workers, Redshift at 8 RPU capped at 4 RPU-hours a day.""" + import inspect + + assert inspect.signature(aws.run_glue).parameters["workers"].default == 2 + redshift = inspect.signature(aws.ensure_redshift).parameters + assert redshift["base_capacity"].default == 8 + assert redshift["daily_rpu_hours"].default == 4 + assert inspect.signature(aws.execute).parameters["timeout"].default == 900 + + +def test_an_hour_on_two_glue_workers_is_two_dpu_hours(clients) -> None: + clients["glue"] = FakeClient( + get_job={"Job": {}}, + start_job_run={"JobRunId": "jr"}, + get_job_run={"JobRun": {"JobRunState": "SUCCEEDED", "ExecutionTime": 3600}}, + ) + assert aws.run_glue(LAKE, "arn:role", workers=2).usage == "2.000 DPU-hours" + + +def test_a_run_that_reports_no_time_is_billed_the_minimum(clients) -> None: + clients["glue"] = FakeClient( + get_job={"Job": {}}, start_job_run={"JobRunId": "jr"}, get_job_run={"JobRun": {"JobRunState": "FAILED"}} + ) + run = aws.run_glue(LAKE, "arn:role", workers=2) + assert run.seconds == 0.0 + assert run.usage == f"{60 / 3600 * 2:.3f} DPU-hours" + + +def test_emr_usage_missing_from_the_response_is_reported_as_zero(clients) -> None: + clients["emr-serverless"] = FakeClient( + list_applications={"applications": [{"name": aws.EMR_APP, "state": "STARTED", "id": "app"}]}, + start_job_run={"jobRunId": "run"}, + get_job_run={"jobRun": {"state": "FAILED", "stateDetails": "quota"}}, + ) + run = aws.run_emr(LAKE, "arn:emr", "s3://x/out") + assert (run.seconds, run.usage, run.error) == (0.0, "0.000 vCPU-hours, 0.000 GB-hours", "quota") + + +def test_a_statement_finishing_exactly_at_the_timeout_is_not_cancelled(clients, monkeypatch) -> None: + """The limit is exceeded only after the timeout, not on it. + + The clock starts at 100 rather than 0, so the elapsed time is measured, not + the clock's reading. + """ + data = clients["redshift-data"] = FakeClient( + execute_statement={"Id": "s"}, + describe_statement=sequence({"Status": "STARTED"}, {"Status": "FINISHED"}), + ) + clock = iter([100.0, 110.0]) + monkeypatch.setattr(aws.time, "time", lambda: next(clock)) + assert aws.execute(["SELECT 1"], timeout=10) == ["s"] + assert not data.called("cancel_statement") + + +def test_an_analytical_query_runs_translated_inside_its_schema(clients) -> None: + data = clients["redshift-data"] = FakeClient( + batch_execute_statement={"Id": "b"}, + describe_statement={"Status": "FINISHED", "SubStatements": [{"Id": "b:1"}, {"Id": "b:2"}]}, + get_statement_result={"ColumnMetadata": [{"name": "week_start"}], "Records": [[{"stringValue": "2011-11-28"}]]}, + ) + frame = aws.query_redshift("sample", "weekly_movers") + + set_path, sql = data.called("batch_execute_statement")[0]["Sqls"] + assert set_path == "SET search_path TO sample" + assert "DATE_TRUNC('week'" in sql and "weekday 0" not in sql + assert frame["week_start"].tolist() == ["2011-11-28"] diff --git a/tests/test_lake.py b/tests/test_lake.py new file mode 100644 index 0000000..125c36c --- /dev/null +++ b/tests/test_lake.py @@ -0,0 +1,88 @@ +"""The lake behaves the same on a local folder and on S3. + +S3 is simulated with moto, so these run without an AWS account. +""" + +from __future__ import annotations + +import pytest + +from rootsignal.pipeline import CURATED, RAW, Lake + + +def test_a_local_lake_is_a_folder(tmp_path) -> None: + lake = Lake(str(tmp_path)) + source = tmp_path / "source.csv" + source.write_text("a,b\n1,2\n") + + uri = lake.put(source, RAW, "invoices.csv") + + assert uri == str(tmp_path / RAW / "invoices.csv") + assert lake.exists(RAW, "invoices.csv") + assert lake.fetch(RAW, "invoices.csv", into=tmp_path / "unused").read_text() == "a,b\n1,2\n" + + +def test_an_s3_lake_splits_bucket_and_prefix() -> None: + lake = Lake("s3://rootsignal-lake/project/") + assert lake.bucket == "rootsignal-lake" + assert lake.uri(CURATED, "fact_sales") == "s3://rootsignal-lake/project/curated/online_retail/fact_sales" + assert Lake("s3://rootsignal-lake").uri(RAW) == "s3://rootsignal-lake/raw/online_retail" + + +def test_an_s3_prefix_can_be_several_folders_deep() -> None: + lake = Lake("s3://rootsignal-lake/teams/analytics/rootsignal") + assert lake.bucket == "rootsignal-lake" + assert lake.prefix == "teams/analytics/rootsignal" + assert lake.key(RAW, "invoices.csv") == "teams/analytics/rootsignal/raw/online_retail/invoices.csv" + + +@pytest.fixture +def s3(monkeypatch): + moto = pytest.importorskip("moto") + import boto3 + + for name, value in { + "AWS_ACCESS_KEY_ID": "testing", + "AWS_SECRET_ACCESS_KEY": "testing", + "AWS_DEFAULT_REGION": "ap-south-1", + }.items(): + monkeypatch.setenv(name, value) + with moto.mock_aws(): + client = boto3.client("s3") + client.create_bucket( + Bucket="rootsignal-lake", CreateBucketConfiguration={"LocationConstraint": "ap-south-1"} + ) + yield client + + +def test_an_s3_lake_round_trips_a_prefix(s3, tmp_path) -> None: + lake = Lake("s3://rootsignal-lake/project") + for part in ("part-0000.parquet", "part-0001.parquet"): + source = tmp_path / part + source.write_bytes(part.encode()) + lake.put(source, CURATED, "fact_sales", part) + + assert lake.exists(CURATED, "fact_sales") + assert not lake.exists(CURATED, "dim_sku") + + fetched = lake.fetch(CURATED, "fact_sales", into=tmp_path / "copy") + assert sorted(path.name for path in fetched.iterdir()) == ["part-0000.parquet", "part-0001.parquet"] + assert (fetched / "part-0001.parquet").read_bytes() == b"part-0001.parquet" + + +def test_a_folder_does_not_pick_up_a_sibling_sharing_its_name(s3, tmp_path) -> None: + lake = Lake("s3://rootsignal-lake") + source = tmp_path / "part-0000.parquet" + source.write_bytes(b"x") + lake.put(source, CURATED, "fact_sales", "part-0000.parquet") + lake.put(source, CURATED, "fact_sales_backup", "part-0000.parquet") + + fetched = lake.fetch(CURATED, "fact_sales", into=tmp_path / "copy") + + assert [path.name for path in fetched.iterdir()] == ["part-0000.parquet"] + assert not (tmp_path / "copy" / CURATED / "fact_sales_backup").exists() + + +def test_fetching_nothing_from_s3_is_an_error_not_an_empty_folder(s3, tmp_path) -> None: + with pytest.raises(FileNotFoundError, match="curated/online_retail/returns"): + Lake("s3://rootsignal-lake").fetch(CURATED, "returns", into=tmp_path) diff --git a/tests/test_pipeline_stages.py b/tests/test_pipeline_stages.py new file mode 100644 index 0000000..61c6350 --- /dev/null +++ b/tests/test_pipeline_stages.py @@ -0,0 +1,241 @@ +"""The load and publish stages: what they write, and when they refuse to. + +The clean zone is what Redshift loads with COPY, which rejects a Parquet type +that does not match the table, so its types are pinned to schema.sql here. +The report tables must be plain enough for a SQL table to hold. +""" + +from __future__ import annotations + +import datetime as dt +import os +import uuid + +import pandas as pd +import pytest + +from rootsignal.adapters.base import AdaptedDataset, DatasetCapabilities +from rootsignal.pipeline import CLEAN, Lake, load, write_clean_zone +from rootsignal.pipeline.publish import flatten +from rootsignal.sql.database import LOAD_ORDER, SCHEMA_PATH +from rootsignal.sql.dialect import schema_columns +from rootsignal.sql.warehouse import URL_VARIABLE, connect, write_report_table +from rootsignal.validation.validator import ValidationIssue, ValidationReport + + +def test_the_clean_zone_is_typed_exactly_as_the_schema_declares(cleaned_dataset, tmp_path) -> None: + import pyarrow as pa + import pyarrow.parquet as pq + + expected_types = {"TEXT": pa.string(), "INTEGER": pa.int32(), "REAL": pa.float64(), "DATE": pa.date32()} + lake = Lake(str(tmp_path)) + write_clean_zone(lake, cleaned_dataset.tables) + + declared = schema_columns(SCHEMA_PATH.read_text(encoding="utf-8")) + for name in LOAD_ORDER: + schema = pq.read_schema(tmp_path / CLEAN / name / "part-00000.parquet") + assert schema.names == [column for column, _ in declared[name]], name + for column, kind in declared[name]: + assert schema.field(column).type == expected_types[kind], f"{name}.{column}" + + +def test_flags_are_written_as_zero_or_one(cleaned_dataset, tmp_path) -> None: + import pyarrow.parquet as pq + + write_clean_zone(Lake(str(tmp_path)), cleaned_dataset.tables) + flags = pq.read_table(tmp_path / CLEAN / "dim_sku" / "part-00000.parquet").column("active_flag") + assert set(flags.to_pylist()) <= {0, 1} + + +def test_row_counts_survive_the_clean_zone(cleaned_dataset, tmp_path) -> None: + import pyarrow.parquet as pq + + write_clean_zone(Lake(str(tmp_path)), cleaned_dataset.tables) + for name in LOAD_ORDER: + rows = pq.read_metadata(tmp_path / CLEAN / name / "part-00000.parquet").num_rows + assert rows == len(cleaned_dataset.tables[name]), name + + +def test_report_tables_are_flattened_to_plain_columns() -> None: + frame = pd.DataFrame( + { + "period_start": pd.PeriodIndex(["2011-10", "2011-11"], freq="M"), + "stamp": pd.to_datetime(["2011-10-03 09:00", "2011-11-04 10:00"]), + "value": [1.5, float("nan")], + "count": [3, 4], + "complete": [True, False], + "segments": [("United Kingdom", "France"), None], + }, + index=pd.Index(["a", "b"], name="kept_out"), + ) + + flat = flatten(frame) + + assert list(flat.columns) == ["period_start", "stamp", "value", "count", "complete", "segments"] + assert flat["period_start"].tolist() == [dt.date(2011, 10, 1), dt.date(2011, 11, 1)] + assert flat["stamp"].tolist() == [dt.date(2011, 10, 3), dt.date(2011, 11, 4)] + assert flat["segments"].tolist() == ["('United Kingdom', 'France')", None] + assert flat["count"].tolist() == [3, 4] and flat["complete"].tolist() == [True, False] + + +def test_the_spark_job_is_where_the_pipeline_looks_for_it() -> None: + from rootsignal.pipeline.stages import JOB_PATH + + assert JOB_PATH.is_file() and JOB_PATH.name == "conform_online_retail.py" + + +def test_a_failed_spark_run_stops_the_pipeline_with_its_error(monkeypatch, tmp_path) -> None: + """The job's own last words are what someone debugging it needs.""" + import subprocess + + from rootsignal.pipeline import conform_locally, stages + + failed = subprocess.CompletedProcess(args=[], returncode=1, stdout="", stderr="x" * 6000 + "Py4JJavaError: disk full") + monkeypatch.setattr(stages.subprocess, "run", lambda *args, **kwargs: failed) + + with pytest.raises(RuntimeError, match="disk full") as raised: + conform_locally(Lake(str(tmp_path))) + assert len(str(raised.value)) < 4100, "the whole log is not pasted into the error" + + +def test_a_run_that_exits_cleanly_is_not_reported_as_failed(monkeypatch, tmp_path) -> None: + import subprocess + + from rootsignal.pipeline import conform_locally, stages + + monkeypatch.setattr( + stages.subprocess, "run", + lambda *args, **kwargs: subprocess.CompletedProcess(args=[], returncode=0, stdout="", stderr="WARN only"), + ) + monkeypatch.setattr(stages, "read_summary", lambda lake: {"source_rows": 3}) + assert conform_locally(Lake(str(tmp_path))) == {"source_rows": 3} + + +def write_curated_zone(root, cleaned_tables) -> None: + """A curated zone laid out as the Spark job writes it, without running Spark.""" + import json + + from rootsignal.pipeline import CURATED + from rootsignal.pipeline.stages import CURATED_TABLES + + for name in CURATED_TABLES: + folder = root / CURATED / name + folder.mkdir(parents=True) + cleaned_tables[name].to_parquet(folder / "part-00000.parquet", index=False) + returns = pd.DataFrame( + {"order_id": ["C1"], "sku_id": ["A"], "region_code": ["France"], "customer_id": ["UNKNOWN"], + "date": [dt.date(2024, 1, 3)], "units": [-2], "returned_units": [2]} + ) + (root / CURATED / "returns").mkdir(parents=True) + returns.to_parquet(root / CURATED / "returns" / "part-00000.parquet", index=False) + (root / CURATED / "_summary").mkdir(parents=True) + summary = {"source_rows": 1234, "returned_lines": 56, "consolidated_lines": 7, "tables": {}} + (root / CURATED / "_summary" / "part-00000.txt").write_text(json.dumps(summary), encoding="utf-8") + + +def test_the_curated_zone_reads_back_as_the_adapter_would_return_it(cleaned_dataset, tmp_path) -> None: + from rootsignal.pipeline import read_curated + from rootsignal.pipeline.stages import EMPTY_FACTS + + write_curated_zone(tmp_path, cleaned_dataset.tables) + dataset, returns = read_curated(Lake(str(tmp_path))) + + assert dataset.source_rows == 1234 + sales = dataset.tables["fact_sales"] + assert len(sales) == len(cleaned_dataset.tables["fact_sales"]) + assert isinstance(sales["date"].iloc[0], dt.date) and not isinstance(sales["date"].iloc[0], dt.datetime) + assert isinstance(returns["date"].iloc[0], dt.date) + for name in EMPTY_FACTS: + assert name in dataset.tables and dataset.tables[name].empty + notes = " ".join(dataset.capabilities.notes) + assert "56 of 1,234 lines are returns" in notes + assert "7 invoice lines repeat a product" in notes + assert "Conformed in Spark" in notes + + +def test_the_report_backtest_is_the_one_behind_the_published_figure() -> None: + """The README's 38.8% over 49 folds comes from these settings. + + scripts/run_external_dataset.py produces that figure and the report table + must match it, so both read one constant rather than two copies. + """ + from pathlib import Path + + from rootsignal.pipeline.publish import BACKTEST + + assert BACKTEST == {"horizon": 7, "initial_train": 56, "step": 14} + script = (Path(__file__).resolve().parents[1] / "scripts" / "run_external_dataset.py").read_text(encoding="utf-8") + assert "rolling_origin_evaluate(series, **BACKTEST)" in script + + +def test_nothing_is_loaded_while_validation_still_fails(monkeypatch, cleaned_dataset) -> None: + import rootsignal.sql.warehouse as warehouse + from rootsignal.validation import DatasetValidator + + failing = ValidationReport( + issues=(ValidationIssue("fact_sales", "keys", "ERROR", 2, "duplicate order lines"),) + ) + monkeypatch.setattr(DatasetValidator, "validate", lambda self, tables: failing) + built = [] + monkeypatch.setattr(warehouse, "build_warehouse", lambda *args, **kwargs: built.append(True)) + dataset = AdaptedDataset( + tables=cleaned_dataset.tables, + capabilities=DatasetCapabilities(name="test", available_facts=("fact_sales",)), + source_rows=0, + ) + + with pytest.raises(ValueError, match="fact_sales.keys: duplicate order lines"): + load(dataset, url="postgresql://unused") + assert not built + + +needs_warehouse = pytest.mark.skipif( + not os.environ.get(URL_VARIABLE), + reason=f"{URL_VARIABLE} is not set; start one with `docker compose up -d warehouse`", +) + + +@needs_warehouse +def test_a_report_table_round_trips_through_the_warehouse() -> None: + schema = f"test_{uuid.uuid4().hex[:10]}" + conn = connect(schema="public") + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute(f'SET search_path TO "{schema}"') + conn.commit() + try: + frame = pd.DataFrame( + { + "segment": ["United Kingdom", None], + "contribution": [386371.2, float("nan")], + "folds": [49, 7], + "is_complete": [True, False], + "period_start": [dt.date(2011, 11, 1), dt.date(2011, 10, 1)], + } + ) + assert write_report_table(conn, "rpt_example", frame) == 2 + types = dict( + conn.execute( + "SELECT column_name, data_type FROM information_schema.columns " + "WHERE table_schema = %s AND table_name = 'rpt_example'", + (schema,), + ).fetchall() + ) + assert types == { + "segment": "text", + "contribution": "double precision", + "folds": "bigint", + "is_complete": "boolean", + "period_start": "date", + } + rows = conn.execute("SELECT * FROM rpt_example ORDER BY folds DESC").fetchall() + assert rows == [ + ("United Kingdom", 386371.2, 49, True, dt.date(2011, 11, 1)), + (None, None, 7, False, dt.date(2011, 10, 1)), + ] + assert write_report_table(conn, "rpt_example", frame.head(1)) == 1, "a rerun replaces the table" + assert conn.execute("SELECT COUNT(*) FROM rpt_example").fetchone()[0] == 1 + finally: + conn.rollback() + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + conn.commit() + conn.close() diff --git a/tests/test_spark_conform.py b/tests/test_spark_conform.py new file mode 100644 index 0000000..d97467f --- /dev/null +++ b/tests/test_spark_conform.py @@ -0,0 +1,227 @@ +"""The Spark conform job produces the tables the pandas adapter produces. + +jobs/conform_online_retail.py is a second implementation of +rootsignal.adapters.online_retail.adapt, so it is held to the first one table by +table. The raw lines below plant every case the adapter handles: returns by +invoice prefix and by negative quantity, a missing customer, a product repeated +on one invoice, a zero price, descriptions that are missing, blank, quoted, +multi-line or padded, and a revenue that lands exactly on a half cent. + +Spark needs a Java runtime, and on Windows a Hadoop helper binary as well, so +these tests run where Spark can: in CI and in the project's container. +""" + +from __future__ import annotations + +import importlib.util +import os +import shutil + +import numpy as np +import pandas as pd +import pytest + +pyspark = pytest.importorskip("pyspark") + +if shutil.which("java") is None and not os.environ.get("JAVA_HOME"): + pytest.skip("Spark needs a Java runtime", allow_module_level=True) +if os.name == "nt" and not os.environ.get("HADOOP_HOME"): + pytest.skip("Spark cannot write files on Windows without HADOOP_HOME", allow_module_level=True) + +from rootsignal.adapters.online_retail import adapt, returns_frame # noqa: E402 +from rootsignal.pipeline import CURATED, RAW, Lake, read_curated, write_raw_csv # noqa: E402 +from rootsignal.pipeline.stages import JOB_PATH, RAW_CSV # noqa: E402 + +KEYS = { + "fact_sales": ["order_id", "sku_id", "sales_type"], + "dim_date": ["date"], + "dim_sku": ["sku_id"], + "dim_customer": ["customer_id"], + "dim_region": ["region_code"], + "dim_kam": ["kam_id"], +} + + +def raw_lines(repeat_a_product: bool = True) -> pd.DataFrame: + """Invoice lines shaped as pandas reads them from the workbook.""" + rows = [ + # Invoice, StockCode, Description, Quantity, InvoiceDate, Price, Customer ID, Country + (489434, 85048, "15CM CHRISTMAS GLASS BALL 20 LIGHTS", 12, "2009-12-01 07:45", 6.95, 13085.0, "United Kingdom"), + (489434, "79323P", "PINK CHERRY LIGHTS", 12, "2009-12-01 07:45", 6.75, 13085.0, "United Kingdom"), + (489434, "79323W", ' white "cherry" lights, boxed', 12, "2009-12-01 07:45", 6.75, 13085.0, "United Kingdom"), + (489435, 22350, "CAT BOWL\nwith lid", 12, "2009-12-01 07:46", 2.55, np.nan, "France"), + (489435, 22349, None, 24, "2009-12-01 07:46", 3.75, np.nan, "France"), + (489435, 22195, " ", 1, "2009-12-02 09:00", 0.0, np.nan, "France"), + (489436, 21755, "LOVE BUILDING BLOCK WORD", 2, "2009-12-02 10:03", 0.5025, 13078.0, "Germany"), + (489436, 85048, "15CM CHRISTMAS GLASS BALL 20 LIGHTS", 6, "2009-12-02 10:03", 7.95, 13078.0, "Germany"), + ("C489449", 22087, "PAPER BUNTING WHITE LACE", -12, "2009-12-03 10:33", 2.95, 16321.0, "Australia"), + (489450, 21755, "LOVE BUILDING BLOCK WORD", -2, "2009-12-03 11:00", 5.95, 13078.0, "Germany"), + (489451, "POST", "POSTAGE", 1, "2009-12-04 12:00", 18.0, 12682.0, "France"), + # A credit note with a positive quantity: a return by its prefix alone. + ("C489460", 22423, "REGENCY CAKESTAND 3 TIER", 1, "2009-12-05 09:00", 12.75, 12682.0, "France"), + # 22195 already sold at a zero price above; its list price must ignore it. + (489461, 22195, "HEART MEASURING SPOONS", 3, "2009-12-05 10:00", 1.65, 12682.0, "France"), + # A trailing space on a code that is the same product without it. + (489463, "47503J ", "SET/3 FLORAL GARDEN TOOLS IN BAG", 1, "2009-12-05 12:00", 3.75, 13085.0, "United Kingdom"), + # A non-breaking space before the first word, which Python's split() + # treats as whitespace and an ASCII-only regex would not. + (489464, 20725, "\xa0LUNCH BAG RED RETROSPOT", 2, "2009-12-05 13:00", 1.65, 13085.0, "United Kingdom"), + # A zero quantity is a sale of nothing, not a return. + (489465, 21212, "PACK OF 72 RETROSPOT CAKE CASES", 0, "2009-12-05 14:00", 0.55, 13085.0, "United Kingdom"), + ] + if repeat_a_product: + # The same product twice on one invoice, at a different quantity and price. + rows.append((489434, 85048, "15CM CHRISTMAS GLASS BALL 20 LIGHTS", 5, "2009-12-01 07:45", 6.45, 13085.0, "United Kingdom")) + # The same product on one invoice under two cases of its code. + rows.append((489462, "15056bl", "EDWARDIAN PARASOL BLACK", 2, "2009-12-05 11:00", 5.95, 13085.0, "United Kingdom")) + rows.append((489462, "15056BL", "EDWARDIAN PARASOL BLACK", 3, "2009-12-05 11:00", 5.95, 13085.0, "United Kingdom")) + + frame = pd.DataFrame( + rows, + columns=["Invoice", "StockCode", "Description", "Quantity", "InvoiceDate", "Price", "Customer ID", "Country"], + ) + frame["InvoiceDate"] = pd.to_datetime(frame["InvoiceDate"]) + return frame + + +@pytest.fixture(scope="module") +def spark(): + from pyspark.sql import SparkSession + + session = ( + SparkSession.builder.master("local[1]") + .appName("rootsignal-tests") + .config("spark.sql.session.timeZone", "UTC") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + yield session + session.stop() + + +def load_job(): + spec = importlib.util.spec_from_file_location("conform_online_retail", JOB_PATH) + job = importlib.util.module_from_spec(spec) + spec.loader.exec_module(job) + return job + + +def run_job(spark, raw: pd.DataFrame, tmp_path): + """Land the lines as the pipeline does, run the job's entry point, read its output. + + main() picks up the test's session through getOrCreate, so this is the same + code path spark-submit and Glue run. + """ + lake = Lake(str(tmp_path / "lake")) + write_raw_csv(raw, lake.uri(RAW, RAW_CSV)) + load_job().main(["--input", lake.uri(RAW, RAW_CSV), "--output", lake.uri(CURATED)]) + return read_curated(lake) + + +def assert_tables_match(spark_table: pd.DataFrame, pandas_table: pd.DataFrame, key: list[str]) -> None: + assert sorted(spark_table.columns) == sorted(pandas_table.columns) + left = spark_table.sort_values(key).reset_index(drop=True) + right = pandas_table[list(spark_table.columns)].sort_values(key).reset_index(drop=True) + assert len(left) == len(right) + for column in left.columns: + a, b = left[column], right[column] + if pd.api.types.is_numeric_dtype(a) and pd.api.types.is_numeric_dtype(b) and a.dtype != bool: + np.testing.assert_allclose( + a.astype(float), b.astype(float), rtol=0, atol=1e-9, equal_nan=True, err_msg=column + ) + else: + assert a.astype(str).tolist() == b.astype(str).tolist(), column + + + +# Each Spark run takes tens of seconds, so each fixture is conformed once and +# shared by the tests that read it. + + +@pytest.fixture(scope="module") +def conformed(spark, tmp_path_factory): + raw = raw_lines() + dataset, returns = run_job(spark, raw, tmp_path_factory.mktemp("with_repeats")) + return raw, dataset, returns + + +@pytest.fixture(scope="module") +def conformed_without_repeats(spark, tmp_path_factory): + raw = raw_lines(repeat_a_product=False) + dataset, _ = run_job(spark, raw, tmp_path_factory.mktemp("without_repeats")) + return raw, dataset + + +@pytest.mark.parametrize("table", list(KEYS)) +def test_spark_matches_the_pandas_adapter(conformed, table) -> None: + raw, spark_dataset, _ = conformed + assert_tables_match(spark_dataset.tables[table], adapt(raw=raw).tables[table], KEYS[table]) + + +def test_without_repeated_lines_the_published_price_is_kept(conformed_without_repeats) -> None: + """The adapter only recomputes prices when it had lines to combine.""" + raw, spark_dataset = conformed_without_repeats + sales = spark_dataset.tables["fact_sales"] + assert sorted(sales["unit_price"].tolist()) == sorted( + raw.loc[(raw["Quantity"] >= 0) & ~raw["Invoice"].astype(str).str.startswith("C"), "Price"].tolist() + ) + assert_tables_match(sales, adapt(raw=raw).tables["fact_sales"], KEYS["fact_sales"]) + assert "0 invoice lines repeat a product" in " ".join(spark_dataset.capabilities.notes) + + +def test_half_cent_revenue_rounds_the_way_pandas_rounds(spark, conformed) -> None: + """2 x 0.5025 is 1.005. pandas gives 1.0; Spark's own round() gives 1.01.""" + from pyspark.sql import functions as F + + assert spark.range(1).select(F.round(F.lit(2) * F.lit(0.5025), 2)).first()[0] == pytest.approx(1.01) + + _, spark_dataset, _ = conformed + sales = spark_dataset.tables["fact_sales"].set_index(["order_id", "sku_id"]) + assert sales.loc[("489436", "21755"), "net_sales"] == pytest.approx(1.0) + + +def test_returns_are_split_out_the_same_way(conformed) -> None: + raw, _, spark_returns = conformed + expected = returns_frame(raw=raw) + assert sorted(spark_returns["order_id"]) == sorted(expected["order_id"]) + assert spark_returns["returned_units"].sum() == expected["returned_units"].sum() + + +def test_summary_counts_what_the_job_did(conformed) -> None: + _, spark_dataset, _ = conformed + assert spark_dataset.source_rows == 19 + notes = " ".join(spark_dataset.capabilities.notes) + assert "3 of 19 lines are returns" in notes + assert "2 invoice lines repeat a product" in notes + + +def test_a_zero_quantity_is_a_sale_of_nothing_not_a_return(conformed) -> None: + _, spark_dataset, spark_returns = conformed + sales = spark_dataset.tables["fact_sales"].set_index(["order_id", "sku_id"]) + assert sales.loc[("489465", "21212"), "units"] == 0 + assert "489465" not in set(spark_returns["order_id"]) + + +def test_the_first_value_follows_the_file_not_the_arrival_order(spark) -> None: + """pandas' groupby 'first' is the first non-null value in file order. + + On one partition, rows arrive in file order anyway, so the fixture above + cannot tell file order from arrival order. Here the rows arrive in a + different order from their position in the file, as they can on a cluster. + """ + job = load_job() + rows = [("k", "arrived first, later in the file", 5), ("k", None, 0), ("k", "earliest non-null", 1)] + frame = spark.createDataFrame(rows, "key string, value string, _row long").coalesce(1) + picked = frame.groupBy("key").agg(job.first_in_file_order("value").alias("value")).first()["value"] + assert picked == "earliest non-null" + + +def test_the_job_runs_from_the_command_line(tmp_path) -> None: + """The entry point spark-submit, Glue and EMR use, run as a separate process.""" + from rootsignal.pipeline import conform_locally, read_summary + + lake = Lake(str(tmp_path / "lake")) + write_raw_csv(raw_lines(), lake.uri(RAW, RAW_CSV)) + summary = conform_locally(lake) + assert summary == read_summary(lake) + assert summary["source_rows"] == 19 and summary["returned_lines"] == 3 diff --git a/tests/test_warehouse.py b/tests/test_warehouse.py new file mode 100644 index 0000000..ef413a4 --- /dev/null +++ b/tests/test_warehouse.py @@ -0,0 +1,258 @@ +"""The PostgreSQL warehouse returns what the SQLite database returns. + +The translation tests run everywhere. The warehouse tests need a server and run +when ROOTSIGNAL_WAREHOUSE_URL is set, which CI does with a PostgreSQL service; +locally, `docker compose up -d warehouse` provides one. +""" + +from __future__ import annotations + +import os +import re +import uuid + +import numpy as np +import pandas as pd +import pytest + +from rootsignal.modeling import build_commercial_mart +from rootsignal.sql import available_queries, build_database, query, run_query_file +from rootsignal.sql.database import ANALYTICS_DIR, LOAD_ORDER, MARTS_DIR, SCHEMA_PATH, STAGING_DIR +from rootsignal.sql.dialect import schema_columns, to_postgres, to_redshift +from rootsignal.sql.warehouse import ( + URL_VARIABLE, + build_warehouse, + query_warehouse, + run_warehouse_query_file, +) + +MART_KEY = ["date", "region_code", "category", "channel", "sales_type"] +STAGING_VIEWS = ["stg_sales", "stg_orders", "stg_inventory", "stg_targets", "stg_kam_targets"] + + +# -------------------------------------------------------------------------- +# Translation +# -------------------------------------------------------------------------- + + +def test_week_start_becomes_the_iso_monday() -> None: + sql = "SELECT DATE(m.date, 'weekday 0', '-6 days') AS week_start FROM t" + assert to_postgres(sql) == "SELECT CAST(DATE_TRUNC('week', m.date) AS DATE) AS week_start FROM t" + + +def test_year_month_and_real_are_translated() -> None: + sql = "SELECT STRFTIME('%Y-%m', date), CAST(x AS REAL) FROM t" + assert to_postgres(sql) == "SELECT TO_CHAR(date, 'YYYY-MM'), CAST(x AS DOUBLE PRECISION) FROM t" + + +def test_pragma_is_dropped() -> None: + assert "PRAGMA" not in to_postgres("PRAGMA foreign_keys = ON;\nCREATE TABLE t (a INTEGER);") + + +def test_an_untranslatable_construct_is_refused_not_passed_through() -> None: + with pytest.raises(ValueError, match="JULIANDAY"): + to_postgres("SELECT JULIANDAY(date) FROM t") + + +def test_a_construct_mentioned_only_in_a_comment_is_not_refused() -> None: + assert to_postgres("-- STRFTIME( is SQLite only\nSELECT 1") == "-- STRFTIME( is SQLite only\nSELECT 1" + + +@pytest.mark.parametrize( + "path", + [SCHEMA_PATH, *sorted(STAGING_DIR.glob("*.sql")), *sorted(MARTS_DIR.glob("*.sql")), + *sorted(ANALYTICS_DIR.glob("*.sql"))], + ids=lambda path: path.stem, +) +def test_every_sql_file_in_the_repository_translates(path) -> None: + """A new query using an unhandled construct fails here, not on the server.""" + to_postgres(path.read_text(encoding="utf-8")) + + +def test_redshift_translation_removes_checks_and_nothing_else() -> None: + sql = ( + "CREATE TABLE t (\n" + " a TEXT NOT NULL CHECK (a IN ('x', 'y')),\n" + " b REAL NOT NULL CHECK (b >= 0 AND (b <= 1)),\n" + " c INTEGER\n" + ");" + ) + assert to_redshift(sql) == ( + "CREATE TABLE t (\n" + " a VARCHAR(1024) NOT NULL,\n" + " b DOUBLE PRECISION NOT NULL,\n" + " c INTEGER\n" + ");" + ) + + +def test_schema_columns_reads_names_and_types_in_order() -> None: + declared = schema_columns(SCHEMA_PATH.read_text(encoding="utf-8")) + assert set(declared) == set(LOAD_ORDER) + assert declared["dim_date"] == [ + ("date", "DATE"), ("week", "INTEGER"), ("month", "INTEGER"), ("quarter", "TEXT"), ("year", "INTEGER") + ] + + +def _sql_code(path) -> str: + return re.sub(r"--[^\n]*", "", path.read_text(encoding="utf-8")) + + +ALL_SQL = [*sorted(STAGING_DIR.glob("*.sql")), *sorted(MARTS_DIR.glob("*.sql")), *sorted(ANALYTICS_DIR.glob("*.sql"))] + + +@pytest.mark.parametrize("path", ALL_SQL, ids=lambda path: path.stem) +def test_every_average_is_of_a_real_number(path) -> None: + """Redshift's AVG of an integer column is an integer. + + SQLite and PostgreSQL return a fraction, so the warehouse tests cannot see + the difference; a stockout rate over 0/1 flags would read 0 on Redshift + alone. Every AVG therefore averages a cast, or a column declared REAL. + """ + real_columns = { + column for columns in schema_columns(SCHEMA_PATH.read_text(encoding="utf-8")).values() + for column, kind in columns if kind == "REAL" + } + code = _sql_code(path) + for match in re.finditer(r"\bAVG\(\s*", code, re.IGNORECASE): + rest = code[match.end():] + if rest.upper().startswith("CAST("): + continue + column = re.match(r"([\w.]+)\s*\)", rest).group(1).split(".")[-1] + assert column in real_columns, f"{path.stem}: AVG({column}) averages an integer column" + + +# The columns that identify one output row of each query. Sorting on anything +# less leaves tied rows in whatever order the engine produces, and SQLite and +# PostgreSQL produce different ones. +QUERY_GRAIN = { + "daily_sales_tracker": ["date"], + "fill_rate_by_segment": ["week_start", "region_code", "category"], + "kam_scorecard": ["kam_id"], + "primary_secondary_mix": ["month", "region_code"], + "supply_watchlist": ["week_start", "region_code", "category"], + "target_variance": ["region_code", "category", "channel"], + "weekly_movers": ["week_start", "region_code", "category"], +} + + +def test_every_query_has_a_declared_grain() -> None: + assert set(QUERY_GRAIN) == set(available_queries()) + + +@pytest.mark.parametrize("name", sorted(QUERY_GRAIN)) +def test_every_query_sorts_on_its_full_grain(name) -> None: + code = _sql_code(ANALYTICS_DIR / f"{name}.sql") + order_by = code[code.upper().rindex("ORDER BY"):] + sorted_on = {term.strip().split()[0].split(".")[-1] for term in order_by[len("ORDER BY"):].rstrip("; \n").split(",")} + missing = [column for column in QUERY_GRAIN[name] if column not in sorted_on] + assert not missing, f"{name} does not sort on {missing}" + + +# -------------------------------------------------------------------------- +# The warehouse against the SQLite database +# -------------------------------------------------------------------------- + +needs_warehouse = pytest.mark.skipif( + not os.environ.get(URL_VARIABLE), + reason=f"{URL_VARIABLE} is not set; start one with `docker compose up -d warehouse`", +) + + +@pytest.fixture(scope="module") +def engines(tmp_path_factory, generated_dataset_dir): + """Both engines loaded from the same cleaned tables, in a throwaway schema.""" + from rootsignal.cleaning import clean_dataset + from rootsignal.ingestion import load_dataset + + tables = clean_dataset(load_dataset(generated_dataset_dir)).tables + schema = f"test_{uuid.uuid4().hex[:10]}" + warehouse = build_warehouse(tables, schema=schema) + yield tables, build_database(tables), warehouse + warehouse.rollback() + warehouse.execute(f'DROP SCHEMA "{schema}" CASCADE') + warehouse.commit() + warehouse.close() + + +def _assert_same(left: pd.DataFrame, right: pd.DataFrame, tolerance: dict[str, float] | None = None) -> None: + """Same columns, same row count, same values in the same order.""" + tolerance = tolerance or {} + assert list(left.columns) == list(right.columns) + assert len(left) == len(right) + left, right = left.reset_index(drop=True), right.reset_index(drop=True) + for column in left.columns: + a, b = left[column], right[column] + if pd.api.types.is_numeric_dtype(a) and pd.api.types.is_numeric_dtype(b): + np.testing.assert_allclose( + a.astype(float).to_numpy(), + b.astype(float).to_numpy(), + rtol=0, + atol=tolerance.get(column, 1e-9), + equal_nan=True, + err_msg=column, + ) + else: + assert a.astype(str).tolist() == b.astype(str).tolist(), column + + +@needs_warehouse +@pytest.mark.parametrize("view", STAGING_VIEWS) +def test_staging_views_match_sqlite(engines, view) -> None: + _, sqlite, warehouse = engines + lite = query(sqlite, f"SELECT * FROM {view}") + pg = query_warehouse(warehouse, f"SELECT * FROM {view}") + key = list(lite.columns) + _assert_same(lite.sort_values(key).reset_index(drop=True), pg.sort_values(key).reset_index(drop=True)) + + +@needs_warehouse +def test_warehouse_mart_matches_the_python_mart(engines) -> None: + """The warehouse is held to the same definition as the SQLite mart. + + aov may differ by one cent on exact half-cent values. The Python mart rounds + half to even, SQLite rounds the binary float, and PostgreSQL rounds the + decimal value half up, so 79.065 is 79.06, 79.06 and 79.07. Every other + column must agree exactly. + """ + tables, _, warehouse = engines + pg = query_warehouse(warehouse, "SELECT * FROM mart_commercial_daily") + python = build_commercial_mart(tables) + python["date"] = python["date"].astype(str) + pg = pg.sort_values(MART_KEY).reset_index(drop=True) + python = python.sort_values(MART_KEY).reset_index(drop=True)[pg.columns] + _assert_same(pg, python, tolerance={"aov": 0.01 + 1e-9}) + + +@needs_warehouse +@pytest.mark.parametrize("name", available_queries()) +def test_every_analytical_query_matches_sqlite_row_for_row(engines, name) -> None: + """Same rows in the same order: the queries sort on a full key.""" + _, sqlite, warehouse = engines + _assert_same(run_query_file(sqlite, name), run_warehouse_query_file(warehouse, name)) + + +@needs_warehouse +def test_a_column_of_nulls_comes_back_as_numbers_not_objects(engines) -> None: + """psycopg returns None for NULL, which pandas would hold as an object column.""" + _, _, warehouse = engines + frame = query_warehouse( + warehouse, "SELECT CAST(NULL AS DOUBLE PRECISION) AS x UNION ALL SELECT CAST(NULL AS DOUBLE PRECISION)" + ) + assert frame["x"].dtype == float and frame["x"].isna().all() + + +@needs_warehouse +def test_a_load_that_breaks_a_foreign_key_leaves_the_warehouse_untouched(engines) -> None: + tables, _, warehouse = engines + schema = warehouse.execute("SELECT current_schema()").fetchone()[0] + before = warehouse.execute("SELECT COUNT(*) FROM fact_sales").fetchone()[0] + warehouse.rollback() + + broken = {name: frame.copy() for name, frame in tables.items()} + broken["fact_sales"].loc[broken["fact_sales"].index[0], "sku_id"] = "NO-SUCH-SKU" + with pytest.raises(Exception, match="foreign key"): + build_warehouse(broken, schema=schema) + + assert warehouse.execute("SELECT COUNT(*) FROM fact_sales").fetchone()[0] == before + warehouse.rollback()