Engineering Report
PySpark · Apache Airflow 3 · Great Expectations · XGBoost · MLflow · Docker
This report documents the commissioning of a customer propensity to purchase pipeline, rebuilt from a single node pandas and XGBoost workflow onto a distributed, orchestrated architecture. PySpark performs extract, transform and load and feature engineering, Apache Airflow 3 orchestrates the workflow as a directed acyclic graph, Great Expectations Core acts as a data quality gate, and XGBoost with MLflow handles training, evaluation and the model registry. The entire stack runs in Docker on a single Apple Silicon workstation.
The codebase was delivered as a scaffold that had never executed against a live Docker daemon. Bringing it to a first successful run surfaced eight distinct defects spanning infrastructure, orchestration, distributed compute, data quality and the machine learning framework layer. Every defect was root caused from primary evidence and corrected at source rather than worked around. The pipeline now completes end to end on the full 42.4 million event month, producing a trained model registered for production.
Starting from a machine with no Docker, no Java and no Python data stack, the pipeline was taken to a verified green run on the full month of REES46 data. The deployed model separates buyers from non buyers with a validation AUC of 0.935, and is registered in MLflow as propensity-model version 1 under the production alias.
The goal is the orchestration and distributed processing layer, not a new model. The brief is explicit that the modelling approach and XGBoost hyperparameters should stay close to a standard pandas pipeline. The claim being supported is migrated to a distributed, orchestrated architecture, not built a better model. A clean, demonstrable, green pipeline run is therefore the primary deliverable, and model quality is treated as a gate to pass.
The source is the REES46 e-commerce behavior dataset from Kaggle, processed one monthly file at a time. The October 2019 file was used: a 5.3 GB CSV containing 42,448,764 behavioural events. Each row is a single event (view, cart, purchase, or remove from cart) carrying a timestamp, product and category identifiers, a brand, a price, a user id and a session id.
| Layer | Component | Version | Role |
|---|---|---|---|
| Orchestration | Apache Airflow | 3.2.2 | DAG scheduling, Task Execution API, web UI |
| Distributed compute | PySpark | 4.1.2 | CSV to Parquet, feature engineering in local[*] mode |
| Runtime | OpenJDK / Python | 17 / 3.13 | JVM for Spark, Python for tasks |
| Data quality | Great Expectations Core | 1.18.1 | Semantic expectations on the feature table |
| Modelling | XGBoost | 3.3.0 | Gradient boosted classifier |
| ML lifecycle | MLflow | 3.14.0 | Tracking, metrics, model registry, aliases |
| Metadata DB | PostgreSQL | 16 | Airflow metadata backend |
The workflow is a strict linear chain of seven tasks. A failure at any stage stops the run and prevents bad data or a weak model from progressing.
| # | Task | Engine | Responsibility |
|---|---|---|---|
| 1 | ingest | PySpark | Read raw CSV, stage as Parquet with an explicit schema |
| 2 | validate | PySpark | Structural checks: required columns, non empty, allowed event types |
| 3 | transform | PySpark | Aggregate events to one row per session, build features and target |
| 4 | quality_gate | GX Core | Semantic checks: uniqueness, nulls, value ranges on features |
| 5 | train | XGBoost | Bridge to pandas, train classifier, log to MLflow |
| 6 | evaluate | MLflow | Read logged AUC, gate on a 0.90 threshold |
| 7 | deploy | MLflow | Register model version, point the production alias at it |
Airflow XCom is backed by the metadata database and is not built for bulk data. Every task writes its output to disk and passes only a small string (a file path or an MLflow run id) to the next task.
Eight defects were identified and corrected, each root caused from primary evidence (container logs, supervisor signals, tracebacks) and fixed at source. They are ordered roughly by the sequence in which they blocked the pipeline.
Cause. The DAG used a try and except fallback import. The airflow.sdk path is the canonical TaskFlow import for 3.2.2.
Fix. Imported directly from airflow.sdk, verified by listing the DAG with no import errors.
Symptom. The ingest task was killed within seconds. A Spark smoke test reported JAVA_GATEWAY_EXITED, the Java gateway process exited before sending its port number.
Cause. The Dockerfile set JAVA_HOME to the amd64 path, but on Apple Silicon the JDK installs to the arm64 path. PySpark launches the JVM via $JAVA_HOME/bin/java, which did not exist.
Fix. An architecture aware symlink created at build time with dpkg --print-architecture, pointing a stable default-java path at the correct JDK.
Symptom. Every task was terminated by the supervisor with SIGKILL about nine seconds in, after repeated Connection refused retries.
Cause. Airflow 3 runs all tasks through the Task Execution API. The task defaulted to localhost:8080, where no api-server listens inside the scheduler container, so it could not call home and was killed. A shared JWT secret was also missing.
Fix. Added EXECUTION_API_SERVER_URL pointing at the api-server service and a shared JWT_SECRET across all services, mirroring the official Airflow 3.2.2 compose.
Cause. Real timestamps are formatted "2019-10-01 00:00:00 UTC" while fixtures use ISO. to_timestamp with no format returns null on the UTC text, which would null session duration and cause training to drop every row.
Fix. Parse both formats explicitly with a coalesce, so a format mismatch can never silently null the timestamps.
Symptom. transform failed after about twenty seconds with py4j errors as the Spark JVM died mid shuffle.
Cause. In local[*] mode the driver JVM is also the executor, and the default 1 GB heap cannot shuffle 42 million rows through the groupBy and countDistinct aggregations.
Fix. Configured spark.driver.memory to 4 GB, which fits inside the 7.8 GB Docker VM alongside the Airflow processes.
Symptom. quality_gate failed: the uniqueness and not null expectations on the session key did not hold.
Cause. The aggregation grouped by session and user id together, emitting duplicate session rows for the ~350 REES46 sessions carrying more than one user id, plus two null sessions. The gate correctly caught both.
Fix. Aligned the code with its stated contract: drop null sessions, group by session alone, keep a single representative user id. The gate now passes honestly rather than being weakened.
Cause. MLflow 3 raises by default on the file store and removed the old model stage API. Both would have failed the last two tasks.
Fix. Set MLFLOW_ALLOW_FILE_STORE and rewrote deploy to register the model and point a production alias with set_registered_model_alias.
Symptom. evaluate failed with "the key run_id in args is reserved", appearing only once every upstream task was finally green.
Cause. Airflow injects run_id as a reserved task context key. A TaskFlow parameter of the same name collided during keyword resolution and raised before the function body ran.
Fix. Renamed the parameter to train_run_id and documented the constraint.
The month of raw events was aggregated to one row per user session. The Great Expectations gate is not decoration: on the real data it caught two null and 350 duplicate session keys, resolved at the feature engineering layer.
| Stage | Rows | Note |
|---|---|---|
| Raw events (CSV) | 42,448,764 | view, cart, purchase, remove_from_cart |
| Session features (pre clean) | 9,244,772 | included 2 null and 350 duplicate session keys |
| Session features (modelled) | 9,244,421 | unique, non null sessions used for training |
Purchase is a rare event: 6.81 percent of sessions are buyers against 93.19 percent non buyers, which is why training applies scale_pos_weight. Comparing mean feature values shows where the signal lives. Buyers add to cart far more often and, notably, spend less time per session, consistent with decisive, intent driven visits.
| Feature | Buyers (mean) | Non buyers (mean) | Read |
|---|---|---|---|
| cart_count | 0.792 | 0.050 | about 16x higher for buyers |
| intent_score | 3.169 | 1.727 | strong separation |
| engagement_score | 6.469 | 5.183 | buyers more engaged |
| view_count | 5.561 | 4.327 | buyers view slightly more |
| session_duration_seconds | 658.7 | 1069.9 | buyers spend less time, decisive |
The model is an XGBoost classifier kept deliberately close to a standard pandas configuration: 200 estimators, max depth 6, learning rate 0.1, positive class up weighted to counter the 6.81 percent base rate, evaluated by AUC on an 80 / 20 stratified split.
Validation AUC 0.935, comfortably above the 0.90 deployment gate. The model is registered as propensity-model version 1 with the production alias.
| Rank | Feature | Gain |
|---|---|---|
| 1 | cart_count | 15,541 |
| 2 | session_duration_seconds | 12,747 |
| 3 | view_count | 5,743 |
| 4 | distinct_brands_viewed | 4,021 |
| 5 | engagement_score | 1,116 |
The registered production model was loaded and used to score real sessions. The table contrasts known buyers with known non buyers. Aggregated over the whole month, the mean predicted value is 0.86 for actual buyers against 0.16 for non buyers, a clear separation consistent with the 0.935 AUC.
| Session | views | carts | intent | score | actual |
|---|---|---|---|---|---|
| b1e86763 | 3 | 0 | 1 | 1.00 | buyer |
| 500e0542 | 1 | 1 | 3 | 1.00 | buyer |
| e3d7ff27 | 17 | 0 | 3 | 1.00 | non buyer (false positive) |
| 7de68edd | 4 | 0 | 3 | 0.00 | buyer (false negative) |
| 664c1a7a | 1 | 0 | 1 | 0.00 | non buyer |
On a 500,000 session sample, at the deployed decision threshold: recall 0.861, precision 0.283, F1 0.426. The high recall and modest precision are the expected shape for rare event propensity, the model catches most buyers at the cost of a number of false positives.
The deployed model returns class labels (0 or 1), not calibrated probabilities, which is why every score above is exactly 0.00 or 1.00 and precision and recall do not move with the threshold. The genuine ranking metric, AUC 0.935, comes from predict_proba at training time. For a true, tunable propensity score, deploy should expose predict_proba. This is captured as a recommendation rather than presented as a probability.
| Priority | Recommendation | Rationale |
|---|---|---|
| High | Expose predict_proba in deploy and add a predict task | Turn the binary label into a tunable propensity score and a visible predictions output |
| Medium | Add an MLflow UI service to the compose stack | Browse runs, metrics and the registry visually |
| Medium | Add task level retries to the DAG | Self heal transient supervisor or store hiccups |
| Medium | Calibrate the classifier and tune the decision threshold | Lift precision at the chosen operating point |
| Low | Parameterise the month and process multiple files | Scale beyond a single October snapshot |
The rebuild achieves its stated objective. A pandas scale workflow now runs as a distributed, orchestrated, quality gated pipeline on tens of millions of events, with training, evaluation and registry handled through an industry standard ML lifecycle tool. The defects encountered were characteristic of moving a scaffold onto real infrastructure, and all were resolved at source and documented.