diff --git a/changelog.d/full-scale-string-storage-canonicalization.fixed.md b/changelog.d/full-scale-string-storage-canonicalization.fixed.md new file mode 100644 index 00000000..82530e96 --- /dev/null +++ b/changelog.d/full-scale-string-storage-canonicalization.fixed.md @@ -0,0 +1 @@ +Canonicalize pandas string storage at the ACS PUMS parse and ASEC raw-stage checkpoint load boundaries, and render dtype reprs in the spine-assembly refusal: pyarrow-carrying environments resolved fresh string casts to pyarrow storage, which failed shared-column dtype validation against checkpoint-restored python-storage channels with both sides printing as 'str'. diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/acs_pums.py b/packages/microcosm-build/src/microcosm/build/us_runtime/acs_pums.py index 81ebd047..dd0456fc 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/acs_pums.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/acs_pums.py @@ -32,6 +32,7 @@ import numpy as np import pandas as pd +from microcosm.build.serialization_dtypes import canonicalize_frame_string_dtypes from microcosm.frame import US_SCHEMA, Frame, WeightKind, Weights from microcosm.frame.units import assign_us_unit_structure @@ -309,6 +310,14 @@ def build_acs_pums_unit_frame( ), } ) + # This parse is a serialization boundary: freshly cast string columns + # otherwise carry the environment-resolved pandas string storage, which + # diverges from checkpoint-restored channels when pyarrow is installed. + frame = canonicalize_frame_string_dtypes( + frame, + boundary="ACS PUMS source parse", + in_place=True, + ) return frame, metadata diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/asec_checkpoint.py b/packages/microcosm-build/src/microcosm/build/us_runtime/asec_checkpoint.py index 9e7b21f8..0c6ebe0f 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/asec_checkpoint.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/asec_checkpoint.py @@ -23,6 +23,7 @@ FrameIdentity, frame_identity, ) +from microcosm.build.serialization_dtypes import canonicalize_frame_string_dtypes from microcosm.build.us_runtime.operator_boundary import ( assert_operator_free_source_frame, ) @@ -187,7 +188,16 @@ def load_asec_raw_stage_checkpoint( ) metadata["identity"] = stored_identity.to_payload() metadata["source_construction_identity"] = source_construction_identity.to_payload() - return loaded.frame, metadata + # Canonicalize only after every identity and binding check has passed on + # the restored representation: this load is a serialization boundary, and + # downstream spine assembly requires one physical string dtype per shared + # column regardless of the storage the checkpoint was written under. + frame = canonicalize_frame_string_dtypes( + loaded.frame, + boundary="ASEC raw-stage checkpoint load", + in_place=True, + ) + return frame, metadata def _validate_outer_stage_binding( diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/spine_assembly.py b/packages/microcosm-build/src/microcosm/build/us_runtime/spine_assembly.py index 21a4c24d..ce6b6e29 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/spine_assembly.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/spine_assembly.py @@ -296,8 +296,10 @@ def _validate_shared_column_dtypes(frames: Mapping[str, Frame]) -> None: if not all(dtype == values[0][1] for _, dtype in values[1:]) } if mismatches: + # repr, not str: pandas string dtypes with different storages all + # str() as 'str', which renders the mismatch invisible. details = ", ".join( - f"{column}={[(channel, str(dtype)) for channel, dtype in values]}" + f"{column}={[(channel, repr(dtype)) for channel, dtype in values]}" for column, values in sorted(mismatches.items()) ) raise ValueError( diff --git a/packages/microcosm-build/tests/test_us_acs_pums.py b/packages/microcosm-build/tests/test_us_acs_pums.py index 0c7fd22d..48b656d9 100644 --- a/packages/microcosm-build/tests/test_us_acs_pums.py +++ b/packages/microcosm-build/tests/test_us_acs_pums.py @@ -7,6 +7,7 @@ import pandas as pd import pytest +from microcosm.build.serialization_dtypes import CANONICAL_STRING_DTYPE from microcosm.build.us_runtime.acs_pums import ( ACS_2024_1YR_SPINE, AcsPumsSource, @@ -138,7 +139,12 @@ def _asec_shaped_frame() -> Frame: "person_marital_unit_id": np.asarray([601], dtype=np.int64), "source_year": np.asarray([2024], dtype=np.int64), "source_household_id": np.asarray([7], dtype=np.int64), - "source_person_id": pd.Series(["7-1"]).astype(str), + # Pinned to the canonical storage, exactly as checkpoint restore + # delivers the real ASEC channel. A bare astype(str) would follow + # the environment's storage default and mask the cross-channel + # divergence this fixture exists to model (pyarrow environments + # resolve fresh 'str' casts to pyarrow storage). + "source_person_id": pd.Series(["7-1"]).astype(CANONICAL_STRING_DTYPE), "source_row_id": np.asarray([0], dtype=np.int64), "A_AGE": np.asarray([55], dtype=np.int64), } @@ -287,6 +293,31 @@ def test_built_acs_lineage_assembles_with_asec_without_measured_coercion( ) +def test_build_acs_pums_unit_frame_canonicalizes_string_storage( + tmp_path: Path, +) -> None: + """The parse boundary must not leak environment-resolved string storage. + + With pyarrow installed, pandas resolves fresh ``astype(str)`` casts to + pyarrow-backed storage while checkpoint-restored channels carry the + canonical python-backed dtype — the exact split that failed the first + full-scale spine assembly (both sides printed as 'str'). + """ + + pytest.importorskip("microunit") + frame, _metadata = build_acs_pums_unit_frame(_source(tmp_path), chunksize=1) + + for entity in frame.entities: + table = frame.table(entity) + offending = { + column: dtype + for column, dtype in table.dtypes.items() + if isinstance(dtype, pd.StringDtype) and dtype != CANONICAL_STRING_DTYPE + } + assert not offending, f"{entity} carries non-canonical string storage" + assert frame.table("person")["source_person_id"].dtype == CANONICAL_STRING_DTYPE + + def test_load_acs_pums_tables_rejects_duplicate_household_serialno( tmp_path: Path, ) -> None: diff --git a/packages/microcosm-build/tests/test_us_asec_checkpoint.py b/packages/microcosm-build/tests/test_us_asec_checkpoint.py index d0ce5d6e..4a604a3f 100644 --- a/packages/microcosm-build/tests/test_us_asec_checkpoint.py +++ b/packages/microcosm-build/tests/test_us_asec_checkpoint.py @@ -12,6 +12,7 @@ OUTER_STAGE_CONTEXT_SCHEMA_VERSION, frame_identity, ) +from microcosm.build.serialization_dtypes import CANONICAL_STRING_DTYPE from microcosm.build.us_runtime import ( ASEC_RAW_STAGE_ARTIFACT_KIND, ASEC_RAW_STAGE_OPERATOR_STATUS, @@ -236,6 +237,13 @@ def test_loads_operator_untouched_raw_stage_checkpoint(tmp_path: Path) -> None: assert loaded_metadata["artifact_kind"] == ASEC_RAW_STAGE_ARTIFACT_KIND assert loaded_metadata["stage"] == ASEC_RAW_STAGE_STAGE assert loaded_metadata["operator_status"] == ASEC_RAW_STAGE_OPERATOR_STATUS + # The load is a declared string-storage canonicalization boundary: no + # entity may leak a non-canonical pandas string dtype downstream, + # whatever storage the checkpoint was written under. + for entity in frame.entities: + for column, dtype in frame.table(entity).dtypes.items(): + if isinstance(dtype, pd.StringDtype): + assert dtype == CANONICAL_STRING_DTYPE, (entity, column) @pytest.mark.parametrize( diff --git a/packages/microcosm-build/tests/test_us_spine_assembly.py b/packages/microcosm-build/tests/test_us_spine_assembly.py index 1763e27f..947e91af 100644 --- a/packages/microcosm-build/tests/test_us_spine_assembly.py +++ b/packages/microcosm-build/tests/test_us_spine_assembly.py @@ -307,6 +307,52 @@ def test_assemble_spines__rejects_shared_dtype_mismatch() -> None: ) +def test_assemble_spines__string_storage_mismatch_is_named_in_the_error() -> None: + """Divergent pandas string storages must be visible in the refusal. + + Both python- and pyarrow-backed 'str' dtypes print identically under + str(); the first full-scale run failed with an error whose two sides both + read 'str'. The refusal must render repr so the storages are legible. + """ + + pytest.importorskip("pyarrow") + + def _with_probe(frame: Frame, storage: str, value: str) -> Frame: + tables = {entity: frame.table(entity).copy() for entity in frame.entities} + person = tables["person"] + person["stacked_probe_label"] = pd.Series( + [value] * len(person), + index=person.index, + ).astype(pd.StringDtype(storage=storage, na_value=np.nan)) + return Frame( + tables, + US_SCHEMA, + {"household": frame.weights_for("household")}, + frame.strata, + ) + + with pytest.raises( + ValueError, match="identical dtypes.*stacked_probe_label" + ) as excinfo: + assemble_spines( + { + "asec": _with_probe(_asec_frame(), "python", "x"), + "acs": _with_probe(_acs_frame(), "pyarrow", "y"), + }, + household_mass_shares={"asec": 0.5, "acs": 0.5}, + ) + message = str(excinfo.value) + # The python-backed side names its storage explicitly; pandas omits the + # storage kwarg from the repr of the environment-default (pyarrow) side. + # Both reprs must appear and must differ — under str() both sides + # rendered as 'str' and the mismatch was unreadable. + asec_repr = repr(pd.StringDtype(storage="python", na_value=np.nan)) + acs_repr = repr(pd.StringDtype(storage="pyarrow", na_value=np.nan)) + assert asec_repr != acs_repr + assert asec_repr in message + assert acs_repr in message + + def test_assemble_spines__owns_support_provenance() -> None: asec = _asec_frame() tables = {entity: asec.table(entity).copy() for entity in asec.entities}