datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/query/dataset.py

@@ -1,25 +1,18 @@
 import contextlib
+import hashlib
 import inspect
 import logging
 import os
-import random
+import secrets
 import string
 import subprocess
 import sys
 from abc import ABC, abstractmethod
-from collections.abc import Generator, Iterable, Iterator, Sequence
+from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
 from copy import copy
 from functools import wraps
 from types import GeneratorType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Optional,
-    Protocol,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
 
 import attrs
 import sqlalchemy
@@ -44,20 +37,21 @@ from datachain.data_storage.schema import (
 from datachain.dataset import DatasetDependency, DatasetStatus, RowDict
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
 from datachain.func.base import Function
+from datachain.hash_utils import hash_column_elements
 from datachain.lib.listing import is_listing_dataset, listing_dataset_expired
-from datachain.lib.signal_schema import SignalSchema
+from datachain.lib.signal_schema import SignalSchema, generate_merge_root_mapping
 from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
 from datachain.project import Project
-from datachain.query.schema import C, UDFParamSpec, normalize_param
+from datachain.query.schema import DEFAULT_DELIMITER, C, UDFParamSpec, normalize_param
 from datachain.query.session import Session
 from datachain.query.udf import UdfInfo
 from datachain.sql.functions.random import rand
 from datachain.sql.types import SQLType
 from datachain.utils import (
-    batched,
     determine_processes,
     determine_workers,
+    ensure_sequence,
     filtered_cloudpickle_dumps,
     get_datachain_executable,
     safe_closing,
@@ -65,11 +59,12 @@ from datachain.utils import (
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
+    from typing import Concatenate
 
-    from sqlalchemy.sql.elements import ClauseElement
+    from sqlalchemy.sql.elements import ClauseElement, KeyedColumnElement
     from sqlalchemy.sql.schema import Table
     from sqlalchemy.sql.selectable import GenerativeSelect
-    from typing_extensions import Concatenate, ParamSpec, Self
+    from typing_extensions import ParamSpec, Self
 
     from datachain.catalog import Catalog
     from datachain.data_storage import AbstractWarehouse
@@ -81,13 +76,10 @@ if TYPE_CHECKING:
 
 INSERT_BATCH_SIZE = 10000
 
-PartitionByType = Union[
-    str,
-    Function,
-    ColumnElement,
-    Sequence[Union[str, Function, ColumnElement]],
-]
-JoinPredicateType = Union[str, ColumnClause, ColumnElement]
+PartitionByType = (
+    str | Function | ColumnElement | Sequence[str | Function | ColumnElement]
+)
+JoinPredicateType = str | ColumnClause | ColumnElement
 DatasetDependencyType = tuple["DatasetRecord", str]
 
 logger = logging.getLogger("datachain")
@@ -168,6 +160,18 @@ class Step(ABC):
     ) -> "StepResult":
         """Apply the processing step."""
 
+    @abstractmethod
+    def hash_inputs(self) -> str:
+        """Calculates hash of step inputs"""
+
+    def hash(self) -> str:
+        """
+        Calculates hash for step which includes step name and hash of it's inputs
+        """
+        return hashlib.sha256(
+            f"{self.__class__.__name__}|{self.hash_inputs()}".encode()
+        ).hexdigest()
+
 
 @frozen
 class QueryStep:
@@ -187,6 +191,11 @@ class QueryStep:
             q, dr.columns, dependencies=[(self.dataset, self.dataset_version)]
         )
 
+    def hash(self) -> str:
+        return hashlib.sha256(
+            self.dataset.uri(self.dataset_version).encode()
+        ).hexdigest()
+
 
 def generator_then_call(generator, func: Callable):
     """
@@ -222,8 +231,9 @@ class DatasetDiffOperation(Step):
 
     def apply(self, query_generator, temp_tables: list[str]) -> "StepResult":
         source_query = query_generator.exclude(("sys__id",))
+        right_before = len(self.dq.temp_table_names)
         target_query = self.dq.apply_steps().select()
-        temp_tables.extend(self.dq.temp_table_names)
+        temp_tables.extend(self.dq.temp_table_names[right_before:])
 
         # creating temp table that will hold subtract results
         temp_table_name = self.catalog.warehouse.temp_table_name()
@@ -257,6 +267,13 @@ class DatasetDiffOperation(Step):
 class Subtract(DatasetDiffOperation):
     on: Sequence[tuple[str, str]]
 
+    def hash_inputs(self) -> str:
+        on_bytes = b"".join(
+            f"{a}:{b}".encode() for a, b in sorted(self.on, key=lambda t: (t[0], t[1]))
+        )
+
+        return hashlib.sha256(bytes.fromhex(self.dq.hash()) + on_bytes).hexdigest()
+
     def query(self, source_query: Select, target_query: Select) -> sa.Selectable:
         sq = source_query.alias("source_query")
         tq = target_query.alias("target_query")
@@ -334,10 +351,10 @@ def process_udf_outputs(
     udf_results: Iterator[Iterable["UDFResult"]],
     udf: "UDFAdapter",
     cb: Callback = DEFAULT_CALLBACK,
+    batch_size: int = INSERT_BATCH_SIZE,
 ) -> None:
     # Optimization: Compute row types once, rather than for every row.
     udf_col_types = get_col_types(warehouse, udf.output)
-    batch_rows = udf.batch_rows or INSERT_BATCH_SIZE
 
     def _insert_rows():
         for udf_output in udf_results:
@@ -349,9 +366,7 @@ def process_udf_outputs(
                     cb.relative_update()
                     yield adjust_outputs(warehouse, row, udf_col_types)
 
-    for row_chunk in batched(_insert_rows(), batch_rows):
-        warehouse.insert_rows(udf_table, row_chunk)
-
+    warehouse.insert_rows(udf_table, _insert_rows(), batch_size=batch_size)
     warehouse.insert_rows_done(udf_table)
 
 
@@ -387,21 +402,34 @@ def get_generated_callback(is_generator: bool = False) -> Callback:
 class UDFStep(Step, ABC):
     udf: "UDFAdapter"
     catalog: "Catalog"
-    partition_by: Optional[PartitionByType] = None
-    parallel: Optional[int] = None
-    workers: Union[bool, int] = False
-    min_task_size: Optional[int] = None
+    partition_by: PartitionByType | None = None
     is_generator = False
+    # Parameters from Settings
     cache: bool = False
-    batch_rows: Optional[int] = None
+    parallel: int | None = None
+    workers: bool | int = False
+    min_task_size: int | None = None
+    batch_size: int | None = None
+
+    def hash_inputs(self) -> str:
+        partition_by = ensure_sequence(self.partition_by or [])
+        parts = [
+            bytes.fromhex(self.udf.hash()),
+            bytes.fromhex(hash_column_elements(partition_by)),
+            str(self.is_generator).encode(),
+        ]
+
+        return hashlib.sha256(b"".join(parts)).hexdigest()
 
     @abstractmethod
     def create_udf_table(self, query: Select) -> "Table":
         """Method that creates a table where temp udf results will be saved"""
 
     def process_input_query(self, query: Select) -> tuple[Select, list["Table"]]:
-        """Apply any necessary processing to the input query"""
-        return query, []
+        """Materialize inputs, ensure sys columns are available, needed for checkpoints,
+        needed for map to work (merge results)"""
+        table = self.catalog.warehouse.create_pre_udf_table(query)
+        return sqlalchemy.select(*table.c), [table]
 
     @abstractmethod
     def create_result_query(
@@ -450,6 +478,7 @@ class UDFStep(Step, ABC):
                         use_cache=self.cache,
                         is_generator=self.is_generator,
                         min_task_size=self.min_task_size,
+                        batch_size=self.batch_size,
                     )
                     udf_distributor()
                     return
@@ -486,6 +515,7 @@ class UDFStep(Step, ABC):
                         is_generator=self.is_generator,
                         cache=self.cache,
                         rows_total=rows_total,
+                        batch_size=self.batch_size or INSERT_BATCH_SIZE,
                     )
 
                     # Run the UDFDispatcher in another process to avoid needing
@@ -534,6 +564,7 @@ class UDFStep(Step, ABC):
                                 udf_results,
                                 self.udf,
                                 cb=generated_cb,
+                                batch_size=self.batch_size or INSERT_BATCH_SIZE,
                             )
                     finally:
                         download_cb.close()
@@ -552,13 +583,10 @@ class UDFStep(Step, ABC):
         """
         Create temporary table with group by partitions.
         """
-        # Check if partition_by is set, we need it to create partitions.
-        assert self.partition_by is not None
-        # Check if sys__id is in the query, we need it to be able to join
-        # the partition table with the udf table later.
-        assert any(c.name == "sys__id" for c in query.selected_columns), (
-            "Query must have sys__id column to use partitioning."
-        )
+        if self.partition_by is None:
+            raise RuntimeError("Query must have partition_by set to use partitioning")
+        if (id_col := query.selected_columns.get("sys__id")) is None:
+            raise RuntimeError("Query must have sys__id column to use partitioning")
 
         if isinstance(self.partition_by, (list, tuple, GeneratorType)):
             list_partition_by = list(self.partition_by)
@@ -574,7 +602,7 @@ class UDFStep(Step, ABC):
 
         # fill table with partitions
         cols = [
-            query.selected_columns.sys__id,
+            id_col,
             f.dense_rank().over(order_by=partition_by).label(PARTITION_COLUMN_ID),
         ]
         self.catalog.warehouse.db.execute(
@@ -586,7 +614,7 @@ class UDFStep(Step, ABC):
 
         return tbl
 
-    def clone(self, partition_by: Optional[PartitionByType] = None) -> "Self":
+    def clone(self, partition_by: PartitionByType | None = None) -> "Self":
         if partition_by is not None:
             return self.__class__(
                 self.udf,
@@ -595,41 +623,25 @@ class UDFStep(Step, ABC):
                 parallel=self.parallel,
                 workers=self.workers,
                 min_task_size=self.min_task_size,
-                batch_rows=self.batch_rows,
+                batch_size=self.batch_size,
             )
         return self.__class__(self.udf, self.catalog)
 
     def apply(
         self, query_generator: QueryGenerator, temp_tables: list[str]
     ) -> "StepResult":
-        _query = query = query_generator.select()
+        query, tables = self.process_input_query(query_generator.select())
+        _query = query
 
         # Apply partitioning if needed.
         if self.partition_by is not None:
-            if not any(c.name == "sys__id" for c in query.selected_columns):
-                # If sys__id is not in the query, we need to create a temp table
-                # to hold the query results, so we can join it with the
-                # partition table later.
-                columns = [
-                    c if isinstance(c, Column) else Column(c.name, c.type)
-                    for c in query.subquery().columns
-                ]
-                temp_table = self.catalog.warehouse.create_dataset_rows_table(
-                    self.catalog.warehouse.temp_table_name(),
-                    columns=columns,
-                )
-                temp_tables.append(temp_table.name)
-                self.catalog.warehouse.copy_table(temp_table, query)
-                _query = query = temp_table.select()
-
             partition_tbl = self.create_partitions_table(query)
-            temp_tables.append(partition_tbl.name)
             query = query.outerjoin(
                 partition_tbl,
                 partition_tbl.c.sys__id == query.selected_columns.sys__id,
             ).add_columns(*partition_columns())
+            tables = [*tables, partition_tbl]
 
-        query, tables = self.process_input_query(query)
         temp_tables.extend(t.name for t in tables)
         udf_table = self.create_udf_table(_query)
         temp_tables.append(udf_table.name)
@@ -641,7 +653,16 @@ class UDFStep(Step, ABC):
 
 @frozen
 class UDFSignal(UDFStep):
+    udf: "UDFAdapter"
+    catalog: "Catalog"
+    partition_by: PartitionByType | None = None
     is_generator = False
+    # Parameters from Settings
+    cache: bool = False
+    parallel: int | None = None
+    workers: bool | int = False
+    min_task_size: int | None = None
+    batch_size: int | None = None
 
     def create_udf_table(self, query: Select) -> "Table":
         udf_output_columns: list[sqlalchemy.Column[Any]] = [
@@ -651,13 +672,6 @@ class UDFSignal(UDFStep):
 
         return self.catalog.warehouse.create_udf_table(udf_output_columns)
 
-    def process_input_query(self, query: Select) -> tuple[Select, list["Table"]]:
-        if os.getenv("DATACHAIN_DISABLE_QUERY_CACHE", "") not in ("", "0"):
-            return query, []
-        table = self.catalog.warehouse.create_pre_udf_table(query)
-        q: Select = sqlalchemy.select(*table.c)
-        return q, [table]
-
     def create_result_query(
         self, udf_table, query
     ) -> tuple[QueryGeneratorFunc, list["sqlalchemy.Column"]]:
@@ -669,11 +683,26 @@ class UDFSignal(UDFStep):
         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):
@@ -711,7 +740,16 @@ class UDFSignal(UDFStep):
 class RowGenerator(UDFStep):
     """Extend dataset with new rows."""
 
+    udf: "UDFAdapter"
+    catalog: "Catalog"
+    partition_by: PartitionByType | None = None
     is_generator = True
+    # Parameters from Settings
+    cache: bool = False
+    parallel: int | None = None
+    workers: bool | int = False
+    min_task_size: int | None = None
+    batch_size: int | None = None
 
     def create_udf_table(self, query: Select) -> "Table":
         warehouse = self.catalog.warehouse
@@ -758,18 +796,42 @@ class SQLClause(Step, ABC):
 
     def parse_cols(
         self,
-        cols: Sequence[Union[Function, ColumnElement]],
+        cols: Sequence[Function | ColumnElement],
     ) -> tuple[ColumnElement, ...]:
         return tuple(c.get_column() if isinstance(c, Function) else c for c in cols)
 
     @abstractmethod
-    def apply_sql_clause(self, query):
+    def apply_sql_clause(self, query: Any) -> Any:
         pass
 
 
+@frozen
+class RegenerateSystemColumns(Step):
+    catalog: "Catalog"
+
+    def hash_inputs(self) -> str:
+        return hashlib.sha256(b"regenerate_system_columns").hexdigest()
+
+    def apply(
+        self, query_generator: QueryGenerator, temp_tables: list[str]
+    ) -> StepResult:
+        query = query_generator.select()
+        new_query = self.catalog.warehouse._regenerate_system_columns(
+            query, keep_existing_columns=True
+        )
+
+        def q(*columns):
+            return new_query.with_only_columns(*columns)
+
+        return step_result(q, new_query.selected_columns)
+
+
 @frozen
 class SQLSelect(SQLClause):
-    args: tuple[Union[Function, ColumnElement], ...]
+    args: tuple[Function | ColumnElement, ...]
+
+    def hash_inputs(self) -> str:
+        return hash_column_elements(self.args)
 
     def apply_sql_clause(self, query) -> Select:
         subquery = query.subquery()
@@ -785,7 +847,10 @@ class SQLSelect(SQLClause):
 
 @frozen
 class SQLSelectExcept(SQLClause):
-    args: tuple[Union[Function, ColumnElement], ...]
+    args: tuple[Function | ColumnElement, ...]
+
+    def hash_inputs(self) -> str:
+        return hash_column_elements(self.args)
 
     def apply_sql_clause(self, query: Select) -> Select:
         subquery = query.subquery()
@@ -798,6 +863,9 @@ class SQLMutate(SQLClause):
     args: tuple[Label, ...]
     new_schema: SignalSchema
 
+    def hash_inputs(self) -> str:
+        return hash_column_elements(self.args)
+
     def apply_sql_clause(self, query: Select) -> Select:
         original_subquery = query.subquery()
         to_mutate = {c.name for c in self.args}
@@ -825,7 +893,10 @@ class SQLMutate(SQLClause):
 
 @frozen
 class SQLFilter(SQLClause):
-    expressions: tuple[Union[Function, ColumnElement], ...]
+    expressions: tuple[Function | ColumnElement, ...]
+
+    def hash_inputs(self) -> str:
+        return hash_column_elements(self.expressions)
 
     def __and__(self, other):
         expressions = self.parse_cols(self.expressions)
@@ -838,7 +909,10 @@ class SQLFilter(SQLClause):
 
 @frozen
 class SQLOrderBy(SQLClause):
-    args: tuple[Union[Function, ColumnElement], ...]
+    args: tuple[Function | ColumnElement, ...]
+
+    def hash_inputs(self) -> str:
+        return hash_column_elements(self.args)
 
     def apply_sql_clause(self, query: Select) -> Select:
         args = self.parse_cols(self.args)
@@ -849,6 +923,9 @@ class SQLOrderBy(SQLClause):
 class SQLLimit(SQLClause):
     n: int
 
+    def hash_inputs(self) -> str:
+        return hashlib.sha256(str(self.n).encode()).hexdigest()
+
     def apply_sql_clause(self, query: Select) -> Select:
         return query.limit(self.n)
 
@@ -857,12 +934,18 @@ class SQLLimit(SQLClause):
 class SQLOffset(SQLClause):
     offset: int
 
+    def hash_inputs(self) -> str:
+        return hashlib.sha256(str(self.offset).encode()).hexdigest()
+
     def apply_sql_clause(self, query: "GenerativeSelect"):
         return query.offset(self.offset)
 
 
 @frozen
 class SQLCount(SQLClause):
+    def hash_inputs(self) -> str:
+        return ""
+
     def apply_sql_clause(self, query):
         return sqlalchemy.select(f.count(1)).select_from(query.subquery())
 
@@ -872,6 +955,9 @@ class SQLDistinct(SQLClause):
     args: tuple[ColumnElement, ...]
     dialect: str
 
+    def hash_inputs(self) -> str:
+        return hash_column_elements(self.args)
+
     def apply_sql_clause(self, query):
         if self.dialect == "sqlite":
             return query.group_by(*self.args)
@@ -884,24 +970,34 @@ class SQLUnion(Step):
     query1: "DatasetQuery"
     query2: "DatasetQuery"
 
+    def hash_inputs(self) -> str:
+        return hashlib.sha256(
+            bytes.fromhex(self.query1.hash()) + bytes.fromhex(self.query2.hash())
+        ).hexdigest()
+
     def apply(
         self, query_generator: QueryGenerator, temp_tables: list[str]
     ) -> StepResult:
+        left_before = len(self.query1.temp_table_names)
         q1 = self.query1.apply_steps().select().subquery()
-        temp_tables.extend(self.query1.temp_table_names)
+        temp_tables.extend(self.query1.temp_table_names[left_before:])
+        right_before = len(self.query2.temp_table_names)
         q2 = self.query2.apply_steps().select().subquery()
-        temp_tables.extend(self.query2.temp_table_names)
+        temp_tables.extend(self.query2.temp_table_names[right_before:])
 
-        columns1, columns2 = _order_columns(q1.columns, q2.columns)
+        columns1 = _drop_system_columns(q1.columns)
+        columns2 = _drop_system_columns(q2.columns)
+        columns1, columns2 = _order_columns(columns1, columns2)
 
         def q(*columns):
-            names = {c.name for c in columns}
-            col1 = [c for c in columns1 if c.name in names]
-            col2 = [c for c in columns2 if c.name in names]
-            res = sqlalchemy.select(*col1).union_all(sqlalchemy.select(*col2))
+            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))
 
-            subquery = res.subquery()
-            return sqlalchemy.select(*subquery.c).select_from(subquery)
+            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,
@@ -915,14 +1011,42 @@ class SQLJoin(Step):
     catalog: "Catalog"
     query1: "DatasetQuery"
     query2: "DatasetQuery"
-    predicates: Union[JoinPredicateType, tuple[JoinPredicateType, ...]]
+    predicates: JoinPredicateType | tuple[JoinPredicateType, ...]
     inner: bool
     full: bool
     rname: str
 
+    @staticmethod
+    def _split_db_name(name: str) -> tuple[str, str]:
+        if DEFAULT_DELIMITER in name:
+            head, tail = name.split(DEFAULT_DELIMITER, 1)
+            return head, tail
+        return name, ""
+
+    @classmethod
+    def _root_name(cls, name: str) -> str:
+        return cls._split_db_name(name)[0]
+
+    def hash_inputs(self) -> str:
+        predicates = (
+            ensure_sequence(self.predicates) if self.predicates is not None else []
+        )
+
+        parts = [
+            bytes.fromhex(self.query1.hash()),
+            bytes.fromhex(self.query2.hash()),
+            bytes.fromhex(hash_column_elements(predicates)),
+            str(self.inner).encode(),
+            str(self.full).encode(),
+            self.rname.encode("utf-8"),
+        ]
+
+        return hashlib.sha256(b"".join(parts)).hexdigest()
+
     def get_query(self, dq: "DatasetQuery", temp_tables: list[str]) -> sa.Subquery:
+        temp_tables_before = len(dq.temp_table_names)
         query = dq.apply_steps().select()
-        temp_tables.extend(dq.temp_table_names)
+        temp_tables.extend(dq.temp_table_names[temp_tables_before:])
 
         if not any(isinstance(step, (SQLJoin, SQLUnion)) for step in dq.steps):
             return query.subquery(dq.table.name)
@@ -978,22 +1102,39 @@ class SQLJoin(Step):
         q1 = self.get_query(self.query1, temp_tables)
         q2 = self.get_query(self.query2, temp_tables)
 
-        q1_columns = list(q1.c)
-        q1_column_names = {c.name for c in q1_columns}
-
-        q2_columns = []
-        for c in q2.c:
-            if c.name.startswith("sys__"):
+        q1_columns = _drop_system_columns(q1.c)
+        existing_column_names = {c.name for c in q1_columns}
+        right_columns: list[KeyedColumnElement[Any]] = []
+        right_column_names: list[str] = []
+        for column in q2.c:
+            if column.name.startswith("sys__"):
                 continue
+            right_columns.append(column)
+            right_column_names.append(column.name)
+
+        root_mapping = generate_merge_root_mapping(
+            existing_column_names,
+            right_column_names,
+            extract_root=self._root_name,
+            prefix=self.rname,
+        )
+
+        q2_columns: list[KeyedColumnElement[Any]] = []
+        for column in right_columns:
+            original_name = column.name
+            column_root, column_tail = self._split_db_name(original_name)
+            mapped_root = root_mapping[column_root]
+
+            new_name = (
+                mapped_root
+                if not column_tail
+                else DEFAULT_DELIMITER.join([mapped_root, column_tail])
+            )
+
+            if new_name != original_name:
+                column = column.label(new_name)
 
-            if c.name in q1_column_names:
-                new_name = self.rname.format(name=c.name)
-                new_name_idx = 0
-                while new_name in q1_column_names:
-                    new_name_idx += 1
-                    new_name = self.rname.format(name=f"{c.name}_{new_name_idx}")
-                c = c.label(new_name)
-            q2_columns.append(c)
+            q2_columns.append(column)
 
         res_columns = q1_columns + q2_columns
         predicates = (
@@ -1038,8 +1179,15 @@ class SQLJoin(Step):
 
 @frozen
 class SQLGroupBy(SQLClause):
-    cols: Sequence[Union[str, Function, ColumnElement]]
-    group_by: Sequence[Union[str, Function, ColumnElement]]
+    cols: Sequence[str | Function | ColumnElement]
+    group_by: Sequence[str | Function | ColumnElement]
+
+    def hash_inputs(self) -> str:
+        return hashlib.sha256(
+            bytes.fromhex(
+                hash_column_elements(self.cols) + hash_column_elements(self.group_by)
+            )
+        ).hexdigest()
 
     def apply_sql_clause(self, query) -> Select:
         if not self.cols:
@@ -1069,46 +1217,52 @@ class SQLGroupBy(SQLClause):
         return sqlalchemy.select(*unique_cols).select_from(subquery).group_by(*group_by)
 
 
-def _validate_columns(
-    left_columns: Iterable[ColumnElement], right_columns: Iterable[ColumnElement]
-) -> set[str]:
-    left_names = {c.name for c in left_columns}
-    right_names = {c.name for c in right_columns}
-
-    if left_names == right_names:
-        return left_names
-
-    missing_right = left_names - right_names
-    missing_left = right_names - left_names
-
-    def _prepare_msg_part(missing_columns: set[str], side: str) -> str:
-        return f"{', '.join(sorted(missing_columns))} only present in {side}"
-
-    msg_parts = [
-        _prepare_msg_part(missing_columns, found_side)
-        for missing_columns, found_side in zip(
-            [
-                missing_right,
-                missing_left,
-            ],
-            ["left", "right"],
-        )
-        if missing_columns
-    ]
-    msg = f"Cannot perform union. {'. '.join(msg_parts)}"
+class UnionSchemaMismatchError(ValueError):
+    """Union input columns mismatch."""
 
-    raise ValueError(msg)
+    @classmethod
+    def from_column_sets(
+        cls,
+        missing_left: set[str],
+        missing_right: set[str],
+    ) -> "UnionSchemaMismatchError":
+        def _describe(cols: set[str], side: str) -> str:
+            return f"{', '.join(sorted(cols))} only present in {side}"
+
+        parts = []
+        if missing_left:
+            parts.append(_describe(missing_left, "left"))
+        if missing_right:
+            parts.append(_describe(missing_right, "right"))
+
+        return cls(f"Cannot perform union. {'. '.join(parts)}")
 
 
 def _order_columns(
     left_columns: Iterable[ColumnElement], right_columns: Iterable[ColumnElement]
 ) -> list[list[ColumnElement]]:
-    column_order = _validate_columns(left_columns, right_columns)
+    left_names = [c.name for c in left_columns]
+    right_names = [c.name for c in right_columns]
+
+    # validate
+    if sorted(left_names) != sorted(right_names):
+        left_names_set = set(left_names)
+        right_names_set = set(right_names)
+        raise UnionSchemaMismatchError.from_column_sets(
+            left_names_set - right_names_set,
+            right_names_set - left_names_set,
+        )
+
+    # Order columns to match left_names order
     column_dicts = [
         {c.name: c for c in columns} for columns in [left_columns, right_columns]
     ]
 
-    return [[d[n] for n in column_order] for d in column_dicts]
+    return [[d[n] for n in left_names] 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
@@ -1124,40 +1278,42 @@ class DatasetQuery:
     def __init__(
         self,
         name: str,
-        version: Optional[str] = None,
-        project_name: Optional[str] = None,
-        namespace_name: Optional[str] = None,
-        catalog: Optional["Catalog"] = None,
-        session: Optional[Session] = None,
+        version: str | None = None,
+        project_name: str | None = None,
+        namespace_name: str | None = None,
+        catalog: "Catalog | None" = None,
+        session: Session | None = None,
         in_memory: bool = False,
         update: bool = False,
     ) -> None:
         self.session = Session.get(session, catalog=catalog, in_memory=in_memory)
         self.catalog = catalog or self.session.catalog
         self.steps: list[Step] = []
-        self._chunk_index: Optional[int] = None
-        self._chunk_total: Optional[int] = None
+        self._chunk_index: int | None = None
+        self._chunk_total: int | None = None
         self.temp_table_names: list[str] = []
         self.dependencies: set[DatasetDependencyType] = set()
         self.table = self.get_table()
-        self.starting_step: Optional[QueryStep] = None
-        self.name: Optional[str] = None
-        self.version: Optional[str] = None
-        self.feature_schema: Optional[dict] = None
-        self.column_types: Optional[dict[str, Any]] = None
+        self.starting_step: QueryStep | None = None
+        self.name: str | None = None
+        self.version: str | None = None
+        self.feature_schema: dict | None = None
+        self.column_types: dict[str, Any] | None = None
         self.before_steps: list[Callable] = []
-        self.listing_fn: Optional[Callable] = None
+        self.listing_fn: Callable | None = None
         self.update = update
 
-        self.list_ds_name: Optional[str] = None
+        self.list_ds_name: str | None = None
 
         self.name = name
         self.dialect = self.catalog.warehouse.db.dialect
         if version:
             self.version = version
 
-        namespace_name = namespace_name or self.catalog.metastore.default_namespace_name
-        project_name = project_name or self.catalog.metastore.default_project_name
+        if namespace_name is None:
+            namespace_name = self.catalog.metastore.default_namespace_name
+        if project_name is None:
+            project_name = self.catalog.metastore.default_project_name
 
         if is_listing_dataset(name) and not version:
             # not setting query step yet as listing dataset might not exist at
@@ -1194,12 +1350,26 @@ class DatasetQuery:
     def __or__(self, other):
         return self.union(other)
 
+    def hash(self) -> str:
+        """
+        Calculates hash of this class taking into account hash of starting step
+        and hashes of each following steps. Ordering is important.
+        """
+        hasher = hashlib.sha256()
+        if self.starting_step:
+            hasher.update(self.starting_step.hash().encode("utf-8"))
+        else:
+            assert self.list_ds_name
+            hasher.update(self.list_ds_name.encode("utf-8"))
+
+        for step in self.steps:
+            hasher.update(step.hash().encode("utf-8"))
+
+        return hasher.hexdigest()
+
     @staticmethod
     def get_table() -> "TableClause":
-        table_name = "".join(
-            random.choice(string.ascii_letters)  # noqa: S311
-            for _ in range(16)
-        )
+        table_name = "".join(secrets.choice(string.ascii_letters) for _ in range(16))
         return sqlalchemy.table(table_name)
 
     @property
@@ -1216,7 +1386,7 @@ class DatasetQuery:
         """
         return self.name is not None and self.version is not None
 
-    def c(self, column: Union[C, str]) -> "ColumnClause[Any]":
+    def c(self, column: C | str) -> "ColumnClause[Any]":
         col: sqlalchemy.ColumnClause = (
             sqlalchemy.column(column)
             if isinstance(column, str)
@@ -1311,6 +1481,7 @@ class DatasetQuery:
         # This is needed to always use a new connection with all metastore and warehouse
         # implementations, as errors may close or render unusable the existing
         # connections.
+        assert len(self.temp_table_names) == len(set(self.temp_table_names))
         with self.catalog.metastore.clone(use_new_connection=True) as metastore:
             metastore.cleanup_tables(self.temp_table_names)
         with self.catalog.warehouse.clone(use_new_connection=True) as warehouse:
@@ -1325,7 +1496,7 @@ class DatasetQuery:
             return list(result)
 
     def to_db_records(self) -> list[dict[str, Any]]:
-        return self.db_results(lambda cols, row: dict(zip(cols, row)))
+        return self.db_results(lambda cols, row: dict(zip(cols, row, strict=False)))
 
     @contextlib.contextmanager
     def as_iterable(self, **kwargs) -> Iterator[ResultIter]:
@@ -1364,7 +1535,7 @@ class DatasetQuery:
                         yield from rows
 
             async def get_params(row: Sequence) -> tuple:
-                row_dict = RowDict(zip(query_fields, row))
+                row_dict = RowDict(zip(query_fields, row, strict=False))
                 return tuple(  # noqa: C409
                     [
                         await p.get_value_async(
@@ -1381,10 +1552,6 @@ class DatasetQuery:
         finally:
             self.cleanup()
 
-    def shuffle(self) -> "Self":
-        # ToDo: implement shaffle based on seed and/or generating random column
-        return self.order_by(C.sys__rand)
-
     def sample(self, n) -> "Self":
         """
         Return a random sample from the dataset.
@@ -1404,6 +1571,7 @@ class DatasetQuery:
         obj.steps = obj.steps.copy()
         if new_table:
             obj.table = self.get_table()
+        obj.temp_table_names = []
         return obj
 
     @detach
@@ -1584,10 +1752,10 @@ class DatasetQuery:
     def join(
         self,
         dataset_query: "DatasetQuery",
-        predicates: Union[JoinPredicateType, Sequence[JoinPredicateType]],
+        predicates: JoinPredicateType | Sequence[JoinPredicateType],
         inner=False,
         full=False,
-        rname="{name}_right",
+        rname="right_",
     ) -> "Self":
         left = self.clone(new_table=False)
         if self.table.name == dataset_query.table.name:
@@ -1626,12 +1794,17 @@ class DatasetQuery:
     def add_signals(
         self,
         udf: "UDFAdapter",
-        parallel: Optional[int] = None,
-        workers: Union[bool, int] = False,
-        min_task_size: Optional[int] = None,
-        partition_by: Optional[PartitionByType] = None,
+        partition_by: PartitionByType | None = None,
+        # Parameters from Settings
         cache: bool = False,
-        batch_rows: Optional[int] = None,
+        parallel: int | None = None,
+        workers: bool | int = False,
+        min_task_size: int | None = None,
+        batch_size: int | None = None,
+        # Parameters are unused, kept only to match the signature of Settings.to_dict
+        prefetch: int | None = None,
+        namespace: str | None = None,
+        project: str | None = None,
     ) -> "Self":
         """
         Adds one or more signals based on the results from the provided UDF.
@@ -1657,7 +1830,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
-                batch_rows=batch_rows,
+                batch_size=batch_size,
             )
         )
         return query
@@ -1672,14 +1845,17 @@ class DatasetQuery:
     def generate(
         self,
         udf: "UDFAdapter",
-        parallel: Optional[int] = None,
-        workers: Union[bool, int] = False,
-        min_task_size: Optional[int] = None,
-        partition_by: Optional[PartitionByType] = None,
-        namespace: Optional[str] = None,
-        project: Optional[str] = None,
+        partition_by: PartitionByType | None = None,
+        # Parameters from Settings
         cache: bool = False,
-        batch_rows: Optional[int] = None,
+        parallel: int | None = None,
+        workers: bool | int = False,
+        min_task_size: int | None = None,
+        batch_size: int | None = None,
+        # Parameters are unused, kept only to match the signature of Settings.to_dict:
+        prefetch: int | None = None,
+        namespace: str | None = None,
+        project: str | None = None,
     ) -> "Self":
         query = self.clone()
         steps = query.steps
@@ -1692,7 +1868,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
-                batch_rows=batch_rows,
+                batch_size=batch_size,
             )
         )
         return query
@@ -1735,26 +1911,30 @@ class DatasetQuery:
 
     def exec(self) -> "Self":
         """Execute the query."""
+        query = self.clone()
         try:
-            query = self.clone()
             query.apply_steps()
         finally:
-            self.cleanup()
+            query.cleanup()
         return query
 
     def save(
         self,
-        name: Optional[str] = None,
-        version: Optional[str] = None,
-        project: Optional[Project] = None,
-        feature_schema: Optional[dict] = None,
-        dependencies: Optional[list[DatasetDependency]] = None,
-        description: Optional[str] = None,
-        attrs: Optional[list[str]] = None,
-        update_version: Optional[str] = "patch",
+        name: str | None = None,
+        version: str | None = None,
+        project: Project | None = None,
+        feature_schema: dict | None = None,
+        dependencies: list[DatasetDependency] | None = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
+        update_version: str | None = "patch",
         **kwargs,
     ) -> "Self":
         """Save the query as a dataset."""
+        # Get job from session to link dataset version to job
+        job = self.session.get_or_create_job()
+        job_id = job.id
+
         project = project or self.catalog.metastore.default_project
         try:
             if (
@@ -1797,14 +1977,11 @@ class DatasetQuery:
                 description=description,
                 attrs=attrs,
                 update_version=update_version,
+                job_id=job_id,
                 **kwargs,
             )
             version = version or dataset.latest_version
 
-            self.session.add_dataset_version(
-                dataset=dataset, version=version, listing=kwargs.get("listing", False)
-            )
-
             dr = self.catalog.warehouse.dataset_rows(dataset)
 
             self.catalog.warehouse.copy_table(dr.get_table(), query.select())
@@ -1814,6 +1991,11 @@ class DatasetQuery:
             )
             self.catalog.update_dataset_version_with_warehouse_info(dataset, version)
 
+            # Link this dataset version to the job that created it
+            self.catalog.metastore.link_dataset_version_to_job(
+                dataset.get_version(version).id, job_id, is_creator=True
+            )
+
             if dependencies:
                 # overriding dependencies
                 self.dependencies = set()
@@ -1845,5 +2027,5 @@ class DatasetQuery:
         return isinstance(self.last_step, SQLOrderBy)
 
     @property
-    def last_step(self) -> Optional[Step]:
+    def last_step(self) -> Step | None:
         return self.steps[-1] if self.steps else None