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

datachain/func/conditional.py

@@ -1,5 +1,3 @@
-from typing import Optional, Union
-
 from sqlalchemy import ColumnElement
 from sqlalchemy import and_ as sql_and
 from sqlalchemy import case as sql_case
@@ -12,10 +10,10 @@ from datachain.sql.functions import conditional
 
 from .func import Func
 
-CaseT = Union[int, float, complex, bool, str, Func, ColumnElement]
+CaseT = int | float | complex | bool | str | Func | ColumnElement
 
 
-def greatest(*args: Union[str, Column, Func, float]) -> Func:
+def greatest(*args: str | Column | Func | float) -> Func:
     """
     Returns the greatest (largest) value from the given input values.
 
@@ -56,7 +54,7 @@ def greatest(*args: Union[str, Column, Func, float]) -> Func:
     )
 
 
-def least(*args: Union[str, Column, Func, float]) -> Func:
+def least(*args: str | Column | Func | float) -> Func:
     """
     Returns the least (smallest) value from the given input values.
 
@@ -94,7 +92,7 @@ def least(*args: Union[str, Column, Func, float]) -> Func:
 
 
 def case(
-    *args: tuple[Union[ColumnElement, Func, bool], CaseT], else_: Optional[CaseT] = None
+    *args: tuple[ColumnElement | Func | bool, CaseT], else_: CaseT | None = None
 ) -> Func:
     """
     Returns a case expression that evaluates a list of conditions and returns
@@ -163,9 +161,7 @@ def case(
     return Func("case", inner=sql_case, cols=args, kwargs=kwargs, result_type=type_)
 
 
-def ifelse(
-    condition: Union[ColumnElement, Func], if_val: CaseT, else_val: CaseT
-) -> Func:
+def ifelse(condition: ColumnElement | Func, if_val: CaseT, else_val: CaseT) -> Func:
     """
     Returns an if-else expression that evaluates a condition and returns one
     of two values based on the result. Values can be Python primitives
@@ -193,7 +189,7 @@ def ifelse(
     return case((condition, if_val), else_=else_val)
 
 
-def isnone(col: Union[str, ColumnElement]) -> Func:
+def isnone(col: str | ColumnElement) -> Func:
     """
     Returns a function that checks if the column value is `None` (NULL in DB).
 
@@ -221,7 +217,7 @@ def isnone(col: Union[str, ColumnElement]) -> Func:
     return case((col.is_(None) if col is not None else True, True), else_=False)
 
 
-def or_(*args: Union[ColumnElement, Func]) -> Func:
+def or_(*args: ColumnElement | Func) -> Func:
     """
     Returns the function that produces conjunction of expressions joined by OR
     logical operator.
@@ -256,7 +252,7 @@ def or_(*args: Union[ColumnElement, Func]) -> Func:
     return Func("or", inner=sql_or, cols=cols, args=func_args, result_type=bool)
 
 
-def and_(*args: Union[ColumnElement, Func]) -> Func:
+def and_(*args: ColumnElement | Func) -> Func:
     """
     Returns the function that produces conjunction of expressions joined by AND
     logical operator.
@@ -291,7 +287,7 @@ def and_(*args: Union[ColumnElement, Func]) -> Func:
     return Func("and", inner=sql_and, cols=cols, args=func_args, result_type=bool)
 
 
-def not_(arg: Union[ColumnElement, Func]) -> Func:
+def not_(arg: ColumnElement | Func) -> Func:
     """
     Returns the function that produces NOT of the given expressions.
 

datachain/func/func.py

@@ -1,12 +1,13 @@
 import inspect
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union, get_args, get_origin
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any, Union, get_args, get_origin
 
 from sqlalchemy import BindParameter, Case, ColumnElement, Integer, cast, desc
 from sqlalchemy.sql import func as sa_func
 
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.sql_to_python import sql_to_python
+from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.query.schema import Column, ColumnMeta
 from datachain.sql.functions import numeric
@@ -22,26 +23,29 @@ if TYPE_CHECKING:
     from .window import Window
 
 
-ColT = Union[str, Column, ColumnElement, "Func", tuple]
+ColT = Union[str, tuple, Column, ColumnElement, "Func"]
 
 
 class Func(Function):  # noqa: PLW1641
     """Represents a function to be applied to a column in a SQL query."""
 
+    cols: Sequence[ColT]
+    args: Sequence[Any]
+
     def __init__(
         self,
         name: str,
         inner: Callable,
-        cols: Optional[Sequence[ColT]] = None,
-        args: Optional[Sequence[Any]] = None,
-        kwargs: Optional[dict[str, Any]] = None,
-        result_type: Optional["DataType"] = None,
-        type_from_args: Optional[Callable[..., "DataType"]] = None,
+        cols: Sequence[ColT] | None = None,
+        args: Sequence[Any] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        result_type: "DataType | None" = None,
+        type_from_args: Callable[..., "DataType"] | None = None,
         is_array: bool = False,
         from_array: bool = False,
         is_window: bool = False,
-        window: Optional["Window"] = None,
-        label: Optional[str] = None,
+        window: "Window | None" = None,
+        label: str | None = None,
     ) -> None:
         self.name = name
         self.inner = inner
@@ -95,7 +99,7 @@ class Func(Function):  # noqa: PLW1641
             else []
         )
 
-    def _db_col_type(self, signals_schema: "SignalSchema") -> Optional["DataType"]:
+    def _db_col_type(self, signals_schema: "SignalSchema") -> "DataType | None":
         if not self._db_cols:
             return None
 
@@ -125,51 +129,51 @@ class Func(Function):  # noqa: PLW1641
 
         return list[col_type] if self.is_array else col_type  # type: ignore[valid-type]
 
-    def __add__(self, other: Union[ColT, float]) -> "Func":
+    def __add__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("add", lambda a: a + other, [self])
         return Func("add", lambda a1, a2: a1 + a2, [self, other])
 
-    def __radd__(self, other: Union[ColT, float]) -> "Func":
+    def __radd__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("add", lambda a: other + a, [self])
         return Func("add", lambda a1, a2: a1 + a2, [other, self])
 
-    def __sub__(self, other: Union[ColT, float]) -> "Func":
+    def __sub__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("sub", lambda a: a - other, [self])
         return Func("sub", lambda a1, a2: a1 - a2, [self, other])
 
-    def __rsub__(self, other: Union[ColT, float]) -> "Func":
+    def __rsub__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("sub", lambda a: other - a, [self])
         return Func("sub", lambda a1, a2: a1 - a2, [other, self])
 
-    def __mul__(self, other: Union[ColT, float]) -> "Func":
+    def __mul__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mul", lambda a: a * other, [self])
         return Func("mul", lambda a1, a2: a1 * a2, [self, other])
 
-    def __rmul__(self, other: Union[ColT, float]) -> "Func":
+    def __rmul__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mul", lambda a: other * a, [self])
         return Func("mul", lambda a1, a2: a1 * a2, [other, self])
 
-    def __truediv__(self, other: Union[ColT, float]) -> "Func":
+    def __truediv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("div", lambda a: _truediv(a, other), [self], result_type=float)
         return Func(
             "div", lambda a1, a2: _truediv(a1, a2), [self, other], result_type=float
         )
 
-    def __rtruediv__(self, other: Union[ColT, float]) -> "Func":
+    def __rtruediv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("div", lambda a: _truediv(other, a), [self], result_type=float)
         return Func(
             "div", lambda a1, a2: _truediv(a1, a2), [other, self], result_type=float
         )
 
-    def __floordiv__(self, other: Union[ColT, float]) -> "Func":
+    def __floordiv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "floordiv", lambda a: _floordiv(a, other), [self], result_type=int
@@ -178,7 +182,7 @@ class Func(Function):  # noqa: PLW1641
             "floordiv", lambda a1, a2: _floordiv(a1, a2), [self, other], result_type=int
         )
 
-    def __rfloordiv__(self, other: Union[ColT, float]) -> "Func":
+    def __rfloordiv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "floordiv", lambda a: _floordiv(other, a), [self], result_type=int
@@ -187,17 +191,17 @@ class Func(Function):  # noqa: PLW1641
             "floordiv", lambda a1, a2: _floordiv(a1, a2), [other, self], result_type=int
         )
 
-    def __mod__(self, other: Union[ColT, float]) -> "Func":
+    def __mod__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mod", lambda a: a % other, [self], result_type=int)
         return Func("mod", lambda a1, a2: a1 % a2, [self, other], result_type=int)
 
-    def __rmod__(self, other: Union[ColT, float]) -> "Func":
+    def __rmod__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mod", lambda a: other % a, [self], result_type=int)
         return Func("mod", lambda a1, a2: a1 % a2, [other, self], result_type=int)
 
-    def __and__(self, other: Union[ColT, float]) -> "Func":
+    def __and__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "and", lambda a: numeric.bit_and(a, other), [self], result_type=int
@@ -209,7 +213,7 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __rand__(self, other: Union[ColT, float]) -> "Func":
+    def __rand__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "and", lambda a: numeric.bit_and(other, a), [self], result_type=int
@@ -221,7 +225,7 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __or__(self, other: Union[ColT, float]) -> "Func":
+    def __or__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "or", lambda a: numeric.bit_or(a, other), [self], result_type=int
@@ -230,7 +234,7 @@ class Func(Function):  # noqa: PLW1641
             "or", lambda a1, a2: numeric.bit_or(a1, a2), [self, other], result_type=int
         )
 
-    def __ror__(self, other: Union[ColT, float]) -> "Func":
+    def __ror__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "or", lambda a: numeric.bit_or(other, a), [self], result_type=int
@@ -239,7 +243,7 @@ class Func(Function):  # noqa: PLW1641
             "or", lambda a1, a2: numeric.bit_or(a1, a2), [other, self], result_type=int
         )
 
-    def __xor__(self, other: Union[ColT, float]) -> "Func":
+    def __xor__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "xor", lambda a: numeric.bit_xor(a, other), [self], result_type=int
@@ -251,7 +255,7 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __rxor__(self, other: Union[ColT, float]) -> "Func":
+    def __rxor__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "xor", lambda a: numeric.bit_xor(other, a), [self], result_type=int
@@ -263,7 +267,7 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __rshift__(self, other: Union[ColT, float]) -> "Func":
+    def __rshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "rshift",
@@ -278,7 +282,7 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __rrshift__(self, other: Union[ColT, float]) -> "Func":
+    def __rrshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "rshift",
@@ -293,7 +297,7 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __lshift__(self, other: Union[ColT, float]) -> "Func":
+    def __lshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "lshift",
@@ -308,7 +312,7 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __rlshift__(self, other: Union[ColT, float]) -> "Func":
+    def __rlshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "lshift",
@@ -323,12 +327,12 @@ class Func(Function):  # noqa: PLW1641
             result_type=int,
         )
 
-    def __lt__(self, other: Union[ColT, float]) -> "Func":
+    def __lt__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("lt", lambda a: a < other, [self], result_type=bool)
         return Func("lt", lambda a1, a2: a1 < a2, [self, other], result_type=bool)
 
-    def __le__(self, other: Union[ColT, float]) -> "Func":
+    def __le__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("le", lambda a: a <= other, [self], result_type=bool)
         return Func("le", lambda a1, a2: a1 <= a2, [self, other], result_type=bool)
@@ -343,12 +347,12 @@ class Func(Function):  # noqa: PLW1641
             return Func("ne", lambda a: a != other, [self], result_type=bool)
         return Func("ne", lambda a1, a2: a1 != a2, [self, other], result_type=bool)
 
-    def __gt__(self, other: Union[ColT, float]) -> "Func":
+    def __gt__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("gt", lambda a: a > other, [self], result_type=bool)
         return Func("gt", lambda a1, a2: a1 > a2, [self, other], result_type=bool)
 
-    def __ge__(self, other: Union[ColT, float]) -> "Func":
+    def __ge__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("ge", lambda a: a >= other, [self], result_type=bool)
         return Func("ge", lambda a1, a2: a1 >= a2, [self, other], result_type=bool)
@@ -369,7 +373,7 @@ class Func(Function):  # noqa: PLW1641
             label,
         )
 
-    def get_col_name(self, label: Optional[str] = None) -> str:
+    def get_col_name(self, label: str | None = None) -> str:
         if label:
             return label
         if self.col_label:
@@ -384,7 +388,7 @@ class Func(Function):  # noqa: PLW1641
         return self.name
 
     def get_result_type(
-        self, signals_schema: Optional["SignalSchema"] = None
+        self, signals_schema: "SignalSchema | None" = None
     ) -> "DataType":
         if self.result_type:
             return self.result_type
@@ -408,10 +412,24 @@ class Func(Function):  # noqa: PLW1641
 
     def get_column(
         self,
-        signals_schema: Optional["SignalSchema"] = None,
-        label: Optional[str] = None,
-        table: Optional["TableClause"] = None,
+        signals_schema: "SignalSchema | None" = None,
+        label: str | None = None,
+        table: "TableClause | None" = None,
     ) -> Column:
+        # Guard against using complex (pydantic) object columns in SQL funcs
+        if signals_schema and self._db_cols:
+            for arg in self._db_cols:
+                # _db_cols normalizes known columns to strings; skip non-string args
+                if not isinstance(arg, str):
+                    continue
+                t_with_sub = signals_schema.get_column_type(arg, with_subtree=True)
+                if ModelStore.is_pydantic(t_with_sub):
+                    raise DataChainParamsError(
+                        f"Function {self.name} doesn't support complex object "
+                        f"columns like '{arg}'. Use a leaf field (e.g., "
+                        f"'{arg}.path') or use UDFs to operate on complex objects."
+                    )
+
         col_type = self.get_result_type(signals_schema)
         sql_type = python_to_sql(col_type)
 
@@ -431,6 +449,7 @@ class Func(Function):  # noqa: PLW1641
             return col
 
         cols = [get_col(col) for col in self._db_cols]
+
         kwargs = {k: get_col(v, string_as_literal=True) for k, v in self.kwargs.items()}
         func_col = self.inner(*cols, *self.args, **kwargs)
 
@@ -467,9 +486,8 @@ def get_db_col_type(signals_schema: "SignalSchema", col: ColT) -> "DataType":
     if isinstance(col, ColumnElement) and not hasattr(col, "name"):
         return sql_to_python(col)
 
-    return signals_schema.get_column_type(
-        col.name if isinstance(col, ColumnElement) else col  # type: ignore[arg-type]
-    )
+    name = col.name if isinstance(col, ColumnElement) else col  # type: ignore[assignment]
+    return signals_schema.get_column_type(name)  # type: ignore[arg-type]
 
 
 def _truediv(a, b):

datachain/func/numeric.py

@@ -1,12 +1,10 @@
-from typing import Union
-
 from datachain.query.schema import Column
 from datachain.sql.functions import numeric
 
 from .func import Func
 
 
-def bit_and(*args: Union[str, Column, Func, int]) -> Func:
+def bit_and(*args: str | Column | Func | int) -> Func:
     """
     Returns a function that computes the bitwise AND operation between two values.
 
@@ -51,7 +49,7 @@ def bit_and(*args: Union[str, Column, Func, int]) -> Func:
     )
 
 
-def bit_or(*args: Union[str, Column, Func, int]) -> Func:
+def bit_or(*args: str | Column | Func | int) -> Func:
     """
     Returns a function that computes the bitwise OR operation between two values.
 
@@ -96,7 +94,7 @@ def bit_or(*args: Union[str, Column, Func, int]) -> Func:
     )
 
 
-def bit_xor(*args: Union[str, Column, Func, int]) -> Func:
+def bit_xor(*args: str | Column | Func | int) -> Func:
     """
     Returns a function that computes the bitwise XOR operation between two values.
 
@@ -141,7 +139,7 @@ def bit_xor(*args: Union[str, Column, Func, int]) -> Func:
     )
 
 
-def int_hash_64(col: Union[str, Column, Func, int]) -> Func:
+def int_hash_64(col: str | Column | Func | int) -> Func:
     """
     Returns a function that computes the 64-bit hash of an integer.
 
@@ -177,7 +175,7 @@ def int_hash_64(col: Union[str, Column, Func, int]) -> Func:
     )
 
 
-def bit_hamming_distance(*args: Union[str, Column, Func, int]) -> Func:
+def bit_hamming_distance(*args: str | Column | Func | int) -> Func:
     """
     Returns a function that computes the Hamming distance between two integers.
 

datachain/func/string.py

@@ -1,4 +1,4 @@
-from typing import Optional, get_origin
+from typing import get_origin
 
 from sqlalchemy import literal
 
@@ -44,7 +44,7 @@ def length(col: ColT) -> Func:
     return Func("length", inner=string.length, cols=[col], result_type=int)
 
 
-def split(col: ColT, sep: str, limit: Optional[int] = None) -> Func:
+def split(col: ColT, sep: str, limit: int | None = None) -> Func:
     """
     Takes a column and split character and returns an array of the parts.
 

datachain/hash_utils.py

@@ -0,0 +1,123 @@
+import hashlib
+import inspect
+import textwrap
+from collections.abc import Sequence
+from typing import TypeAlias, TypeVar
+
+from sqlalchemy.sql.elements import ClauseElement, ColumnElement
+
+from datachain import json
+
+T = TypeVar("T", bound=ColumnElement)
+ColumnLike: TypeAlias = str | T
+
+
+def _serialize_value(val):  # noqa: PLR0911
+    """Helper to serialize arbitrary values recursively."""
+    if val is None:
+        return None
+    if isinstance(val, (str, int, float, bool)):
+        return val
+    if isinstance(val, ClauseElement):
+        return serialize_column_element(val)
+    if isinstance(val, dict):
+        # Sort dict keys for deterministic serialization
+        return {k: _serialize_value(v) for k, v in sorted(val.items())}
+    if isinstance(val, (list, tuple)):
+        return [_serialize_value(v) for v in val]
+    if callable(val):
+        return val.__name__ if hasattr(val, "__name__") else str(val)
+    return str(val)
+
+
+def serialize_column_element(expr: str | ColumnElement) -> dict:
+    """
+    Recursively serialize a SQLAlchemy ColumnElement into a deterministic structure.
+    Uses SQLAlchemy's _traverse_internals to automatically handle all expression types.
+    """
+    from sqlalchemy.sql.elements import BindParameter
+
+    # Special case: BindParameter has non-deterministic 'key' attribute, only use value
+    if isinstance(expr, BindParameter):
+        return {"type": "bind", "value": _serialize_value(expr.value)}
+
+    # Generic handling for all ClauseElement types using SQLAlchemy's internals
+    if isinstance(expr, ClauseElement):
+        # All standard SQLAlchemy types have _traverse_internals
+        if hasattr(expr, "_traverse_internals"):
+            result = {"type": expr.__class__.__name__}
+            for attr_name, _ in expr._traverse_internals:
+                # Skip 'table' attribute - table names can be auto-generated/random
+                # and are not semantically important for hashing
+                if attr_name == "table":
+                    continue
+                if hasattr(expr, attr_name):
+                    val = getattr(expr, attr_name)
+                    result[attr_name] = _serialize_value(val)
+            return result
+        # Rare case: custom user-defined ClauseElement without _traverse_internals
+        # We don't know its structure, so just stringify it
+        return {"type": expr.__class__.__name__, "repr": str(expr)}
+
+    # Absolute fallback: stringify completely unknown types
+    return {"type": "other", "repr": str(expr)}
+
+
+def hash_column_elements(columns: ColumnLike | Sequence[ColumnLike]) -> str:
+    """
+    Hash a list of ColumnElements deterministically, dialect agnostic.
+    Only accepts ordered iterables (like list or tuple).
+    """
+    # Handle case where a single ColumnElement is passed instead of a sequence
+    if isinstance(columns, (ColumnElement, str)):
+        columns = (columns,)
+
+    serialized = [serialize_column_element(c) for c in columns]
+    json_str = json.dumps(
+        serialized, sort_keys=True, separators=(", ", ": ")
+    )  # stable JSON
+    return hashlib.sha256(json_str.encode("utf-8")).hexdigest()
+
+
+def hash_callable(func):
+    """
+    Calculate a hash from a callable.
+    Rules:
+    - Named functions (def) → use source code for stable, cross-version hashing
+    - Lambdas → use bytecode (deterministic in same Python runtime)
+    """
+    if not callable(func):
+        raise TypeError("Expected a callable")
+
+    # Determine if it is a lambda
+    is_lambda = func.__name__ == "<lambda>"
+
+    if not is_lambda:
+        # Try to get exact source of named function
+        try:
+            lines, _ = inspect.getsourcelines(func)
+            payload = textwrap.dedent("".join(lines)).strip()
+        except (OSError, TypeError):
+            # Fallback: bytecode if source not available
+            payload = func.__code__.co_code
+    else:
+        # For lambdas, fall back directly to bytecode
+        payload = func.__code__.co_code
+
+    # Normalize annotations
+    annotations = {
+        k: getattr(v, "__name__", str(v)) for k, v in func.__annotations__.items()
+    }
+
+    # Extras to distinguish functions with same code but different metadata
+    extras = {
+        "name": func.__name__,
+        "defaults": func.__defaults__,
+        "annotations": annotations,
+    }
+
+    # Compute SHA256
+    h = hashlib.sha256()
+    h.update(str(payload).encode() if isinstance(payload, str) else payload)
+    h.update(str(extras).encode())
+    return h.hexdigest()

datachain/job.py

@@ -1,8 +1,9 @@
-import json
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, TypeVar
+
+from datachain import json
 
 J = TypeVar("J", bound="Job")
 
@@ -18,27 +19,29 @@ class Job:
     workers: int
     params: dict[str, str]
     metrics: dict[str, Any]
-    finished_at: Optional[datetime] = None
-    python_version: Optional[str] = None
+    finished_at: datetime | None = None
+    python_version: str | None = None
     error_message: str = ""
     error_stack: str = ""
+    parent_job_id: str | None = None
 
     @classmethod
     def parse(
         cls,
-        id: Union[str, uuid.UUID],
+        id: str | uuid.UUID,
         name: str,
         status: int,
         created_at: datetime,
-        finished_at: Optional[datetime],
+        finished_at: datetime | None,
         query: str,
         query_type: int,
         workers: int,
-        python_version: Optional[str],
+        python_version: str | None,
         error_message: str,
         error_stack: str,
         params: str,
         metrics: str,
+        parent_job_id: str | None,
     ) -> "Job":
         return cls(
             str(id),
@@ -54,4 +57,5 @@ class Job:
             python_version,
             error_message,
             error_stack,
+            str(parent_job_id) if parent_job_id else None,
         )