fix(udf): don't allow silent wrong schema merges (#1460)

Author: shcheklein

src/datachain/lib/dc/datachain.py

@@ -850,14 +850,13 @@ def map(
         if (prefetch := self._settings.prefetch) is not None:
             udf_obj.prefetch = prefetch
 
+        sys_schema = SignalSchema({"sys": Sys})
         return self._evolve(
             query=self._query.add_signals(
                 udf_obj.to_udf_wrapper(self._settings.batch_size),
                 **self._settings.to_dict(),
             ),
-            signal_schema=SignalSchema({"sys": Sys})
-            | self.signals_schema
-            | udf_obj.output,
+            signal_schema=sys_schema | self.signals_schema | udf_obj.output,
         )
 
     def gen(

src/datachain/lib/signal_schema.py

@@ -34,7 +34,7 @@
 from datachain.lib.data_model import DataModel, DataType, DataValue
 from datachain.lib.file import File
 from datachain.lib.model_store import ModelStore
-from datachain.lib.utils import DataChainParamsError
+from datachain.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.query.schema import DEFAULT_DELIMITER, C, Column, ColumnMeta
 from datachain.sql.types import SQLType
 
@@ -1038,7 +1038,28 @@ def get_headers_with_length(self, include_hidden: bool = True):
         ], max_length
 
     def __or__(self, other):
-        return self.__class__(self.values | other.values)
+        new_values = dict(self.values)
+
+        for name, new_type in other.values.items():
+            if name in new_values:
+                current_type = new_values[name]
+                if current_type != new_type:
+                    raise DataChainColumnError(
+                        name,
+                        "signal already exists with a different type",
+                    )
+                continue
+
+            root = self._extract_root(name)
+            if any(self._extract_root(existing) == root for existing in new_values):
+                raise DataChainColumnError(
+                    name,
+                    "signal root already exists in schema",
+                )
+
+            new_values[name] = new_type
+
+        return self.__class__(new_values)
 
     def __contains__(self, name: str):
         return name in self.values

src/datachain/query/dataset.py

@@ -683,11 +683,26 @@ def create_result_query(
         signal_name_cols = {c.name: c for c in signal_cols}
         cols = signal_cols
 
-        overlap = {c.name for c in original_cols} & {c.name for c in cols}
+        original_names = {c.name for c in original_cols}
+        new_names = {c.name for c in cols}
+
+        overlap = original_names & new_names
         if overlap:
             raise ValueError(
                 "Column already exists or added in the previous steps: "
-                + ", ".join(overlap)
+                + ", ".join(sorted(overlap))
+            )
+
+        def _root(name: str) -> str:
+            return name.split(DEFAULT_DELIMITER, 1)[0]
+
+        existing_roots = {_root(name) for name in original_names}
+        new_roots = {_root(name) for name in new_names}
+        root_conflicts = existing_roots & new_roots
+        if root_conflicts:
+            raise ValueError(
+                "Signals already exist in the previous steps: "
+                + ", ".join(sorted(root_conflicts))
             )
 
         def q(*columns):

tests/func/test_udf.py

@@ -12,7 +12,7 @@
 from datachain.func import path as pathfunc
 from datachain.lib.file import AudioFile, AudioFragment, File
 from datachain.lib.udf import Mapper
-from datachain.lib.utils import DataChainError
+from datachain.lib.utils import DataChainColumnError, DataChainError
 from tests.utils import LARGE_TREE, NUM_TREE
 
 
@@ -390,6 +390,23 @@ def test_types():
     ]
 
 
+def test_udf_rejects_root_override(test_session):
+    class X(dc.DataModel):
+        x: int
+
+    chain = dc.read_values(x=[X(x=0), X(x=1)], session=test_session)
+
+    with pytest.raises(
+        DataChainColumnError,
+        match="Error for column x: signal already exists with a different type",
+    ):
+        chain.map(
+            lambda x: x.model_dump(),
+            params=["x"],
+            output={"x": dict},
+        )
+
+
 @pytest.mark.parametrize("use_cache", [False, True])
 @pytest.mark.parametrize("prefetch", [0, 2])
 def test_map_file(cloud_test_catalog, use_cache, prefetch, monkeypatch):

tests/unit/lib/test_signal_schema.py

@@ -17,6 +17,7 @@
     SignalSchemaError,
     SignalSchemaWarning,
 )
+from datachain.lib.utils import DataChainColumnError
 from datachain.sql.types import (
     JSON,
     Array,
@@ -185,6 +186,45 @@ def test_feature_schema_serialize_list():
     assert deserialized_schema.values == schema
 
 
+def test_schema_or_rejects_type_change():
+    base = SignalSchema({"foo": int})
+    new = SignalSchema({"foo": str})
+
+    with pytest.raises(DataChainColumnError, match="different type"):
+        _ = base | new
+
+
+def test_schema_or_rejects_root_conflict():
+    base = SignalSchema({"feature": MyType1})
+    new = SignalSchema({"feature.extra": int})
+
+    with pytest.raises(DataChainColumnError, match="root"):
+        _ = base | new
+
+
+def test_schema_or_allows_sys_root():
+    base = SignalSchema({"foo": int})
+    new = SignalSchema({"sys": Sys})
+
+    combined = base | new
+    assert combined.values["sys"] is Sys
+
+
+def test_schema_or_rejects_sys_override():
+    base = SignalSchema({"sys": Sys})
+    new = SignalSchema({"sys": dict})
+
+    with pytest.raises(DataChainColumnError, match="different type"):
+        _ = base | new
+
+
+def test_schema_or_allows_identical_signal():
+    base = SignalSchema({"foo": int})
+    combined = base | SignalSchema({"foo": int})
+
+    assert combined.values["foo"] is int
+
+
 def test_feature_schema_serialize_list_old():
     schema = {
         "name": str | None,