Union types - #1824

Related? 🤔

not really but it's better to keep 🙂
however, it's still messing up with comments 😅

A multi-output UDF can silently reinterpret an unrelated BaseModel as a declared union arm.
This only succeeds when the unrelated model’s flattened fields happen to match the union’s physical column count/order/types. That constraint makes it less common, but the outcome is silent data corruption rather than an error.

Minimal reproducer:

from typing import Union

import datachain as dc
from datachain.lib.data_model import DataModel
from datachain.query import Session


class Foo(DataModel):
    x: int


class Bar(DataModel):
    y: int


class Wrong(DataModel):
    tag: str
    bar_y: int | None
    foo_x: int


def bad(n):
    return n, Wrong(tag="Foo", bar_y=None, foo_x=123)


session = Session.get(in_memory=True)

out = dc.read_values(n=[1], session=session).map(
    bad,
    output={"n2": int, "value": Union[Foo, Bar]},
)

print(out.to_values("value"))

Actual: [Foo(x=123)]
Expected: DataChainParamsError: Wrong(...) does not match any arm of Union[Foo, Bar]

The same Wrong value is correctly rejected when the union is the UDF’s only output. The bug is specific to the multi-output path.

Could we distinguish a per-output union value from the wrapper/Parquet case here, reject unrelated models, and add a multi-output regression test?

fixed. A cover model is now recognised by its fields matching the remaining declared outputs, so a non-arm model raises DataChainParamsError exactly as in the single-output path.

This test keeps value on the preserved left side, so it does not exercise a missing union value. If value: Union[int, str] is on the right, this left merge produces None for the unmatched row while the schema remains Union[int, str]:

left = dc.read_values(id=[1, 2], session=test_session)
right = dc.read_values(
    right_id=[1],
    value=["x"],
    output={"right_id": int, "value": Union[int, str]},
    session=test_session,
)

result = left.merge(
    right, on="id", right_on="right_id", inner=False
).order_by("id")

assert result.to_list("id", "value") == [(1, "x"), (2, None)]
assert result.signals_schema.values["value"] == Union[int, str, None]

Currently the second assertion fails because SignalSchema.merge() only widens direct scalar types. Could we widen a tagged union to Union[..., None] when its merge side is nullable and add this reverse-direction regression test?

Fixed: SignalSchema._nullable() now widens a tagged union as well, and your reverse-direction case is in as test_merge_widens_union_on_the_padded_side.

If a model arm’s stable name is _type_tag, this assignment overwrites the discriminator inserted above:

class _type_tag(DataModel):
    x: int

annotation = Union[_type_tag, str]
schema = SignalSchema({"value": annotation})

assert schema.db_signals() == [
    "value___type_tag",
    "value___type_tag__x",
    "value__str",
]

Currently the first column is missing, and ingestion fails later with a misleading type error because the discriminator value is written into value___type_tag__x.

Could we reject _type_tag as a reserved arm selector during union-layout validation and add a regression test? The validation should use the stable ModelStore base name, since that is what arm_selector() returns.

Good catch - it was actually fixed in one place but not the other.

Fixed: added a validation with an exception.

Returning False here is expected for a collection arm, but on Python 3.10–3.13:

dc.read_values(
    value=[["x"]],
    output={"value": str | list[str]},
    session=test_session,
)

raises: AttributeError: 'types.UnionType' object has no attribute '__name__'

The equivalent typing.Union[str, list[str]] correctly raises UdfSignatureError. The difference comes from _validate_output() formatting the rejected annotation with value.__name__.

Could we format unsupported annotations safely and add both union spellings to a regression test? Collection-arm support can remain out of scope.

str | list[str] now raises UdfSignatureError like the typing.Union spelling

The final layout is selector-keyed (value__int / value__str), so this no longer resolves to a positional value._0 slot, and _type_tag stores the selector string rather than an index.
Could we sweep the other stale positive descriptions in datachain.py:1549, signal_schema.py:1191/1670, test_union_adversarial.py:55/63/301, and test_union_types.py:239/246 plus the _type_tag index comment in the unit suite? Assertions that _0/_1 never appear publicly can remain as regression checks.

Current representation:

from typing import Union

from datachain.lib.signal_schema import SignalSchema


schema = SignalSchema({"value": Union[int, str]})

print(schema.db_signals())

Output:

[
    "value___type_tag",
    "value__int",
    "value__str",
]

There are no value._0 or value._1 slots, and _type_tag stores "int" or "str", not an arm index.

ha, right! Fixed all the comments with old meaning.