dict signals with mixed str/non-str key types silently lose data on save

[shcheklein]

Spotted by @dreadatour while reviewing #1911 (comment). Unrelated to that PR — it predates it and affects the write path, so no read-side decoding rule can recover from it.

Problem

A dict signal whose key type admits both strings and non-strings loses data when saved. json.dumps coerces every key to a string, so 1 and "1" become the same JSON key:

import json
json.dumps({"1": "foo", 1: "bar"})
# '{"1": "foo", "1": "bar"}'     <- duplicate key emitted
json.loads('{"1": "foo", "1": "bar"}')
# {'1': 'bar'}                   <- "foo" is gone

Verified end to end through the schema:

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

m = M(d={"1": "foo", 1: "bar"})
# flattened for storage: [{'1': 'foo', 1: 'bar'}]
# serialized:            {"1": "foo", "1": "bar"}
# read back:             {'1': 'bar'}

The value is destroyed at serialization time, silently. Nothing raises, and there is no warning.

Why the read side cannot fix it

Both keys arrive as the JSON string "1", and the second has already overwritten the first. This is distinct from the key-decoding question handled in #1911 — there the information still exists in the stored value and only the decode rule was wrong. Here the information is gone before the row is written.

Affected types

Any key annotation that can be both a string and something else, e.g. dict[str | int, X], dict[Any, X], dict[object, X]. Uniform key types are fine: dict[str, X] keys never collide, and dict[int, X] keys are JSON-encoded consistently and decode back correctly.

Suggested direction

Reject the ambiguity at schema-build time rather than losing data at write time — treat a key annotation that admits both str and a non-str arm as unsupported, the same way other unrepresentable types are refused. That turns silent corruption into an error the user can act on. datachain.lib.data_model.key_needs_json_decode already reasons about exactly this "can this key be a plain string?" question and would be the natural place to detect it.

Worth checking whether the same collision affects the ClickHouse path, which may serialize maps differently from SQLite.

Status

#1943 catches the collision at
write time in _to_jsonable, where the overwrite actually happens, and reports it as a
JsonSerializationError naming the column and both offending keys.

Detection rather than the annotation rejection suggested above, for two reasons.
dict[str | int, str] holding {"x": "a", 2: "b"} does not collide and still stores —
rejecting the annotation would break it. And a str subclass with its own __eq__
collides while its annotation is perfectly ordinary, so no annotation rule would catch
it. ({"x": "a", 2: "b"} reads back as {"x": "a", "2": "b"}: both values survive, but
the declared int key returns as a string. That is the separate read-side question in
#1918, not this ticket.)

Keys are compared by the string json actually emits, which is str.__str__(k) rather
than str(k). A str-mixin enum serializes as its value, so str(k) would have
written "Key.A" instead of "a" and falsely rejected {Key.A: 1, "Key.A": 2}, whose
emitted keys genuinely differ.

Answering the ClickHouse question above

ClickHouse was already refusing the nested case. Warehouse.python_type resolves a
column type through the dialect, and ClickHouse maps JSON to String, so an
Array(JSON) item never matched the isinstance fast path and was always converted
element by element. SQLite resolved the same column to dict and short-circuited.

#1943 validates the elements on SQLite too, but deliberately does not adopt
ClickHouse's stored form. Writing JSON strings changes what SQLite compares, and the
schema carries no encoding marker to reconcile datasets written on either side of the
change. Measured on one row written before and one after, which read back identically:

operation	correct	with SQLite storing JSON strings
union().distinct()	1	2
union().group_by()	1	2
subtract()	0	1
merge() right side	'gen_pr'	NULL — no match
filter(C(...) == [{"a": 1}])	1	0

merge is worth singling out: it keeps the left row either way, so a count() stays 1
and hides the failure. Only inspecting a right-side column shows the join found nothing.

Physical convergence between the backends therefore needs a migration and literal
conversion, and is out of scope here. Worth knowing for whoever picks that up: on
ClickHouse the equality filter above cannot be expressed at all today — a dict literal
collides with its {name:Type} substitution syntax and the query fails to parse with
DB::Exception: Syntax error ... Expected substitution name (identifier). Pre-existing,
and unrelated to this ticket, but it means the guard for the regression is SQLite-scoped.

The nested path has to read what the encoder emits

Worth recording, because two successive attempts at the nested fix were wrong, each in
a different way. Elements of an Array(JSON) column reach the driver unconverted,
so sqlite3's registered adapter serializes them — adapt_array, which is
datachain.json.dumps — and what it writes cannot be predicted from the input.

First, it spells non-str keys differently from _to_jsonable:

key	_to_jsonable	what the driver stores
(1, 2)	[1,2]	(1, 2)
date(2020, 1, 2)	"2020-01-02" (quoted)	2020-01-02
datetime(...)	"2020-01-02T03:04:00"	2020-01-02 03:04:00
nan	NaN	nan
int, None, bool	1, null, true	same

Validating with the wrong spelling fails both ways:

class Holder(dc.DataModel):
    rows: list[dict[tuple[int, int] | str, str]]

[{(1, 2): "tuple", "[1,2]": "string"}]   # rejected, though the driver keeps both
[{(1, 2): "tuple", "(1, 2)": "string"}]  # accepted, then reloads with only the second

Second — and this is what broke the next attempt — serialize_numpy materializes a
numpy object array into mappings during encoding, so the emitted JSON can hold keys
the input never had. Comparing key counts before and after therefore fails both ways:

# falsely rejected: 1 key becomes 2, with no collision anywhere
[{"payload": np.array([{"a": 1}], dtype=object)}]

# accepted, then loses "first": the collision removed a key and the array added one,
# so both totals were 3
[{"1": "first", 1: "second", "payload": np.array([{"a": 1}], dtype=object)}]

The check therefore parses the emitted JSON with a duplicate-preserving
object_pairs_hook. That sees the duplicate properties the row would actually carry,
with no spelling to restate and no structural comparison to invalidate.

_to_jsonable keeps a check of its own, and the two are complementary: it builds a
dict, so a collision there merges the values before serialization and leaves nothing
in the emitted JSON to find. Neither alone is sufficient.

Covered by #1943

• a top-level dict[str | int, X] signal
• dict[str, dict[str | int, X]] — via _to_jsonable's recursion, not the array path
• list[dict[str | int, X]], list[list[...]] and tuple[..., ...], validated in
  the driver's key spelling so tuple, date, datetime and nan keys agree
• str subclasses that hide a duplicate behind a custom __eq__
• str-mixin enum keys, which serialize by value and must not be rejected
• a colliding mapping inside a numpy object array, in either an Array(JSON) item
  or a top-level JSON column. _to_jsonable does not enter an ndarray, and the
  encoder materializes the mapping afterwards, so the duplicate property reached the
  row unseen. Caught now because the emitted JSON is what gets checked.

Still open

• a live pydantic instance reaching _to_jsonable. It checks is_pydantic
  first and returns model_dump(mode="json"), which collapses colliding keys before the
  check can run. This is one mechanism, not a read_values/UDF split — whether a
  given route happens to flatten its values into plain dicts first is incidental:

  read_values dict[str, list[Inner]]   SAVED  <- value lost
  read_values tuple[Inner, ...]        SAVED  <- value lost
  read_values list[Inner]              COLLISION-ERROR   (this one flattens)
  UDF -> list[Inner]                   SAVED  <- value lost
  UDF -> dict[str, Inner]              SAVED  <- value lost
  

  Pre-existing: main loses the value on all of these, so fix: refuse a dict whose keys collide as JSON instead of losing a value #1943 changed which routes are
  caught rather than introducing the gap. Pinned by a strict xfail in
  test_save_refuses_colliding_dict_keys_inside_a_model, so it converts to a failure the
  moment it is fixed. The fix is to inspect the mapping before delegating to pydantic.