fix: refuse a dict whose keys collide as JSON instead of losing a value

[shcheklein]

Refs #1914 — two shapes are left, both pinned by strict xfail; see Known gaps.

Warning

Draft — not ready to merge. The write-time check works, but it leans on
markers and backend-specific assertions that should not land as they are.

TODO before this is mergeable

Three xfails and five SQLite-only tests. Each is a real gap, not a formality:

• Two xfails: a live Pydantic model still loses colliding keys.
  tuple[Model, ...] and dict[str, list[Model]] reach the converter as
  model instances, and model_dump(mode="json") merges the keys inside
  Pydantic before anything here can look. Closing this needs models converted
  in Pydantic's python mode everywhere, which preserves both keys — and that
  needs this project's encoder to first cover the types only JSON mode knows
  (timedelta, paths, addresses, plain enums, sets), or those regress.
• One xfail: a numpy object array collapses keys before any JSON exists.
  _numpy_to_python rebuilds the mapping while materializing the array, so
  np.datetime64("NaT", "ns") and None land on one key. This is our own
  code and the fix belongs where the merge happens.
• Five skip_if_not_sqlite tests. They assert what SQLite's adapter
  writes, which has no ClickHouse equivalent — there a JSON item resolves to
  str, so dict items never take that path. They should assert a relationship
  that holds on both backends rather than a physical form. The pattern works:
  one test already compares two shapes' stored values instead of naming either.

Until those are addressed this documents a defect rather than fixing it, and the
markers would become permanent furniture in the test suite.

The problem

JSON names are strings, so distinct Python keys can collapse onto one when a dict is serialized. 1 and "1" both become "1", the last one wins, and the row is written with a value missing — silently, with nothing on the read side able to recover it.

class M(dc.DataModel):
    d: dict[str | int, str]

dc.read_values(m=[M(d={"1": "foo", 1: "bar"})]).save("m")
# read back: {"1": "bar"}   <- "foo" is gone

The shape of the fix

Every earlier revision of this PR failed the same way: it tried to predict what a serializer would write. Each prediction drifted from the real thing in both directions — refusing valid data and missing real losses. What works is reading what the writer actually produced.

_to_jsonable checks as it builds. It constructs the dict itself, JSON-encoding non-str keys so key_needs_json_decode can read them back. A collision there merges the values before anything is serialized, so there is no later output to inspect.

Everything stored unconverted is checked from the emitted JSON. Those values are spelled by the driver's own encoder — sqlite3's registered adapt_array, which is datachain.json.dumps. It spells a tuple key (1, 2) where _to_jsonable spells it [1,2], and with serialize_numpy it materializes object arrays into mappings the input never contained. Neither is predictable, so the emitted JSON is parsed with a duplicate-preserving object_pairs_hook.

convert_type calls that before storing an array of dicts. Without it the guard would not fire until sqlite3 adapted the value mid-insert, with no column to name:

UdfError: UDF returned an invalid value for output column 'm__d'. Expected
JSON-serializable dict[Union[str, int], str]. Keys '1' and 1 both serialize
to the JSON key '1', so one of their values would be lost.

Verified that a poisoned row in the middle of a 2000-row save leaves nothing written.

80 source lines in one file: an exception, a hook, and one _dump_json shared by both branches.

Known gaps

Both are pre-existing, both stay on #1914, and both are pinned by strict xfail so they convert to failures when fixed.

A live pydantic instance. model_dump(mode="json") merges colliding keys inside pydantic-core, before anything here can see them. This is reachable for tuple[Model, ...] and dict[str, list[Model]], where the value arrives as a model rather than flattened. (read_values flattens Model and list[Model] to plain dicts first, so those are caught, by _to_jsonable's own check.)

I attempted this and reverted it. model_dump_json() keeps both keys and emits a duplicate property, so reading that JSON detects the collision — but it is a second serialization pass, and that re-runs field serializers:

@field_serializer("items")
def materialize(self, value):
    return list(value)

# given iter([1, 2, 3]) --  main: 1 call, stores [1, 2, 3]
#                          that revision: 2 calls, stores []

A stateful serializer could likewise emit safe keys during the check and colliding ones during storage. model_dump_json also writes nan and the infinities as null, which model_dump preserves. Closing this needs a single-pass hook pydantic does not offer, so it is deferred rather than approximated.

A numpy object array. _numpy_to_python rebuilds a mapping while materializing the array, so np.datetime64("NaT", "ns") and None collapse onto one None key there, before any emitted JSON exists to read. Other numpy-created mappings are covered, since the encoder builds them where the emitted JSON can be read.

Why detect the collision rather than reject the annotation

The issue suggested refusing ambiguous key annotations at schema-build time. Detecting the actual collision is better on three counts:

• No false positives. dict[str | int, str] holding {"x": "a", 2: "b"} keeps working. Rejecting the annotation would break it.
• It catches shapes an annotation rule would miss — a str subclass whose __eq__ hides a duplicate has a perfectly ordinary annotation.
• It fires where the data is lost, so the message names the actual keys.

For the record: is_chain_type already refuses dict[str | int, X]. The loss was reachable because model fields never consult it — the gap tracked in #1942.

Cost

Not measurable where it matters. A saved row costs ~240 µs and these differences are single-digit µs; across repeated end-to-end runs main and this branch are indistinguishable. In isolation, convert_type on a 20-element list[dict] goes 0.74 → 7.7 µs. Element types other than dict are untouched.

Validation

Rebased onto bc3e2c83. tests/unit: 3178 passed, 15 skipped, 18 xfailed. tests/func/test_datachain.py + test_udf.py + test_datasets.py: 449 passed, 117 skipped — one flake in test_udf_parallel_worker_failure_exits_peers, which passes on rerun and fails the same way on main. ruff 0.16.4 and mypy 2.3.0 clean, matching the pins in .pre-commit-config.yaml.

Tests go through convert_type or the chain API. Coverage includes the collision (int/str, None/"null", bool/"true", tuple/str, str subclass, str-mixin enum), the four nested-collection shapes, the encoder's key spelling in both directions, the numpy-created mapping that is caught, the flattened model shapes, and a round trip of nan/inf/-inf through save.

Five tests carry skip_if_not_sqlite: they assert what sqlite3's adapter stores, which has no ClickHouse equivalent — there a JSON item resolves to str, so dict items never take that path. The one test that derives its expectation from warehouse.python_type(JSON()) runs on both.

🤖 Generated with Claude Code

[cloudflare-workers-and-pages]

Deploying datachain with    Cloudflare Pages

Latest commit:	16b4317
Status:	✅  Deploy successful!
Preview URL:	https://2d4f832e.datachain-2g6.pages.dev
Branch Preview URL:	https://fix-reject-ambiguous-dict-ke.datachain-2g6.pages.dev

View logs

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[Copilot]

Pull request overview

Adds write-time detection for Python dictionary keys that collide when serialized as JSON, preventing silent data loss.

Changes:

• Detects duplicate JSON property names and raises column-aware errors.
• Validates nested dictionary arrays using backend encoder output.
• Adds collision, round-trip, backend, and known-gap tests.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
src/datachain/data_storage/warehouse.py	Implements collision-safe JSON serialization.
tests/unit/test_warehouse.py	Tests conversion and backend serialization behavior.
tests/unit/lib/test_datachain.py	Tests end-to-end save behavior and known gaps.

Suppressed comments (10)

src/datachain/data_storage/warehouse.py:162

• The SQLite adapter does not call _dump_json; it calls datachain.json.dumps independently. State that the options currently match rather than claiming this helper is shared, otherwise future maintainers may incorrectly assume adapter changes automatically update validation.

        JSON columns store this text. Array items are instead stored as they are
        and serialized later, by sqlite3's registered adapter, which calls this
        same function -- so dumping here shows what that will write.

tests/unit/test_warehouse.py:352

• AGENT.md:109-112 requires tests to express intent through their names and omit test docstrings/comments. Remove this historical explanation; the parameter IDs and test name already identify the covered collision.

    """The keys an object array materializes are created by the encoder, after
    _to_jsonable has run. Counting keys let one of these cancel out: the collision
    removed a key and the array added one, so the totals matched.
    """

tests/unit/test_warehouse.py:361

• AGENT.md:109-112 requires tests to rely on descriptive names rather than test docstrings. The test name and exception assertion already state this contract, so remove the docstring.

    """The check runs the driver's encoder, so it sees the driver's TypeError first.
    It must arrive named rather than as a bare TypeError from the conversion.
    """

tests/unit/test_warehouse.py:374

• AGENT.md:109-112 says tests should use clear names and skip docstrings/comments. This test name plus the SQLite-only marker captures the behavior, so remove the explanatory docstring.

    """sqlite3's registered adapt_array spells a tuple key "(1, 2)", so this does
    not collide there. Scoped to SQLite because that adapter is the subject: on
    ClickHouse a JSON item resolves to str, so dict items never take this path.
    _to_jsonable spells it "[1,2]", which collides with the string key and had
    this valid value rejected.

tests/unit/test_warehouse.py:393

• AGENT.md:109-112 requires tests to omit explanatory docstrings and rely on their names. The test name and parameter IDs already distinguish the driver-collision cases, so remove this docstring.

    """The mirror: these do collide in the driver's spelling. _to_jsonable spells
    them "[1,2]" and '"2020-01-02"' — both distinct from the string key — so the
    JSON-encoded spelling let the value be lost.

tests/unit/test_warehouse.py:404

• AGENT.md:109-112 requires tests to use descriptive names instead of test docstrings/comments. The test name and backend-conditional assertions already encode this invariant, so remove the docstring.

    """Dict items are checked for colliding keys but never re-encoded, so each
    backend keeps the form it already stored. Re-encoding them on SQLite stops
    equality, distinct and merge from matching rows written earlier. ClickHouse
    resolves a JSON column to String and has always stored strings there.
    """

tests/unit/lib/test_datachain.py:1024

• AGENT.md:109-112 requires test intent to be conveyed by the test name rather than a docstring. This name and the NaN/Inf parameter IDs are sufficient, so remove the docstring.

    """These survive a round trip, which reading pydantic's JSON instead of its
    dump would have written as null.
    """

tests/unit/lib/test_datachain.py:1040

• AGENT.md:109-112 requires tests to omit explanatory docstrings/comments. The descriptive test name and SQLite-only marker establish the contract; remove this docstring.

    """Guards SQLite's stored form: encoding the items as JSON strings makes this
    filter stop matching. ClickHouse stores strings already and cannot express the
    comparison at all — a dict literal collides with its {name:Type} substitution
    syntax and the query fails to parse.
    """

tests/unit/lib/test_datachain.py:1065

• AGENT.md:109-112 says tests should rely on clear names and skip comments. This restates the test name and expected exception, so remove it.

    # both keys become the JSON key "1", so one value would vanish with nothing
    # on the read side able to recover it

tests/unit/lib/test_datachain.py:1078

• AGENT.md:109-112 says tests should rely on their names and assertions rather than explanatory comments. The following assertions already show that both keys return as strings, so remove this comment.

    # both entries survive, but the declared int key reads back as a JSON string

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

Suppressed comments (3)

src/datachain/data_storage/warehouse.py:122

• This comment describes the reverted model_dump_json() approach, but the code returns model_dump(mode="json"); the strict xfail in test_datachain.py confirms that this path has already merged colliding keys. Please document the actual limitation rather than claiming the JSON output is read here.

            # Pydantic decides how its own keys are written, and model_dump merges
            # any that collide. The JSON it writes keeps both, so read that: it
            # settles the datetime and Decimal spellings too, which the encoder
            # here does not reproduce.

tests/unit/lib/test_datachain.py:1067

• This single-row case does not cover a collision discovered after earlier rows have been buffered or flushed. The PR description says the 2,000-row poisoned-save case was manually verified, and AGENT.md:121-129 requires such probes to become permanent tests; add that regression and assert that the failed save exposes no partial dataset or rows.

    # both keys become the JSON key "1", so one value would vanish with nothing
    # on the read side able to recover it
    with pytest.raises(UdfError, match="would be lost"):
        chain.save("ambiguous_keys")

tests/unit/test_warehouse.py:316

• This test name already states the invariant, while the docstring records a discarded key-counting implementation. AGENT.md:99-110 says tests should rely on clear names and omit change-narrative docstrings; remove it.

    """serialize_numpy materializes an object array into mappings during encoding,
    so the emitted JSON holds more keys than the input. Nothing collides, so this
    must store: counting keys before and after would have called it a collision.
    """

[shcheklein]

I checked this before writing the tests, and the four extra operations would not catch the regression they are meant to guard.

Within one test run every dataset is written by the same code, so distinct/group_by/subtract/merge compare stored against stored and agree with themselves under either representation. Applying the exact change they are supposed to catch (storing array items as JSON strings):

operation	correct	with JSON-string storage
distinct	1	1
group_by	1	1
subtract	0	0
merge right side	present	present
filter == literal	2	0

Only the filter test moves, because it is the one that compares stored data against a Python literal. Adding the other four would give four tests that pass under the very regression they name — the vacuous-test failure AGENT.md:130-133 warns about.

The manual checks in the description did show those operations breaking, but only across generations: a dataset written before the change compared against one written after. A single test run cannot produce that without a fixture of pre-encoded rows.

That fixture is buildable — insert a row with the old representation via raw SQL, then assert distinct/merge treat it as equal to a normally written one. It would genuinely catch a representation change, at the cost of being SQLite-only and pinning a byte format in a test. I have not added it, and I would rather you decide whether that trade is worth it than have me pick unilaterally. Leaving this thread open for that.