12. Data engineering for machine learning

ML data engineering protects one invariant: at training time the model must see exactly the information it will have at prediction time, computed exactly the same way. The second duty is being able to rebuild any training set, months later, byte for byte.

How a training set is built

Start with a spine: one row per training example, holding the entity ID, the prediction timestamp and the label. Then attach features to each row using only information available at that timestamp. The spine fixes the grain; the joins fill in the columns.

Leakage: the four kinds

Leakage always shows up the same way: offline metrics look excellent, production results disappoint. Treat suspiciously good results as a bug report.

Splitting data

• Random split: valid only when examples are truly independent and the future looks like the past.
• Time-based split: train on the past, validate and test on later periods. This mirrors deployment and is the default for anything time-dependent.
• Group split: keep all rows of one user, patient or document in the same split.
• Stratified split: preserve class proportions in each split, important for rare classes.

The test set is used once, at the end. Every time you look at it and adjust something, it becomes part of training in spirit.

Point-in-time correctness

For each spine row, a feature must take the latest value whose timestamp is at or before the prediction time. This is an as-of join. Many engines support it directly; in DuckDB:

SELECT s.customer_id, s.prediction_ts, s.label,
       f.plan, f.orders_last_30d
FROM spine s
ASOF LEFT JOIN customer_features f
  ON s.customer_id = f.customer_id
 AND s.prediction_ts >= f.available_at;

Note the join is on available_at, not on when the event happened. If yesterday’s order counts are only computed at 03:00 today, a prediction at 01:00 could not have used them. Modeling that availability lag is what separates correct pipelines from subtly leaky ones.

Training-serving skew

Skew happens when the features a model sees in production differ from those it was trained on. Typical causes: training features computed in SQL and serving features reimplemented in another language; different source systems; different freshness; different handling of nulls.

Three defences, from simplest to most robust:

1. Share one feature definition that generates both the training and the serving values.
2. Use a feature store that materializes the same definition into both an offline and an online store.
3. Log and wait: log the exact feature values used at serving time, and train future models on those logs.

Feature stores

A feature store has four parts:

• A registry of feature definitions with owners and versions.
• An offline store holding feature history, used to build training sets with as-of joins.
• An online store, typically a key-value database, holding the latest values for millisecond lookups at inference.
• Materialization jobs that compute each feature once and write it to both stores.

Examples include Feast (open source), Tecton, Hopsworks, and the feature stores built into Databricks, Vertex AI and SageMaker. You likely do not need one if you have a few models with batch predictions; you likely do if many models share features and serve online.

Batch versus online prediction

• Batch prediction scores everything on a schedule and writes results to a table. It is simple and cheap, but predictions are only as fresh as the last run.
• Online prediction scores on request. It needs features available within a latency budget, which is where online stores and streaming features come in.

Reproducibility and dataset versioning

To reproduce a model you need five things recorded together: the data version, the code commit, the configuration, the random seeds and the software environment. Data is the hardest of the five.

Whatever tool you use, the test is simple: given a model ID, can you produce the exact training data?

Labeling pipelines

• Annotation tools such as Label Studio, Argilla and CVAT, with written guidelines and examples.
• Quality control: overlapping assignments, agreement metrics and gold questions.
• Active learning: send the examples the current model is least sure about to humans first.
• Weak supervision: write labeling functions (heuristics, rules) and combine their noisy votes, as in Snorkel.
• LLMs as labelers: fast and cheap, but validate against a human-labeled sample, and record which model and prompt produced each label.

Version labels like data: a relabeling campaign creates a new dataset version.

Sampling and imbalance

Downsampling the common class makes training cheaper and balances classes, but it changes the base rate the model learns. Record the sampling scheme in the dataset metadata, and correct predicted probabilities or apply example weights accordingly.

After deployment

Log every input, prediction and model version. Join in the true labels when they arrive, which may be days or months later. Compute metrics per slice (country, device, customer segment), because an average can hide a segment that collapsed.

Questions to ask

• At the moment of prediction, was every feature value actually available?
• Is the split random, time-based or grouped, and does it match deployment?
• Is the serving feature computed by the same code as the training feature?
• Given this model ID, can I rebuild its exact training data?
• How were the labels produced, and how good are they?