Skip to content

Repository files navigation

feature-store

A feature store with point-in-time correct training data, a streaming materializer, a serving API, and drift monitoring — built around the one bug feature stores exist to prevent.

pip install -r requirements-dev.txt
python pipelines/generate.py --cards 1200 --days 45     # event stream
python pipelines/train.py                               # measures the leakage gap
python pipelines/materialize.py                         # stream -> online store, drift
python pipelines/leakage_study.py                       # when leakage matters
pytest                                                  # 44 tests

The bug this is about

You want to train a fraud model. For a transaction at 14:32 you want features like "how many times has this card been used in the last hour". The obvious implementation groups the event table by card and joins it back:

features = events.groupby("card_id").agg({"amount": ["count", "sum", "mean"]})
training = events.merge(features, on="card_id")     # this is the bug

That aggregate includes events that happened after 14:32. The model learns from information that did not exist at prediction time. Offline AUC goes up, production AUC does not, and nothing anywhere throws an error.

The fix is an as-of join: for every (entity, timestamp) pair, aggregate only events strictly before that timestamp. Both implementations live in fs/pit.py — the correct one and the leaky one — because being able to measure the difference is what makes the argument.

What leakage actually costs

Same model, same aggregations, same time-based train/test split. The only variable is time discipline. 65,190 transactions over 45 simulated days, 2.2% fraud:

pipeline AUC avg precision
point-in-time correct 0.6664 0.0533
leaky, evaluated on leaky features 0.6583 0.0535
leaky model, evaluated on real features 0.5989 0.0421

The middle row is what a leaky pipeline puts in the slide deck. The bottom row is what that same model does when production hands it features that can only see the past. The offline number is inflated by 0.059 AUC, and shipping it instead of doing the join properly costs 0.068 AUC.

The third row is the one nobody computes. A leaky pipeline evaluates its leaky model on leaky test features and never finds out.

When does leakage matter?

One number from one dataset is an anecdote. Leakage inflates scores by letting the model see the future, so the damage should scale with how much of the real signal lives in the historical features. That's a testable prediction, so leakage_study.py sweeps the fraud regime, 3 seeds each:

regime fraud rate history adds offline inflation
burst-dominated 0.53% +0.0197 AUC +0.0564
mixed 2.22% -0.0053 AUC +0.0447
noise-dominated 5.41% +0.0016 AUC +0.0173

"history adds" is pit_auc - request_only_auc: how much the stored features contribute over the transaction's own attributes. The relationship holds — when the signal is in the card's recent behaviour, leaking that behaviour inflates the score three times as much as when it isn't.

Worth being blunt: history contributes little in absolute terms here (and is slightly negative in one regime). The generated fraud is mostly detectable from the transaction itself. That weakens the demo and it's what the measurement says, so it goes in the README rather than being tuned away.

Train/serve consistency

The other way these systems fail: the batch job that builds training data and the streaming job that feeds serving compute the same feature slightly differently. The model then trains on one distribution and scores on another, and again nothing errors.

This isn't solved by testing after the fact, it's solved by not having two implementations. Aggregations are declared as data in fs/definitions.py, and the batch and stream engines both call the same _apply:

@dataclass(frozen=True)
class Aggregation:
    name: str
    agg: Agg                      # COUNT | SUM | MEAN | MIN | MAX | STDDEV | DISTINCT
    field: str | None = None
    window_seconds: int = 3600

Same for the request-time encoding and the on-demand transforms. Training and serving both call one build_row, so there is nothing to keep in sync.

The tests then check the invariant actually holds:

  • streamed features match a batch replay after every single event, not just at the end
  • micro-batch size (1, 7, 1000 events) doesn't change any value
  • both online backends produce identical values
  • and the end-to-end one: a transaction's training vector and its serving vector are compared element by element
def test_training_and_serving_produce_identical_vectors():
    stored_offline = point_in_time_join([VIEW], {"v": log}, [(card, ts)])[0]
    train_vector = [build_row(scored, stored_offline).get(n, 0.0) for n in FEATURES]

    for ev in events:
        engine.apply(ev, ev["ts"])          # stream into the online store
    body = client(ctx).post("/predict", json=scored).json()
    serve_vector = [body["features"].get(n, 0.0) for n in FEATURES]

    assert serve_vector == pytest.approx(train_vector, rel=1e-9, abs=1e-9)

The off-by-one

Batch computes features excluding the row being scored. Stream folds an event in and reports state including it. Both are right for their purpose, and comparing them means nudging the batch as-of just past the event. Getting this edge wrong is the leak in miniature — at serving time the transaction hasn't happened yet, so training must not count it either.

Monitoring, and what it can't see

materialize.py replays the whole stream and then checks two things.

Freshness. A dead materializer looks exactly like a quiet one from the read path. Every online row carries its event timestamp, and serving compares that against the view's TTL. On the 45-day replay the monitor reports 332 of 500 cards stale against a 6h TTL, which is correct — most cards' last transaction really was days before the end of the window. In strict mode /predict returns 503 rather than scoring on state that old.

Drift. PSI with quantile bins (equal-width bins on a skewed feature put everything in one bucket and PSI stops seeing anything), plus a KS statistic.

And here's the useful part:

0 of 11 features drifted: none

label rate: 1.555% -> 2.850% (1.83x)
  concept drift with no feature drift: the inputs look the same,
  the relationship changed. only delayed labels can catch this.

The fraud rate nearly doubled and PSI saw nothing. That's not a bug in the monitor, it's a structural limit: feature drift measures P(x) moving, and this is P(y|x) moving. Every input distribution is unchanged; what changed is what those inputs mean. A dashboard that only watches feature drift will show all green through exactly this failure, which is why the pipeline reports the label rate next to it.

Serving

uvicorn fs.serving:build_default_app --factory --port 8000

POST /predict assembles the vector through the same build_row training uses, and refuses to be quiet about problems:

{
  "card_id": "card_00042",
  "fraud_probability": 0.31,
  "features": {"card_velocity__txn_count_1h": 6.0, "od__velocity_ratio": 14.4, ...},
  "warnings": ["card_spend: features are 91230s old, ttl is 21600s"],
  "took_ms": 0.41
}

Two guards worth calling out:

  • staleness — with FS_STRICT_TTL=1 a stale feature is a 503, not a confident prediction on week-old state
  • version skew — each feature view hashes its own definition. A model trained against v1 must not be served features materialized under v2, because amount_1h is a name plus a window plus an aggregation, and only the name is stable

Also POST /features, GET /registry, GET /metrics, GET /health.

On-demand transforms

The interesting fraud features are ratios. "£400" means nothing; "£400 on a card whose 7-day average is £22" means a great deal. The amount arrives in the request, the average comes from the store, and the feature is a function of both — so it can't be precomputed and has to run at request time:

@ondemand.add("od__amount_zscore_7d",
              ("card_spend__amount_mean_7d", "card_spend__amount_std_7d"))
def _amount_z(request, stored):
    return (request["amount"] - stored["..._mean_7d"]) / (stored["..._std_7d"] + 1.0)

This is exactly where train/serve skew is most likely, which is why there's one registry and both paths read from it. A missing stored input yields a documented 0.0 fallback that serving reports as a warning rather than silently passing off as a real value.

Layout

fs/definitions.py   entities, aggregations, feature views, versioned registry
fs/pit.py           as-of join (correct) and groupby join (leaky), side by side
fs/stream.py        append-only log with consumer offsets, streaming engine
fs/online.py        online store: in-memory and SQLite, both reject stale writes
fs/ondemand.py      request-time transforms over stored features
fs/serving.py       FastAPI, TTL and version-skew guards
fs/monitoring.py    PSI, KS, freshness
pipelines/          generate, materialize, train, leakage_study
docker/             image + compose for the materialize -> serve path

Being straight about what this is

  • The data is synthetic, deliberately. Proving that a leaky pipeline inflates its score requires knowing exactly which signals a model is entitled to see at each instant. A public fraud dataset gives you labels but not that, so a gap measured on one is an argument rather than a proof. The generator does have the properties that make time correctness hard: fraud bursts on compromised cards, a diurnal transaction rhythm, per-card lognormal amounts, and a fraud rate that drifts upward.
  • AUC ~0.67 is not a good fraud model. It isn't meant to be. The deliverable is the gap between pipelines, and both sides of every comparison use the same model and the same features.
  • No Kafka, no Redis, no Postgres. EventStream implements the semantics that matter (ordered log, offsets, independent consumer groups) in process, and OnlineStore has in-memory and SQLite backends. Both are interfaces a real backend slots behind. I didn't write those adapters, because shipping an adapter I have no way to run is worse than not shipping one — Docker isn't available on the machine this was built on.
  • The stream engine keeps raw events in each window and re-aggregates, rather than maintaining O(1) incremental sketches. That's what makes batch and stream values identical by construction instead of by luck. It's the right trade at this scale and the wrong one for very large windows, where you'd move to incremental aggregates and accept approximation on stddev and distinct counts.
  • Backfill is a full replay. There's no incremental materialization, no compaction, and no partial-window checkpointing.

Tests

44 tests, pytest.

Point-in-time correctness: the window excludes the row itself, excludes events after it, excludes events older than the window; entities don't bleed into each other; insertion order doesn't matter; recomputation is stable; and the leaky join is caught red-handed reading a value from 8,700 seconds in the future.

Consistency: stream matches batch after every event, across micro-batch sizes, across both online backends, and end to end from training vector to HTTP response.

Serving: stale features warn, strict mode returns 503, version skew is flagged, unknown entities produce a warning rather than a fabricated zero.

Plus the operational edges — window eviction keeps state bounded, late events are counted rather than hidden, and a late write never clobbers fresher state.

About

Feature store with point-in-time correct training data, streaming materialization, train/serve consistency guarantees, and drift monitoring

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages