Skip to content

Carry UK build datasets as populace Frames: retire the UKNationalDataset / UKRowwiseDatasetResult shadow types #612

Description

@juaristi22

Where this sits

DESIGN.md makes populace.frame.Frame the atom of the stack: entity tables with explicit linkage where "structure is established once, at assembly, and every operator works on the bundle. No operator ever re-derives person↔unit attachment from a flat frame", carrying typed weights "with conservation invariants the kernel enforces".

The US path honours that. The UK path carries a shadow kernel type instead — UKNationalDataset (uk_runtime/national_build.py:81): three pandas DataFrames plus household_weight_kind and mass_log, with with_tables() standing in for Frame's functional update and UKNationalStage.run() (:161) standing in for the Frame -> Frame stage contract. It already imports WeightKind and MassChangeRecord from populace-frame. It is a Frame with the invariants, strata, metadata receipts, and links removed. #495's rowwise work added a second one, UKRowwiseDatasetResult (rowwise_dataset.py:62).

No issue or PR covers this. The nearest homes assume it rather than charter it: #378's identity spine already lists a "canonical post-source frame hash" and sets migration order US → UK; #144 owns source-stage and target ownership; #495 is explicitly "mostly wiring, fences, and targets".

This is not a tidiness issue. The divergence is compounding: between #156 (2026-06-20, still an unmerged draft) and #553 (2026-07-25) the UK path gained ~5,550 additions on the shadow carrier and went from one parallel type to two.

The cost, measured

1. Every helper that touches dataset structure gets written twice. The shared Frame-typed input_mass_totals is 78 LOC (build/input_mass.py:28). The UK equivalent, weighted_integrity.py, is 626 LOC — and its own docstring calls uk_dataset_input_mass_totals (:415) "the UK analog of populace.build.input_mass.input_mass_totals". Because it cannot accept a Frame it takes dataset: Any and hand-rolls _entity_table (:329) and _uk_entity_weights (:342), re-deriving the household→person/benunit weight broadcast and re-validating structural columns the kernel guarantees at construction. The #609 changelog records the split precisely: the gate comparison was promoted and shared; the evidence plumbing had to be rewritten.

The tax recurs across uk_runtime: 15 dataset: Any parameters, 11 hand-rolled entity-extraction helpers, 15 files re-deriving the household_weight broadcast.

2. UK is locked out of country-agnostic machinery built to serve it.

module LOC UK imports us_runtime imports
outer_stage_runtime.py 907 no yes
frame_checkpoint.py 858 no yes
plan.py 244 no yes
input_mass.py 78 gate only, not the extractor yes

outer_stage_runtime.py's docstring: "The runtime deliberately knows nothing about a country rules-engine H5. Every stage boundary is a lossless populace.frame.Frame checkpoint." It was written country-agnostic. UK is excluded by one type. In its place UK stage identity is _UKSourceFileFingerprint (national_build.py:60) — device, inode, size, mtime — which cannot survive a file copy, where frame_identity() hashes content.

3. The UK local solve reaches into private calibration internals. local_solver.py:18-30 imports _optimize, _search_l0_lambda_for_budget, and _torch_constraint_matrix from populace.calibrate.solve, bypassing the public calibrate(targets, frame) -> Frame (calibrate/solve.py:1306) because it has no Frame to pass. Consequences, all live today:

4. One live defect, independent of this migration. rowwise_dataset.py:723:

weight_kind = getattr(dataset, "household_weight_kind", WeightKind.DESIGN)

Any carrier that does not expose that attribute is silently downgraded to DESIGN, written to a new H5 with the wrong populace_household_weight_kind attr, and never errors. Fixed in increment 1 step 3, but it is a hazard now.

Why the lift is bounded

The seam count, not the line count, sizes this:

  • UKNationalDataset( is constructed in exactly 2 places, both in national_build.py (:127 inside with_tables, :227 in the H5 loader). with_tables() has exactly 2 call sites (frs_hmrc_leaves.py:295, hmrc_restoration.py:352).
  • The production national pipeline is 2 stagesfrs_hmrc_retained_leaves and hmrc_spi_income (tools/build_uk_national_dataset.py:339,343). Every other UKNationalStage( in the tree is a test.
  • The linkage convention already matches Frame's exactly. UK carries person_id, person_household_id, person_benunit_id, benunit_id, household_id; Frame's convention is person_{group}_id + {group}_id. No renaming. EntitySchema(group_entities=("benunit","household")) expresses UK directly.
  • household_weight as a table column is permitted, not rejected — Frame reserves {entity}_weight only on entities lacking typed weights (bundle.py:303-332), and permits it on one that has them precisely "for round-trip convenience". No column surgery for the H5 round-trip.
  • The published artifact contract does not change. populace.data "does not depend on the Frame kernel: a published population is an engine-native dataset". The Frame lives inside the build; the export boundary stays engine-native.
  • Leaf helpers that take and return pd.DataFrame (QRF imputation, geography merges, sparse assembly) do not change. Frame is per-entity pandas tables internally; the doctrine binds stages, not helpers.

Read sites to convert: ~56 dataset.person/.benunit/.household (43 in populace-build/src, 13 in tools/) plus ~27 under other receiver names (result., candidate., staging.) — call it ~80, plus ~35 in tests across 9 UK test files.

End state

One carrier for UK builds. Stages typed Frame -> Frame. Provenance travelling beside the Frame as its own typed record. The duplicated evidence plumbing deleted, and the UK path calling populace's own front doors — input_mass_totals, calibrate, frame_checkpoint, frame_identity. This is also #378's stated UK prerequisite, delivered as a side effect rather than a separate epic.

Increment 1 has little standalone value: it is the key that unlocks 2 and 3, not a win by itself.

Increment 0 — Linkage preflight (credentialed, no code change)

PR CI is secrets-free and never touches restricted microdata, and the UK H5 is private UKDS-licensed data. So prove Frame construction succeeds before migrating anything. tools/preflight_uk_frame_linkage.py, following tools/preflight_us_release_gates.py for exit codes and tools/measure_uk_weighted_integrity_baselines.py for disclosure control.

Assert each invariant Frame.__init__ enforces (bundle.py:151-276):

  • household_id / benunit_id non-null, unique, sorted ascending — Frame raises, it does not reorder.
  • Person membership set exactly equals the group id set, in both directions. This is the top unknown: Frame errors on orphaned group ids (a benunit or household referenced by no person) as well as dangling references, and nothing in the UK path checks the orphaned direction todayvalidate_uk_national_dataset (national_build.py:285), spi_support (:915), and rowwise_dataset (:673) all check dangling only.
  • No cross-entity column-name collision (_validate_global_columns, bundle.py:261). Unproven: the SPI stage's metadata columns are entity-prefixed and safe, but clone_index, region, source_year, source_household_id, household_is_capital_gains_clone are unprefixed household columns.
  • household_weight finite, non-negative, any(> 0) — individual zeros are fine, so the 200k zero-weight SPI stratum passes — and its dtype. Weights forces float64 (weights.py:89), so a float32 column today would widen on materialization and change the payload.
  • Dry Frame construction; column-vs-typed-vector agreement.
  • Payload probe per table: row count, column count, {dtype: count} histogram, index type, is_range_index.

Run twice: on the certified candidate, and on the last successful staging output. The second matters more — post-SPI tables are what Frame will validate.

The household side is already proven in production. Every successful HMRC or CGT calibration run constructs a Frame over the full person table plus household (hmrc_calibration.py:286, cgt_calibration.py:178), so household_id unique + sorted + exactly-referenced already holds. benunit is the only genuine unknown — and #609's fdbcbb0 added a benunit-nesting validator that raises when a benunit has no member persons, which is functionally Frame's orphaned-benunit check. Running the #609 battery on a real artifact is free evidence for the single most likely failure.

Disclosure discipline. Frame's error messages embed real ids — bundle.py:196 emits duplicated.head(5).tolist(), :211/:216 emit orphaned[:5] / dangling[:5]. Never print the exception text. Catch ValueError, classify to an invariant name, print the classification plus a threshold-guarded count. A violation count is itself a small-cell frequency, so apply the --sdc-minimum-count rule (default 10) and report below_threshold rather than "3 orphan benunits". The load-bearing output is the boolean.

Increment 1 — Carrier swap

Acceptance is payload identity, not byte identity

Byte identity is unattainable and always was: HDFStore.put defaults track_times=True (pandas/io/pytables.py:1145, forwarded at :4591), so HDF5 stamps object-header ctime/mtime and two consecutive runs of today's unchanged writer already differ. Nothing depends on staging bytes either — the fences bind the input candidate (hmrc_restoration.py:254, :589), #495's clone reads at key/table/dtype level (rowwise_dataset.py:741), and the staging sha is only recorded, retention: "local_untracked" (tools/build_uk_national_dataset.py:450).

Payload identity = same store keys in the same write order, same per-table column list in order, same per-column dtype, same index type and values, same row order, same values, same HDF5 root attributes in the same set order. Assertable in PR CI on synthetic tables, which byte identity never could be.

Steps

1. Add the engine-tables materializer, standalone, zero callers. Frame -> dict[str, pd.DataFrame], copying tables and overwriting household_weight from the typed Weights. Prefer promotion over a third copy: adapters/axiom.py:470 is already country-agnostic (iterates bundle.weighted_entities) where adapters/policyengine_us.py:797 hardcodes household. Lift the generic body into a shared helper all three call. Both existing copies carry the invariant worth keeping verbatim — typed weights are authoritative, any existing {entity}_weight column is overwritten, never trusted.

2. Add the payload-identity test, no production change. Synthetic UK tables written both ways; assert equal keys-as-ordered-list, per-table column order, per-column dtype, index, row order, values, and root attribute names/order/values. This is the gate that makes step 4 safe.

3. Harden the three duck-typed consumers, while the carrier is still the dataclass. terminal_gates._entity_tables (:814), weighted_integrity._entity_table (:329), rowwise_dataset._dataset_tables (:708) all resolve tables via getattr(dataset, entity, None). Two would raise a misleading TypeError; the third silently downgrades WeightKind (above). Fix before the carrier moves so step 4 cannot regress it silently.

4. Retire UKNationalDataset; thread Frame plus a provenance record.

Provenance cannot live on the Frame: Frame.__slots__ (bundle.py:90) blocks attaching _source_h5, and Frame.metadata cannot hold it either — _freeze_metadata_value (bundle.py:1356) accepts only mappings/sequences/sets plus scalars, so a Path raises, and metadata is public and caller-suppliable, which would turn an unforgeable init=False field into a forgeable one.

But provenance already has its own class. UKCertifiedCandidateIdentity (hmrc_restoration.py:93) is a frozen record carrying path, filename, tier, revision, sha256, size_bytes, _source_file_fingerprint, and an unforgeable _verification_token sentinel (object() at :90, compared with is not at :558). "Provenance as its own typed value beside the data" is therefore the established idiom, and UKNationalDataset._source_h5 / _source_file_fingerprint partly duplicate it. So: a small UKStagingProvenance following that shape — or a fold into it, if inspection shows they are the same concept.

Do not subclass Frame. All four Frame operations hard-construct Frame(...) (bundle.py:742, :921, :1164, :1241) rather than type(self), so a subclass silently degrades to a base Frame and drops provenance on any operation.

Rejected: an envelope with delegating .person/.benunit/.household properties. It buys the single-weights-object and validated-linkage goals while leaving the shadow type in place permanently, and a wrapper around a Frame is worse than a shadow type beside one — two things claim to be the source of truth, which is what this issue exists to end. With 2 production stages the direct conversion is not a big-bang risk; sequence it per helper cluster so the suite stays green between commits.

5. Route both writers through the materializer. write_uk_national_dataset (:327) keeps its own temp-path + atomic .replace() + _write_weight_metadata two-phase wrapper: UKSingleYearDataset knows nothing about the root attrs, and its save() does not unlink (policyengine_uk/data/dataset_schema.py:122, unlike the multi-year variant at :228). write_uk_rowwise_dataset (rowwise_dataset.py:634) is a verbatim second copy of the same four store.put calls — route it in the same commit or ship two writers that drift.

Three concrete traps:

  • Pass time_period as the string. policyengine_us.py:830 does int(period) and hmrc_calibration.py:424 already does int(dataset.time_period), which raises on "2023-24".
  • Keep household_weight in the household table. Assigning to an existing column preserves its position; a dropped-then-materialized column reappends last and changes the payload.
  • Never load via UKSingleYearDataset(file_path=...). It does pd.HDFStore(file_path) at dataset_schema.py:54 with default mode "a" — a write open, which would bump st_mtime_ns/st_ctime_ns and break _uk_source_file_fingerprint (national_build.py:69) and the certified-candidate fence on licensed data. Populace's own loader correctly uses mode="r" (:217); keep it that way.

6. One test must change, and must not be "fixed" by sorting. tests/test_uk_frs_hmrc_leaves.py:122-129 deliberately shuffles household, person, and benunit (.sample(frac=1.0, random_state=...)) to prove lineage resolution is id-keyed, not positional. Frame requires group ids sorted, so it will raise. Correct fix: sort household and benunit, keep person shuffled. Frame places no ordering constraint on the person table (_validate_tables only requires person_id unique, bundle.py:151) and _resolve_candidate_lineage reindexes by id (frs_hmrc_leaves.py:488), so person row i still does not correspond to household row i and a positional implementation is still caught. The test keeps its teeth.

Increment 2 — Retire the remaining duplication

The gate-evidence half belongs to #611, not here. Deleting uk_dataset_input_mass_totals (:415), _entity_table, and _uk_entity_weights in favour of the shared input_mass_totals(frame) — ~400 LOC out of weighted_integrity.py — is exactly what #611's increment 1 does to the UK battery, only correctly typed because the carrier is a Frame by then. Do not charter it twice.

What stays here:

  • Replace the three private populace.calibrate.solve imports (local_solver.py:18-30) with the public calibrate() front door, so UK weights take the WeightKind.CALIBRATED transition through the kernel. Adjudicate max_weight_ratio against max_weight_ratio anchors differently per arm: 5x vs design (dense) but ~25x effective (sparse refit re-anchors) — declare one stretch contract #493's "one declared stretch contract" in the same PR — the 100.0 default exists only because this is a second entry point.
  • Retire the remaining dataset: Any signatures outside the gate battery.
  • Drop the household_weight column from the carrier only after its readers move (diagnostics.py:422, release_input_coverage.py:691, weighted_integrity.py:393, spi_support.py:274/705/738/911, local_runner.py:294, rowwise_geography.py:331, cgt_calibration.py:145/186, hmrc_calibration.py:228/294/392). Not in increment 1 — see the column-order trap.

Increment 3 — Adopt the stage runtime

Swap _UKSourceFileFingerprint staging for outer_stage_runtime + frame_checkpoint, giving UK content-addressed stage identity and restartability. Staging becomes a populace-private frame checkpoint between stages while the engine-readable artifact stays materializer-produced — the US shape (l0_refit_export.py:599write_datasetUSSingleYearDataset.save()).

Two honest caveats: the staging format changes, so #495's clone input and national_staging_build_record.json need a one-time re-cut; and StageRuntime/StagePipeline has exactly one consumer today (tools/build_us_puf_support_base.py), so UK would be the second — both the payoff and the risk. #378 frames that as the point: UK "proves the abstraction on the second consumer".

Also gated here: do not let Frame.select or Frame.concat into the UK stage path in increment 1. select keeps the person table's original index (bundle.py:1220) while resetting every group index (:1234), so the two disagree; concat resets (:1159). Today's UK stages produce a RangeIndex via pd.concat(..., ignore_index=True) (spi_support.py:312), and format="table" persists the index, so either operator changes the payload.

Explicitly out of scope

  • Frame strata adoption. UK provenance is multidimensional and household-grain — SPI_REPLACEMENT_STRATA_COLUMNS = ("clone_index", "household_is_capital_gains_clone", "region") (spi_support.py:35), plus household_is_spi_synthetic, consumed by uk_zero_weight_strata (diagnostics.py:167) as a multi-column groupby with weights passed positionally. Frame's strata is a single person-grain label Series with stratum_mass. Mapping one onto the other is a design decision, not a rename, and needs its own adjudication. These stay ordinary household columns. When strata is adopted it must be built off person.index_validated_strata requires .index.equals(person.index) (bundle.py:348).
  • UKRowwiseDatasetResult. National staging is strictly upstream of the rowwise clone in UK dense/local epic: rowwise constituency-grain dataset by cloning the national staging pool (US dense-arm analog) #495's pipeline, so migrating national first makes the clone's input a Frame. Its own carrier swap is a follow-on — but increment 1 steps 3 and 5 still touch rowwise_dataset.py to close the silent-downgrade and duplicate-writer holes.
  • PR Add UK raw-source and local-geography parity #156. Left open. Its uk/source_stages.json (1,007 lines) is its only unlanded piece and is independent of this work: us/ and be/ each carry a full source_stages.json, uk/ has only hmrc_income_source_stages.json. Worth harvesting separately.

Verification

  1. uv run pytest and uv run ruff check . green. Each increment-1 step green on its own.
  2. Increment 1 acceptance: the payload-identity test (step 2), in PR CI on synthetic tables, plus one credentialed rebuild whose payload probe matches the increment-0 receipt.
  3. Weight-metadata round-trip: read back household_weight_kind and mass_log from the H5 root attributes; assert they match the Frame's typed weights and mass log.
  4. Re-run the UK parity increment 4: port the weighted integrity gates (input-mass parity + QRF tail concentration) into the UK terminal battery (#578) #609 terminal-gate battery and diff evidence digests — identical verdicts through the shared helper as through the deleted UK analog. Any delta is a real behavioural difference to explain, not to paper over.
  5. changelog.d/ fragment per increment.
  6. If increment 0 shows the national tables need sorting, that sort lands with an explicit test that weights still map to the right households. local_runner.py:294 reads household_weight positionally and local_geography.stacked_weights_to_long (:290) builds np.tile(np.arange(n_households), n_areas). Sorting inside the carrier is exactly the bug class sort_households_by_id exists to prevent — its docstring names it "the Populace-side guard for the 2024-25 FRS bug class where household IDs were sorted but household attributes were assigned by raw row position" (local_geography.py:72). Let Frame raise; fix the producer.

Suggested slot in the #578 UK parity ladder: item 2.5, after diagnostics standardization and before the consolidated terminal gates, so items 3–6 are written once against shared machinery.

Risks

risk severity handling
An orphaned benunit (no member persons) → Frame raises high Increment 0; #609's fdbcbb0 validator is a free rehearsal. If real, it is a data defect worth its own issue
household_weight float32 today → dtype widens, payload changes medium Reported by increment 0; if float32, materialize preserving source dtype
Sorting for Frame reorders rows under a positional weight vector medium Verification 6; let Frame raise, fix the producer, never sort in the carrier
Silent WeightKind downgrade at rowwise_dataset.py:723 medium Increment 1 step 3, before the carrier moves
Merge conflict with #610 high Do not start until #610 merges
StageRuntime has one consumer; UK is its first portability test low Increment 3 only

Unverified

Everything about the real certified candidate: column set, per-column dtypes, household_weight dtype, index type, whether benunit_id is sorted ascending, whether any orphaned benunit or household exists, whether any column name is shared across entities. The H5 is private UKDS-licensed data and is not in the repo. That is what increment 0 exists to settle — increment 1 does not start without the receipt.

Refs: #144 (UK migration umbrella), #495 (rowwise consumer), #378 (Build K step-1 prerequisite), #609 (the gate-duplication receipt), #611 (complementary gate battery), #493 (max_weight_ratio divergence), #578 (parity ladder).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions