fix(mutate): simplify and drop mutated columns (#1298)

* fix(mutate): simplify and drop mutated columns

* cover more edge cases, add more tests

* fix tests, keep sys columns on mutate

Author: shcheklein

src/datachain/lib/dc/datachain.py

@@ -1262,8 +1262,10 @@ def mutate(self, **kwargs) -> "Self":
                 # adding new signal
                 mutated[name] = value
 
+        new_schema = schema.mutate(kwargs)
         return self._evolve(
-            query=self._query.mutate(**mutated), signal_schema=schema.mutate(kwargs)
+            query=self._query.mutate(new_schema=new_schema, **mutated),
+            signal_schema=new_schema,
         )
 
     @property

src/datachain/lib/signal_schema.py

@@ -34,7 +34,7 @@
 from datachain.lib.file import File
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainParamsError
-from datachain.query.schema import DEFAULT_DELIMITER, Column, ColumnMeta
+from datachain.query.schema import DEFAULT_DELIMITER, C, Column, ColumnMeta
 from datachain.sql.types import SQLType
 
 if TYPE_CHECKING:
@@ -680,35 +680,46 @@ def mutate(self, args_map: dict) -> "SignalSchema":
         primitives = (bool, str, int, float)
 
         for name, value in args_map.items():
+            current_type = None
+
+            if C.is_nested(name):
+                try:
+                    current_type = self.get_column_type(name)
+                except SignalResolvingError as err:
+                    msg = f"Creating new nested columns directly is not allowed: {name}"
+                    raise ValueError(msg) from err
+
             if isinstance(value, Column) and value.name in self.values:
                 # renaming existing signal
+                # Note: it won't touch nested signals here (e.g. file__path)
+                # we don't allow removing nested columns to keep objects consistent
                 del new_values[value.name]
                 new_values[name] = self.values[value.name]
-                continue
-            if isinstance(value, Column):
+            elif isinstance(value, Column):
                 # adding new signal from existing signal field
-                try:
-                    new_values[name] = self.get_column_type(
-                        value.name, with_subtree=True
-                    )
-                    continue
-                except SignalResolvingError:
-                    pass
-            if isinstance(value, Func):
+                new_values[name] = self.get_column_type(value.name, with_subtree=True)
+            elif isinstance(value, Func):
                 # adding new signal with function
                 new_values[name] = value.get_result_type(self)
-                continue
-            if isinstance(value, primitives):
+            elif isinstance(value, primitives):
                 # For primitives, store the type, not the value
                 val = literal(value)
                 val.type = python_to_sql(type(value))()
                 new_values[name] = sql_to_python(val)
-                continue
-            if isinstance(value, ColumnElement):
+            elif isinstance(value, ColumnElement):
                 # adding new signal
                 new_values[name] = sql_to_python(value)
-                continue
-            new_values[name] = value
+            else:
+                new_values[name] = value
+
+            if C.is_nested(name):
+                if current_type != new_values[name]:
+                    msg = (
+                        f"Altering nested column type is not allowed: {name}, "
+                        f"current type: {current_type}, new type: {new_values[name]}"
+                    )
+                    raise ValueError(msg)
+                del new_values[name]
 
         return SignalSchema(new_values)
 

src/datachain/query/dataset.py

@@ -10,7 +10,6 @@
 from collections.abc import Generator, Iterable, Iterator, Sequence
 from copy import copy
 from functools import wraps
-from secrets import token_hex
 from types import GeneratorType
 from typing import (
     TYPE_CHECKING,
@@ -29,7 +28,7 @@
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback, TqdmCallback
 from sqlalchemy import Column
 from sqlalchemy.sql import func as f
-from sqlalchemy.sql.elements import ColumnClause, ColumnElement
+from sqlalchemy.sql.elements import ColumnClause, ColumnElement, Label
 from sqlalchemy.sql.expression import label
 from sqlalchemy.sql.schema import TableClause
 from sqlalchemy.sql.selectable import Select
@@ -46,6 +45,7 @@
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
 from datachain.func.base import Function
 from datachain.lib.listing import is_listing_dataset, listing_dataset_expired
+from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
 from datachain.project import Project
@@ -795,28 +795,32 @@ def apply_sql_clause(self, query: Select) -> Select:
 
 @frozen
 class SQLMutate(SQLClause):
-    args: tuple[Union[Function, ColumnElement], ...]
+    args: tuple[Label, ...]
+    new_schema: SignalSchema
 
     def apply_sql_clause(self, query: Select) -> Select:
         original_subquery = query.subquery()
-        args = [
-            original_subquery.c[str(c)] if isinstance(c, (str, C)) else c
-            for c in self.parse_cols(self.args)
-        ]
-        to_mutate = {c.name for c in args}
+        to_mutate = {c.name for c in self.args}
 
-        prefix = f"mutate{token_hex(8)}_"
-        cols = [
-            c.label(prefix + c.name) if c.name in to_mutate else c
+        # Drop the original versions to avoid name collisions, exclude renamed
+        # columns. Always keep system columns (sys__*) if they exist in original query
+        new_schema_columns = set(self.new_schema.db_signals())
+        base_cols = [
+            c
             for c in original_subquery.c
+            if c.name not in to_mutate
+            and (c.name in new_schema_columns or c.name.startswith("sys__"))
         ]
-        # this is needed for new column to be used in clauses
-        # like ORDER BY, otherwise new column is not recognized
-        subquery = (
-            sqlalchemy.select(*cols, *args).select_from(original_subquery).subquery()
+
+        # Create intermediate subquery to properly handle window functions
+        intermediate_query = sqlalchemy.select(*base_cols, *self.args).select_from(
+            original_subquery
         )
+        intermediate_subquery = intermediate_query.subquery()
 
-        return sqlalchemy.select(*subquery.c).select_from(subquery)
+        return sqlalchemy.select(*intermediate_subquery.c).select_from(
+            intermediate_subquery
+        )
 
 
 @frozen
@@ -1470,7 +1474,7 @@ def select_except(self, *args) -> "Self":
         return query
 
     @detach
-    def mutate(self, *args, **kwargs) -> "Self":
+    def mutate(self, *args, new_schema, **kwargs) -> "Self":
         """
         Add new columns to this query.
 
@@ -1482,7 +1486,7 @@ def mutate(self, *args, **kwargs) -> "Self":
         """
         query_args = [v.label(k) for k, v in dict(args, **kwargs).items()]
         query = self.clone()
-        query.steps.append(SQLMutate((*query_args,)))
+        query.steps.append(SQLMutate((*query_args,), new_schema))
         return query
 
     @detach

src/datachain/query/schema.py

@@ -36,6 +36,10 @@ def to_db_name(name: str) -> str:
     def __getattr__(cls, name: str):
         return cls(ColumnMeta.to_db_name(name))
 
+    @staticmethod
+    def is_nested(name: str) -> bool:
+        return DEFAULT_DELIMITER in name
+
 
 class Column(sa.ColumnClause, metaclass=ColumnMeta):
     inherit_cache: Optional[bool] = True

tests/func/test_data_storage.py

@@ -17,7 +17,7 @@
 from tests.utils import (
     DEFAULT_TREE,
     TARRED_TREE,
-    create_tar_dataset_with_legacy_columns,
+    create_tar_dataset,
 )
 
 COMPLEX_TREE: dict[str, Any] = {
@@ -39,7 +39,7 @@ def test_dir_expansion(cloud_test_catalog, version_aware, cloud_type):
         # we don't want to index things in parent directory
         src_uri += "/"
 
-    chain = create_tar_dataset_with_legacy_columns(session, ctc.src_uri, "dc")
+    chain = create_tar_dataset(session, ctc.src_uri, "dc")
     dataset = catalog.get_dataset(chain.name)
     with catalog.warehouse.clone() as warehouse:
         dr = warehouse.dataset_rows(dataset, column="file")

tests/func/test_datachain.py

@@ -16,7 +16,6 @@
 import pytest
 import pytz
 from PIL import Image
-from sqlalchemy import Column
 
 import datachain as dc
 from datachain import DataModel, func
@@ -236,22 +235,6 @@ def test_read_storage_dependencies(cloud_test_catalog, cloud_type):
     assert dependencies[0].name == dep_name
 
 
-def test_persist_after_mutate(test_session):
-    chain = (
-        dc.read_values(fib=[1, 1, 2, 3, 5, 8, 13, 21], session=test_session)
-        .map(mod3=lambda fib: fib % 3, output=int)
-        .group_by(
-            cnt=dc.func.count(),
-            partition_by="mod3",
-        )
-        .mutate(x=1)
-        .persist()
-    )
-
-    assert chain.count() == 3
-    assert set(chain.to_values("mod3")) == {0, 1, 2}
-
-
 def test_persist_not_affects_dependencies(tmp_dir, test_session):
     for i in range(4):
         (tmp_dir / f"file{i}.txt").write_text(f"file{i}")
@@ -776,59 +759,6 @@ def test_read_storage_check_rows(tmp_dir, test_session):
         )
 
 
-def test_mutate_existing_column(test_session):
-    ds = dc.read_values(ids=[1, 2, 3], session=test_session)
-    ds = ds.mutate(ids=Column("ids") + 1)
-
-    assert ds.order_by("ids").to_list() == [(2,), (3,), (4,)]
-
-
-def test_mutate_with_primitives_save_load(test_session):
-    """Test that mutate with primitive values properly persists schema
-    through save/load cycle."""
-    original_data = [1, 2, 3]
-
-    # Create dataset with multiple primitive columns added via mutate
-    ds = dc.read_values(data=original_data, session=test_session).mutate(
-        str_col="test_string",
-        int_col=42,
-        float_col=3.14,
-        bool_col=True,
-    )
-
-    # Verify schema before saving
-    schema = ds.signals_schema.values
-    assert schema.get("str_col") is str
-    assert schema.get("int_col") is int
-    assert schema.get("float_col") is float
-    assert schema.get("bool_col") is bool
-
-    ds.save("test_mutate_primitives")
-
-    # Load the dataset back
-    loaded_ds = dc.read_dataset("test_mutate_primitives", session=test_session)
-
-    # Verify schema after loading
-    loaded_schema = loaded_ds.signals_schema.values
-    assert loaded_schema.get("str_col") is str
-    assert loaded_schema.get("int_col") is int
-    assert loaded_schema.get("float_col") is float
-    assert loaded_schema.get("bool_col") is bool
-
-    # Verify data integrity
-    results = set(loaded_ds.to_list())
-    assert len(results) == 3
-
-    # Expected tuples: (data, str_col, int_col, float_col, bool_col)
-    expected_results = {
-        (1, "test_string", 42, 3.14, True),
-        (2, "test_string", 42, 3.14, True),
-        (3, "test_string", 42, 3.14, True),
-    }
-
-    assert results == expected_results
-
-
 @pytest.mark.parametrize("processes", [False, 2, True])
 @pytest.mark.xdist_group(name="tmpfile")
 def test_parallel(processes, test_session_tmpfile):

tests/func/test_dataset_query.py

@@ -8,7 +8,9 @@
 
 from datachain.dataset import DatasetDependencyType, DatasetStatus
 from datachain.error import DatasetNotFoundError
+from datachain.lib.file import File
 from datachain.lib.listing import parse_listing_uri
+from datachain.lib.signal_schema import SignalSchema
 from datachain.query import C, DatasetQuery, Object, Stream
 from datachain.sql.functions import path as pathfunc
 from datachain.sql.types import String
@@ -19,6 +21,12 @@ def from_result_row(col_names, row):
     return dict(zip(col_names, row))
 
 
+def create_dataset_query_mutate_schema(**mutations):
+    schema_values = {"file": File}
+    schema_values.update(mutations)
+    return SignalSchema(schema_values)
+
+
 @pytest.fixture
 def dogs_cats_dataset(listed_bucket, cloud_test_catalog, dogs_dataset, cats_dataset):
     dataset_name = uuid.uuid4().hex
@@ -306,12 +314,14 @@ def test_distinct_count(cloud_test_catalog, animal_dataset):
 def test_mutate(cloud_test_catalog, save, animal_dataset):
     catalog = cloud_test_catalog.catalog
     ds = DatasetQuery(animal_dataset.name, catalog=catalog)
+    schema = create_dataset_query_mutate_schema(size10x=int, size1000x=int)
     q = (
-        ds.mutate(size10x=C("file.size") * 10)
-        .mutate(size1000x=C.size10x * 100)
+        ds.mutate(new_schema=schema, size10x=C("file.size") * 10)
+        .mutate(new_schema=schema, size1000x=C.size10x * 100)
         .mutate(
             ("s2", C("file.size") * 2),
             ("s3", C("file.size") * 3),
+            new_schema=schema,
             s4=C("file.size") * 4,
         )
         .filter((C.size10x < 40) | (C.size10x > 100) | C("file.path").glob("cat*"))
@@ -349,8 +359,9 @@ def test_mutate(cloud_test_catalog, save, animal_dataset):
 def test_order_by_after_mutate(cloud_test_catalog, save, animal_dataset):
     catalog = cloud_test_catalog.catalog
     ds = DatasetQuery(animal_dataset.name, catalog=catalog)
+    schema = create_dataset_query_mutate_schema(size10x=int)
     q = (
-        ds.mutate(size10x=C("file.size") * 10)
+        ds.mutate(new_schema=schema, size10x=C("file.size") * 10)
         .filter((C.size10x < 40) | (C.size10x > 100) | C("file.path").glob("cat*"))
         .order_by(C.size10x.desc())
     )
@@ -446,10 +457,12 @@ def test_offset_limit(cloud_test_catalog, save, animal_dataset):
 @pytest.mark.parametrize("save", [True, False])
 def test_mutate_offset_limit(cloud_test_catalog, save, animal_dataset):
     catalog = cloud_test_catalog.catalog
+    base_query = DatasetQuery(animal_dataset.name, catalog=catalog).order_by(
+        C("file.path")
+    )
+    schema = create_dataset_query_mutate_schema(size10x=int)
     q = (
-        DatasetQuery(animal_dataset.name, catalog=catalog)
-        .order_by(C("file.path"))
-        .mutate(size10x=C("file.size") * 10)
+        base_query.mutate(new_schema=schema, size10x=C("file.size") * 10)
         .offset(3)
         .limit(2)
     )