Skip to content

fix: verify registered models against the stored spec on deserialization - #1877

Open
rahul188 wants to merge 2 commits into
datachain-ai:mainfrom
rahul188:fix-1872-custom-type-shadowing
Open

fix: verify registered models against the stored spec on deserialization#1877
rahul188 wants to merge 2 commits into
datachain-ai:mainfrom
rahul188:fix-1872-custom-type-shadowing

Conversation

@rahul188

Copy link
Copy Markdown

Problem

SignalSchema._deserialize_custom_type resolves custom types by name@version alone:

if fr := ModelStore.get(model_name, target_version):
    return fr

Since ModelStore is process-global and all generated models are @v1, two datasets whose inferred models share a name (both explode()d from a column named json, two parquet files with an items list-of-struct column, …) shadow each other: whichever model registers first is used to read every same-named dataset for the rest of the process. Depending on how the shapes overlap that surfaces as a KeyError (missing DB column), a Pydantic ValidationError, or a silent misread. Repro from #1872:

c1 = dc.read_values(json=[{"a": 1}]).explode("json")
c2 = dc.read_values(json=[{"b": "x"}]).explode("json")
c1.save("e1"); c2.save("e2")
dc.read_dataset("e2").to_values("json_expl")  # works
dc.read_dataset("e1").to_values("json_expl")  # KeyError: 'json_expl__b'

Fix

Following the approach sketched in the issue:

  • A registered model resolves a custom type only when its field spec matches the spec stored in the dataset schema. The check reserializes the registered model with the existing _serialize_custom_model machinery and compares fields per custom type, so a divergence in a nested model (identical top-level field strings, different Inner@v1 shapes) is caught too.
  • On mismatch, the model is recreated from the stored spec under a spec-fingerprinted alias ({name}_{fingerprint[:10]}, the same shape to_partial uses for partials) instead of clobbering the registered model's name in ModelStore. Repeated reads of the same dataset hit the alias cache and reuse the recreated class.
  • When there is no registered model, behavior is unchanged (recreate under the stored name and register).

The fingerprint covers the transitive closure of the type's spec (fields, bases, hidden fields of every reachable custom type), so two colliding specs that differ only in a nested type or base still get distinct aliases.

Testing

  • test_deserialize_same_name_different_shape_not_shadowed — the cross-dataset collision, plus re-read identity (same class object on repeated deserialization of both specs)
  • test_deserialize_nested_shape_divergence_not_shadowed — identical top-level specs, diverging nested model
  • test_deserialize_prefers_registered_model_when_spec_matches — round-trip still returns the registered class itself, not a recreation
  • test_explode_same_column_name_across_datasets — the issue's end-to-end repro (explodesaveread_dataset both ways)

All four were run against unmodified main first: the three collision tests fail there (the e2e test with KeyError: 'json_expl__a'), and the registered-model test passes on both, confirming no regression of round-trip identity. Full tests/unit/lib/test_signal_schema.py (120) and tests/unit/lib/test_datachain.py (349) pass; ruff check, ruff format --check, and mypy are clean on the changed files.

Fixes #1872

SignalSchema._deserialize_custom_type returned any registered model whose
name and version matched, so two generated models sharing a name (e.g.
both created by explode() on a column named json) shadowed each other
across datasets: whichever registered first was used to read every
same-named dataset, producing KeyError, ValidationError, or silently
misread values.

A registered model now resolves a custom type only when its field spec,
including nested custom types, equals the spec stored in the dataset
schema. On mismatch the model is recreated from the stored spec under a
spec-fingerprinted alias, so the registered model keeps its name and
repeated reads reuse the recreated class.

Fixes datachain-ai#1872
@rahul188

Copy link
Copy Markdown
Author

Hi 👋 Gentle nudge on this PR — it's mergeable and CI is passing, just waiting on a review. Would really appreciate a look whenever you have bandwidth, and I'm glad to make any changes you'd like. Thanks for your time and for maintaining this project! 🙏

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Custom model name collisions: same-named generated models shadow each other across datasets

1 participant