Feature Store: Point-in-Time Joins, Backfills and Skew

Feature Store: Point-in-Time Joins, Backfills and Skew

You've joined the ML platform team that owns Halcyon, the company's feature store. Data scientists write feature values into it, build training sets by joining those values onto labelled examples, and the product reads the same features online. Last week a churn model looked fantastic offline and fell apart in production; the post-mortem (ML-1840) found the training join was using data that did not exist yet at label time. Two more bugs turned up while digging, and three features are queued behind the fixes. Timestamps and feature values are integers.

Codebase

How your code is run

run(operations: list[list]) -> list executes the operations in order against a fresh store and returns one result per operation. Every operation is [name, ts, ...args]; ts is the integer operation clock. It never decreases, and it is also the ingestion time of anything the operation writes. Errors are returned as strings, never raised.

• Each error is "error: <code>", e.g. "error: unknown feature". An unknown operation name is "error: unknown operation". A wrong number of arguments or a non-integer ts is "error: invalid arguments". Integers never include true/false; names (features, entities, snapshots) are non-empty, non-blank strings, case-sensitive.
• When a rule below lists several errors, they are checked in the order listed.
• Operations that fail change nothing.

Records and the point-in-time rule

A stored record has event_ts (when the value was true in the world), ingested_ts (when the store learned about it), value and an arrival number seq (1, 2, 3, ... across the whole store). A feature value is an integer.

For a lookup of (entity, feature) as of time T (a training label time, or "now" when serving), a record is visible if event_ts <= T and ingested_ts <= T (both inclusive). Using anything else is *leakage*: the model would train on facts that nobody could have known at T.

A raw feature's value as of T is determined like this:

1. Among the visible records pick the one with the greatest event_ts. If several share it, the one that arrived last (greatest seq) wins. Corrections are just new records for the same event_ts.
2. If there is none, or T - event_ts > ttl, the value is the feature's default. A value is valid for exactly ttl seconds: T - event_ts == ttl is still valid.

Operations (Part 1)

• ingest records ingested_ts = ts. event_ts > ts is event in future (event_ts == ts is fine). Old event_ts values are fine: late data is normal.
• labels is a list of [entity_id, label_ts] pairs, or a label grid that expands to many labels: {"entities": [...], "first": t0, "step": dt, "count": n} means timestamps t0, t0 + dt, ... (n of them, dt and n positive), and at each timestamp every entity in the order listed. Any label_ts > ts is label in future. features is a non-empty list of feature names (duplicates allowed); every name is checked before any label.
• Each row is [entity_id, label_ts, value_1, ..., value_k], one per label in label order, with the features' values in the order requested.

Example

[["register", 1, "avg_basket", 100, 0],
 ["ingest", 20, "u1", "avg_basket", 10, 40],
 ["ingest", 60, "u1", "avg_basket", 50, 55],
 ["ingest", 90, "u1", "avg_basket", 52, 99],
 ["training_set", 200, [["u1", 5], ["u1", 25], ["u1", 85], ["u1", 95], ["u1", 152], ["u1", 153]], ["avg_basket"]]]

returns

[true, true, true, true,
 [["u1", 5, 0], ["u1", 25, 40], ["u1", 85, 55], ["u1", 95, 99], ["u1", 152, 99], ["u1", 153, 0]]]

The record with event time 52 arrived at 90: the label at 85 must not see it; the label at 95 must.

Part 1: Fix the training join (ML-1840, ML-1852, ML-1861)

Goal: make training_set follow the point-in-time rule above. Reports from the team:

1. "Training rows contain values that were ingested after the label time." (ML-1840, the churn model)
2. "A value is dropped one second too early; the docs say it lives exactly ttl seconds." (ML-1852)
3. "We re-ingested a corrected value for the same event_ts and the old one still wins." (ML-1861)

tests/test_pit.py::test_value_ingested_after_the_label_is_not_visible currently fails. The rules for ties, visibility and TTL are in the section above; fix them for every lookup, not just the example.

Part 2: Late data, backfills, snapshots and online serving (ML-1893)

Goal: data scientists need to repair history and freeze training sets. Late ingests keep their real ingestion time, so they never change the past. A backfill is a re-run of the pipeline that produces corrected history, so its records count as available from their own event time. Frozen snapshots must not move when history is later repaired, and the online path reads a separate, explicitly published store.

Rules

1. backfill's points is a non-empty list of [event_ts, value]. Each point becomes a record with ingested_ts = event_ts (so it is visible to any label at or after its event time, even one that was built before the backfill ran). Points are written in list order, so for equal event_ts the later point wins. Arrival order, not ingested_ts, breaks ties (see the point-in-time rule): a backfilled correction beats the original it corrects even though its ingested_ts is earlier. The operation is atomic: any point with event_ts > ts is event in future and nothing is written.
2. backfill_series generates count points first_ts + i * step with value base + i % 10 (i = 0..count-1; step and count positive integers) and then behaves exactly like backfill.
3. materialize runs training_set (same labels forms, same errors) and stores the resulting rows under name. Versions are 1, 2, ... per name; materializing an existing name creates the next version. checksum is the sum of every feature value in the rows (not entity ids or label times). A stored snapshot never changes, whatever is ingested or backfilled later.
4. read_snapshot returns a copy of the stored rows of that version; version: null means the latest. A version must be null or an integer >= 1, else invalid arguments; an absent name is unknown snapshot; a version that does not exist is unknown version.
5. The online store starts empty. publish copies every record written offline (ingests and backfills) since the previous publish, keeping their seq/ingested_ts, and returns how many. serve answers as of ts from the online store only: the same point-in-time rule and TTL as offline, applied to the published records (so unpublished writes are invisible, and a feature with nothing published is the default). It returns one value per requested feature, in order.
6. Part 1 behaviour is unchanged. A late record with an *older* event time never displaces a newer value.

Example

[["register", 1, "avg_basket", 100, 0],
 ["ingest", 20, "u1", "avg_basket", 10, 40],
 ["materialize", 30, "train", [["u1", 15], ["u1", 25]], ["avg_basket"]],
 ["backfill", 40, "u1", "avg_basket", [[12, 7]]],
 ["materialize", 41, "train", [["u1", 15], ["u1", 25]], ["avg_basket"]],
 ["read_snapshot", 42, "train", 1],
 ["serve", 43, "u1", ["avg_basket"]],
 ["publish", 44],
 ["serve", 45, "u1", ["avg_basket"]]]

returns

[true, true,
 {"snapshot": "train", "version": 1, "rows": 2, "checksum": 40}, 1,
 {"snapshot": "train", "version": 2, "rows": 2, "checksum": 14},
 [["u1", 15, 0], ["u1", 25, 40]], [0], 2, [7]]

Part 3: Rolling features and a skew check (ML-1927)

Goal: the fraud team wants derived features, rolling sums over a raw feature, available in training sets, snapshots and online serving with exactly the same point-in-time semantics as raw ones. The platform team wants a check that finds training/serving skew. This changes the Part 2 online store: serving a rolling sum needs the underlying records, not just a latest value.

Rules

1. A rolling feature's value for an entity as of T is the sum of value over the source feature's records with T - window < event_ts <= T that are visible at T (rule above), where for each event_ts only one record counts: the latest-arriving visible one. If no record counts, the value is the feature's default (a counted record whose value is 0 makes the sum 0, not the default). TTL does not apply to rolling features; the window does that job.
2. Rolling features work anywhere a feature name is accepted (training_set, materialize, serve, skew_check) with the same visibility and tie rules as raw features, including backfills and late ingests. Online, they are computed from the published records only. A rolling feature cannot be ingested or backfilled (not a raw feature, in the position listed for those operations after unknown feature), and source must be a raw feature.
3. skew_check compares, for every distinct (entity, feature) pair in entities x features, what serve returns at ts with what an offline training_set row for the label (entity, ts) would contain. It returns {"checked": <distinct pairs>, "mismatches": [[entity, feature, online, offline], ...]} listing only differing pairs, sorted by entity then feature name (plain string order). Both lists must be non-empty lists of names.
4. Earlier Parts' behaviour must not change.
5. Hidden tests build tens of thousands of events per entity with backfill_series and materialize thousands of labels at once (materialize over a label grid with raw and rolling features). The whole suite has a few seconds.

Example

[["register", 1, "clicks", 1000, 0],
 ["register_rolling", 2, "clicks_10", "clicks", 10, -1],
 ["ingest", 11, "u1", "clicks", 10, 5],
 ["ingest", 16, "u1", "clicks", 12, 3],
 ["ingest", 21, "u1", "clicks", 12, 4],
 ["ingest", 22, "u1", "clicks", 20, 1],
 ["publish", 23],
 ["ingest", 24, "u1", "clicks", 23, 9],
 ["skew_check", 25, ["u1", "u2"], ["clicks_10", "clicks"]],
 ["serve", 26, "u1", ["clicks_10"]],
 ["training_set", 40, [["u1", 15], ["u1", 19], ["u1", 21], ["u1", 22], ["u1", 32]], ["clicks_10", "clicks"]]]

returns

[true, true, true, true, true, true, 4, true,
 {"checked": 4, "mismatches": [["u1", "clicks", 1, 9], ["u1", "clicks_10", 1, 10]]},
 [1],
 [["u1", 15, 5, 5], ["u1", 19, 8, 3], ["u1", 21, 4, 4], ["u1", 22, 1, 1], ["u1", 32, 9, 9]]]

At label 22 the window is (12, 22], so event 12 is out; at label 21 the correction (arrived at 21) replaces the original 3 rather than adding to it.

Notes

• pit.py is the single source of truth for what a lookup sees; callers should not re-implement it.
• You may use the AI assistant freely; you are expected to verify what it produces.
• Later parts must not break earlier ones; hidden tests re-check earlier behaviour.