fix(steps): drop sys columns when it is not safe to keep them (#1400)

Author: shcheklein

src/datachain/data_storage/schema.py

@@ -11,7 +11,6 @@
     JSON,
     Boolean,
     DateTime,
-    Int,
     Int64,
     SQLType,
     String,
@@ -269,7 +268,7 @@ def delete(self):
     @classmethod
     def sys_columns(cls):
         return [
-            sa.Column("sys__id", Int, primary_key=True),
+            sa.Column("sys__id", UInt64, primary_key=True),
             sa.Column(
                 "sys__rand", UInt64, nullable=False, server_default=f.abs(f.random())
             ),

src/datachain/data_storage/sqlite.py

@@ -868,11 +868,8 @@ def add_left_rows_filter(exp: BinaryExpression):
                 if isinstance(c, BinaryExpression):
                     right_left_join = add_left_rows_filter(c)
 
-        # Use CTE instead of subquery to force SQLite to materialize the result
-        # This breaks deep nesting and prevents parser stack overflow.
         union_cte = sqlalchemy.union(left_right_join, right_left_join).cte()
-
-        return self._regenerate_system_columns(union_cte)
+        return sqlalchemy.select(*union_cte.c).select_from(union_cte)
 
     def _system_row_number_expr(self):
         return func.row_number().over()
@@ -884,11 +881,7 @@ def create_pre_udf_table(self, query: "Select") -> "Table":
         """
         Create a temporary table from a query for use in a UDF.
         """
-        columns = [
-            sqlalchemy.Column(c.name, c.type)
-            for c in query.selected_columns
-            if c.name != "sys__id"
-        ]
+        columns = [sqlalchemy.Column(c.name, c.type) for c in query.selected_columns]
         table = self.create_udf_table(columns)
 
         with tqdm(desc="Preparing", unit=" rows", leave=False) as pbar:

src/datachain/data_storage/warehouse.py

@@ -5,7 +5,7 @@
 import string
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Union, cast
 from urllib.parse import urlparse
 
 import attrs
@@ -23,7 +23,7 @@
 from datachain.query.batch import RowsOutput
 from datachain.query.schema import ColumnMeta
 from datachain.sql.functions import path as pathfunc
-from datachain.sql.types import Int, SQLType
+from datachain.sql.types import SQLType
 from datachain.utils import sql_escape_like
 
 if TYPE_CHECKING:
@@ -32,6 +32,7 @@
         _FromClauseArgument,
         _OnClauseArgument,
     )
+    from sqlalchemy.sql.selectable import FromClause
     from sqlalchemy.types import TypeEngine
 
     from datachain.data_storage import schema
@@ -248,45 +249,56 @@ def dataset_select_paginated(
 
     def _regenerate_system_columns(
         self,
-        selectable: sa.Select | sa.CTE,
+        selectable: sa.Select,
         keep_existing_columns: bool = False,
+        regenerate_columns: Iterable[str] | None = None,
     ) -> sa.Select:
         """
-        Return a SELECT that regenerates sys__id and sys__rand deterministically.
+        Return a SELECT that regenerates system columns deterministically.
 
-        If keep_existing_columns is True, existing sys__id and sys__rand columns
-        will be kept as-is if they exist in the input selectable.
-        """
-        base = selectable.subquery() if hasattr(selectable, "subquery") else selectable
-
-        result_columns: dict[str, sa.ColumnElement] = {}
-        for col in base.c:
-            if col.name in result_columns:
-                raise ValueError(f"Duplicate column name {col.name} in SELECT")
-            if col.name in ("sys__id", "sys__rand"):
-                if keep_existing_columns:
-                    result_columns[col.name] = col
-            else:
-                result_columns[col.name] = col
+        If keep_existing_columns is True, existing system columns will be kept as-is
+        even when they are listed in ``regenerate_columns``.
 
-        system_types: dict[str, sa.types.TypeEngine] = {
+        Args:
+            selectable: Base SELECT
+            keep_existing_columns: When True, reuse existing system columns even if
+                they are part of the regeneration set.
+            regenerate_columns: Names of system columns to regenerate. Defaults to
+                {"sys__id", "sys__rand"}. Columns not listed are left untouched.
+        """
+        system_columns = {
             sys_col.name: sys_col.type
             for sys_col in self.schema.dataset_row_cls.sys_columns()
         }
+        regenerate = set(regenerate_columns or system_columns)
+        generators = {
+            "sys__id": self._system_row_number_expr,
+            "sys__rand": self._system_random_expr,
+        }
+
+        base = cast("FromClause", selectable.subquery())
+
+        def build(name: str) -> sa.ColumnElement:
+            expr = generators[name]()
+            return sa.cast(expr, system_columns[name]).label(name)
+
+        columns: list[sa.ColumnElement] = []
+        present: set[str] = set()
+        changed = False
+
+        for col in base.c:
+            present.add(col.name)
+            regen = col.name in regenerate and not keep_existing_columns
+            columns.append(build(col.name) if regen else col)
+            changed |= regen
+
+        for name in regenerate - present:
+            columns.append(build(name))
+            changed = True
+
+        if not changed:
+            return selectable
 
-        # Add missing system columns if needed
-        if "sys__id" not in result_columns:
-            expr = self._system_row_number_expr()
-            expr = sa.cast(expr, system_types["sys__id"])
-            result_columns["sys__id"] = expr.label("sys__id")
-        if "sys__rand" not in result_columns:
-            expr = self._system_random_expr()
-            expr = sa.cast(expr, system_types["sys__rand"])
-            result_columns["sys__rand"] = expr.label("sys__rand")
-
-        # Wrap in subquery to materialize window functions, then wrap again in SELECT
-        # This ensures window functions are computed before INSERT...FROM SELECT
-        columns = list(result_columns.values())
         inner = sa.select(*columns).select_from(base).subquery()
         return sa.select(*inner.c).select_from(inner)
 
@@ -950,10 +962,15 @@ def create_udf_table(
         SQLite TEMPORARY tables cannot be directly used as they are process-specific,
         and UDFs are run in other processes when run in parallel.
         """
+        columns = [
+            c
+            for c in columns
+            if c.name not in [col.name for col in self.dataset_row_cls.sys_columns()]
+        ]
         tbl = sa.Table(
             name or self.udf_table_name(),
             sa.MetaData(),
-            sa.Column("sys__id", Int, primary_key=True),
+            *self.dataset_row_cls.sys_columns(),
             *columns,
         )
         self.db.create_table(tbl, if_not_exists=True)

src/datachain/diff/__init__.py

@@ -24,7 +24,7 @@ class CompareStatus(str, Enum):
     SAME = "S"
 
 
-def _compare(  # noqa: C901, PLR0912
+def _compare(  # noqa: C901
     left: "DataChain",
     right: "DataChain",
     on: str | Sequence[str],
@@ -151,11 +151,7 @@ def _to_list(obj: str | Sequence[str] | None) -> list[str] | None:
     if status_col:
         cols_select.append(diff_col)
 
-    if not dc_diff._sys:
-        # TODO workaround when sys signal is not available in diff
-        dc_diff = dc_diff.settings(sys=True).select(*cols_select).settings(sys=False)
-    else:
-        dc_diff = dc_diff.select(*cols_select)
+    dc_diff = dc_diff.select(*cols_select)
 
     # final schema is schema from the left chain with status column added if needed
     dc_diff.signals_schema = (

src/datachain/lib/dc/datachain.py

@@ -856,7 +856,9 @@ def map(
                 udf_obj.to_udf_wrapper(self._settings.batch_size),
                 **self._settings.to_dict(),
             ),
-            signal_schema=self.signals_schema | udf_obj.output,
+            signal_schema=SignalSchema({"sys": Sys})
+            | self.signals_schema
+            | udf_obj.output,
         )
 
     def gen(
@@ -894,7 +896,7 @@ def gen(
                 udf_obj.to_udf_wrapper(self._settings.batch_size),
                 **self._settings.to_dict(),
             ),
-            signal_schema=udf_obj.output,
+            signal_schema=SignalSchema({"sys": Sys}) | udf_obj.output,
         )
 
     @delta_disabled
@@ -1031,7 +1033,7 @@ def my_agg(files: list[File]) -> Iterator[tuple[File, int]]:
                 partition_by=processed_partition_by,
                 **self._settings.to_dict(),
             ),
-            signal_schema=udf_obj.output,
+            signal_schema=SignalSchema({"sys": Sys}) | udf_obj.output,
         )
 
     def batch_map(
@@ -1097,11 +1099,7 @@ def _udf_to_obj(
         sign = UdfSignature.parse(name, signal_map, func, params, output, is_generator)
         DataModel.register(list(sign.output_schema.values.values()))
 
-        signals_schema = self.signals_schema
-        if self._sys:
-            signals_schema = SignalSchema({"sys": Sys}) | signals_schema
-
-        params_schema = signals_schema.slice(
+        params_schema = self.signals_schema.slice(
             sign.params, self._setup, is_batch=is_batch
         )
 
@@ -1156,11 +1154,9 @@ def distinct(self, arg: str, *args: str) -> "Self":  # type: ignore[override]
             )
         )
 
-    def select(self, *args: str, _sys: bool = True) -> "Self":
+    def select(self, *args: str) -> "Self":
         """Select only a specified set of signals."""
         new_schema = self.signals_schema.resolve(*args)
-        if self._sys and _sys:
-            new_schema = SignalSchema({"sys": Sys}) | new_schema
         columns = new_schema.db_signals()
         return self._evolve(
             query=self._query.select(*columns), signal_schema=new_schema
@@ -1710,9 +1706,11 @@ def _resolve(
 
         signals_schema = self.signals_schema.clone_without_sys_signals()
         right_signals_schema = right_ds.signals_schema.clone_without_sys_signals()
-        ds.signals_schema = SignalSchema({"sys": Sys}) | signals_schema.merge(
-            right_signals_schema, rname
-        )
+
+        ds.signals_schema = signals_schema.merge(right_signals_schema, rname)
+
+        if not full:
+            ds.signals_schema = SignalSchema({"sys": Sys}) | ds.signals_schema
 
         return ds
 
@@ -1723,6 +1721,7 @@ def union(self, other: "Self") -> "Self":
         Parameters:
             other: chain whose rows will be added to `self`.
         """
+        self.signals_schema = self.signals_schema.clone_without_sys_signals()
         return self._evolve(query=self._query.union(other._query))
 
     def subtract(  # type: ignore[override]

src/datachain/query/dataset.py

@@ -438,9 +438,6 @@ def create_result_query(
         """
 
     def populate_udf_table(self, udf_table: "Table", query: Select) -> None:
-        if "sys__id" not in query.selected_columns:
-            raise RuntimeError("Query must have sys__id column to run UDF")
-
         if (rows_total := self.catalog.warehouse.query_count(query)) == 0:
             return
 
@@ -634,12 +631,11 @@ def apply(
 
         # Apply partitioning if needed.
         if self.partition_by is not None:
-            if "sys__id" not in query.selected_columns:
-                _query = query = self.catalog.warehouse._regenerate_system_columns(
-                    query,
-                    keep_existing_columns=True,
-                )
-
+            _query = query = self.catalog.warehouse._regenerate_system_columns(
+                query_generator.select(),
+                keep_existing_columns=True,
+                regenerate_columns=["sys__id"],
+            )
             partition_tbl = self.create_partitions_table(query)
             temp_tables.append(partition_tbl.name)
             query = query.outerjoin(
@@ -960,28 +956,23 @@ def apply(
         q2 = self.query2.apply_steps().select().subquery()
         temp_tables.extend(self.query2.temp_table_names)
 
-        columns1, columns2 = _order_columns(q1.columns, q2.columns)
-
-        union_select = sqlalchemy.select(*columns1).union_all(
-            sqlalchemy.select(*columns2)
-        )
-        union_cte = union_select.cte()
-        regenerated = self.query1.catalog.warehouse._regenerate_system_columns(
-            union_cte
-        )
-        result_columns = tuple(regenerated.selected_columns)
+        columns1 = _drop_system_columns(q1.columns)
+        columns2 = _drop_system_columns(q2.columns)
+        columns1, columns2 = _order_columns(columns1, columns2)
 
         def q(*columns):
-            if not columns:
-                return regenerated
+            selected_names = [c.name for c in columns]
+            col1 = [c for c in columns1 if c.name in selected_names]
+            col2 = [c for c in columns2 if c.name in selected_names]
+            union_query = sqlalchemy.select(*col1).union_all(sqlalchemy.select(*col2))
 
-            names = {c.name for c in columns}
-            selected = [c for c in result_columns if c.name in names]
-            return regenerated.with_only_columns(*selected)
+            union_cte = union_query.cte()
+            select_cols = [union_cte.c[name] for name in selected_names]
+            return sqlalchemy.select(*select_cols)
 
         return step_result(
             q,
-            result_columns,
+            columns1,
             dependencies=self.query1.dependencies | self.query2.dependencies,
         )
 
@@ -1070,7 +1061,7 @@ def apply(
         q1 = self.get_query(self.query1, temp_tables)
         q2 = self.get_query(self.query2, temp_tables)
 
-        q1_columns = list(q1.c)
+        q1_columns = _drop_system_columns(q1.c) if self.full else list(q1.c)
         q1_column_names = {c.name for c in q1_columns}
 
         q2_columns = []
@@ -1211,6 +1202,10 @@ def _order_columns(
     return [[d[n] for n in column_order] for d in column_dicts]
 
 
+def _drop_system_columns(columns: Iterable[ColumnElement]) -> list[ColumnElement]:
+    return [c for c in columns if not c.name.startswith("sys__")]
+
+
 @attrs.define
 class ResultIter:
     _row_iter: Iterable[Any]

tests/func/test_datachain.py

@@ -1629,7 +1629,7 @@ def test_read_pandas_multiindex(test_session):
 
     # Check the resulting column names and data
     expected_columns = ["a_cat", "b_dog", "b_cat", "a_dog"]
-    assert set(chain.signals_schema.db_signals()) == set(expected_columns)
+    assert set(chain.schema.keys()) == set(expected_columns)
 
     expected_data = [
         {"a_cat": 1, "b_dog": 2, "b_cat": 3, "a_dog": 4},