datachain 0.34.6__py3-none-any.whl → 0.35.0__py3-none-any.whl

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

@@ -3,101 +3,74 @@ import inspect
 import json
 import textwrap
 from collections.abc import Sequence
-from typing import TypeVar, Union
-
-from sqlalchemy.sql.elements import (
-    BinaryExpression,
-    BindParameter,
-    ColumnElement,
-    Label,
-    Over,
-    UnaryExpression,
-)
-from sqlalchemy.sql.functions import Function
-
-T = TypeVar("T", bound=ColumnElement)
-ColumnLike = Union[str, T]
+from typing import TypeAlias, TypeVar
 
+from sqlalchemy.sql.elements import ClauseElement, ColumnElement
 
-def serialize_column_element(expr: Union[str, ColumnElement]) -> dict:  # noqa: PLR0911
+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
 
-    # Binary operations: col > 5, col1 + col2, etc.
-    if isinstance(expr, BinaryExpression):
-        op = (
-            expr.operator.__name__
-            if hasattr(expr.operator, "__name__")
-            else str(expr.operator)
-        )
-        return {
-            "type": "binary",
-            "op": op,
-            "left": serialize_column_element(expr.left),
-            "right": serialize_column_element(expr.right),
-        }
-
-    # Unary operations: -col, NOT col, etc.
-    if isinstance(expr, UnaryExpression):
-        op = (
-            expr.operator.__name__
-            if expr.operator is not None and hasattr(expr.operator, "__name__")
-            else str(expr.operator)
-        )
-
-        return {
-            "type": "unary",
-            "op": op,
-            "element": serialize_column_element(expr.element),  # type: ignore[arg-type]
-        }
-
-    # Function calls: func.lower(col), func.count(col), etc.
-    if isinstance(expr, Function):
-        return {
-            "type": "function",
-            "name": expr.name,
-            "clauses": [serialize_column_element(c) for c in expr.clauses],
-        }
-
-    # Window functions: func.row_number().over(partition_by=..., order_by=...)
-    if isinstance(expr, Over):
-        return {
-            "type": "window",
-            "function": serialize_column_element(expr.element),
-            "partition_by": [
-                serialize_column_element(p) for p in getattr(expr, "partition_by", [])
-            ],
-            "order_by": [
-                serialize_column_element(o) for o in getattr(expr, "order_by", [])
-            ],
-        }
-
-    # Labeled expressions: col.label("alias")
-    if isinstance(expr, Label):
-        return {
-            "type": "label",
-            "name": expr.name,
-            "element": serialize_column_element(expr.element),
-        }
-
-    # Bound values (constants)
+    # Special case: BindParameter has non-deterministic 'key' attribute, only use value
     if isinstance(expr, BindParameter):
-        return {"type": "bind", "value": expr.value}
-
-    # Plain columns
-    if hasattr(expr, "name"):
-        return {"type": "column", "name": expr.name}
-
-    # Fallback: stringify unknown nodes
+        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: Sequence[ColumnLike]) -> str:
+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)  # stable JSON
     return hashlib.sha256(json_str.encode("utf-8")).hexdigest()

datachain/job.py

@@ -2,7 +2,7 @@ import json
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, TypeVar
 
 J = TypeVar("J", bound="Job")
 
@@ -18,29 +18,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: Optional[str] = None
+    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: Optional[str],
+        parent_job_id: str | None,
     ) -> "Job":
         return cls(
             str(id),

datachain/lib/arrow.py

@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from itertools import islice
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
 import ujson as json
@@ -44,10 +44,10 @@ class ArrowGenerator(Generator):
 
     def __init__(
         self,
-        input_schema: Optional["pa.Schema"] = None,
-        output_schema: Optional[type["BaseModel"]] = None,
+        input_schema: pa.Schema | None = None,
+        output_schema: type["BaseModel"] | None = None,
         source: bool = True,
-        nrows: Optional[int] = None,
+        nrows: int | None = None,
         **kwargs,
     ):
         """
@@ -112,7 +112,7 @@ class ArrowGenerator(Generator):
         record: dict[str, Any],
         file: File,
         index: int,
-        hf_schema: Optional[tuple["Features", dict[str, "DataType"]]],
+        hf_schema: tuple["Features", dict[str, "DataType"]] | None,
         use_datachain_schema: bool,
     ):
         if use_datachain_schema and self.output_schema:
@@ -141,7 +141,7 @@ class ArrowGenerator(Generator):
     def _process_non_datachain_record(
         self,
         record: dict[str, Any],
-        hf_schema: Optional[tuple["Features", dict[str, "DataType"]]],
+        hf_schema: tuple["Features", dict[str, "DataType"]] | None,
     ):
         vals = list(record.values())
         if not self.output_schema:
@@ -149,7 +149,9 @@ class ArrowGenerator(Generator):
 
         fields = self.output_schema.model_fields
         vals_dict = {}
-        for i, ((field, field_info), val) in enumerate(zip(fields.items(), vals)):
+        for i, ((field, field_info), val) in enumerate(
+            zip(fields.items(), vals, strict=False)
+        ):
             anno = field_info.annotation
             if hf_schema:
                 from datachain.lib.hf import convert_feature
@@ -180,7 +182,7 @@ def infer_schema(chain: "DataChain", **kwargs) -> pa.Schema:
 
 
 def schema_to_output(
-    schema: pa.Schema, col_names: Optional[Sequence[str]] = None
+    schema: pa.Schema, col_names: Sequence[str] | None = None
 ) -> tuple[dict[str, type], list[str]]:
     """
     Generate UDF output schema from pyarrow schema.
@@ -205,14 +207,15 @@ def schema_to_output(
     hf_schema = _get_hf_schema(schema)
     if hf_schema:
         return {
-            column: hf_type for hf_type, column in zip(hf_schema[1].values(), col_names)
+            column: hf_type
+            for hf_type, column in zip(hf_schema[1].values(), col_names, strict=False)
         }, list(normalized_col_dict.values())
 
     output = {}
-    for field, column in zip(schema, col_names):
+    for field, column in zip(schema, col_names, strict=False):
         dtype = arrow_type_mapper(field.type, column)
         if field.nullable and not ModelStore.is_pydantic(dtype):
-            dtype = Optional[dtype]  # type: ignore[assignment]
+            dtype = dtype | None  # type: ignore[assignment]
         output[column] = dtype
 
     return output, list(normalized_col_dict.values())
@@ -243,7 +246,7 @@ def arrow_type_mapper(col_type: pa.DataType, column: str = "") -> type:  # noqa:
         for field in col_type:
             dtype = arrow_type_mapper(field.type, field.name)
             if field.nullable and not ModelStore.is_pydantic(dtype):
-                dtype = Optional[dtype]  # type: ignore[assignment]
+                dtype = dtype | None  # type: ignore[assignment]
             type_dict[field.name] = dtype
         return dict_to_data_model(f"ArrowDataModel_{column}", type_dict)
     if pa.types.is_map(col_type):
@@ -257,7 +260,7 @@ def arrow_type_mapper(col_type: pa.DataType, column: str = "") -> type:  # noqa:
 
 def _get_hf_schema(
     schema: "pa.Schema",
-) -> Optional[tuple["Features", dict[str, "DataType"]]]:
+) -> tuple["Features", dict[str, "DataType"]] | None:
     if schema.metadata and b"huggingface" in schema.metadata:
         from datachain.lib.hf import get_output_schema, schema_from_arrow
 
@@ -266,7 +269,7 @@ def _get_hf_schema(
     return None
 
 
-def _get_datachain_schema(schema: "pa.Schema") -> Optional[SignalSchema]:
+def _get_datachain_schema(schema: "pa.Schema") -> SignalSchema | None:
     """Return a restored SignalSchema from parquet metadata, if any is found."""
     if schema.metadata and DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY in schema.metadata:
         serialized_signal_schema = json.loads(

datachain/lib/audio.py

@@ -1,5 +1,5 @@
 import posixpath
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from datachain.lib.file import FileError
 
@@ -18,7 +18,7 @@ except ImportError as exc:
     ) from exc
 
 
-def audio_info(file: "Union[File, AudioFile]") -> "Audio":
+def audio_info(file: "File | AudioFile") -> "Audio":
     """Extract metadata like sample rate, channels, duration, and format."""
     from datachain.lib.file import Audio
 
@@ -99,7 +99,7 @@ def _encoding_to_format(encoding: str, file_ext: str) -> str:
 
 
 def audio_to_np(
-    audio: "AudioFile", start: float = 0, duration: Optional[float] = None
+    audio: "AudioFile", start: float = 0, duration: float | None = None
 ) -> "tuple[ndarray, int]":
     """Load audio fragment as numpy array.
     Multi-channel audio is transposed to (samples, channels)."""
@@ -146,7 +146,7 @@ def audio_to_bytes(
     audio: "AudioFile",
     format: str = "wav",
     start: float = 0,
-    duration: Optional[float] = None,
+    duration: float | None = None,
 ) -> bytes:
     """Convert audio to bytes using soundfile.
 
@@ -166,9 +166,9 @@ def audio_to_bytes(
 def save_audio(
     audio: "AudioFile",
     output: str,
-    format: Optional[str] = None,
+    format: str | None = None,
     start: float = 0,
-    end: Optional[float] = None,
+    end: float | None = None,
 ) -> "AudioFile":
     """Save audio file or extract fragment to specified format.
 

datachain/lib/clip.py

@@ -1,5 +1,6 @@
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal, Union
 
 import torch
 from transformers.modeling_utils import PreTrainedModel
@@ -32,14 +33,14 @@ def _get_encoder(model: Any, type: Literal["image", "text"]) -> Callable:
 
 
 def clip_similarity_scores(
-    images: Union[None, "Image.Image", list["Image.Image"]],
-    text: Union[None, str, list[str]],
+    images: Union["Image.Image", list["Image.Image"]] | None,
+    text: str | list[str] | None,
     model: Any,
     preprocess: Callable,
     tokenizer: Callable,
     prob: bool = False,
     image_to_text: bool = True,
-    device: Optional[Union[str, torch.device]] = None,
+    device: str | torch.device | None = None,
 ) -> list[list[float]]:
     """
     Calculate CLIP similarity scores between one or more images and/or text.

datachain/lib/convert/python_to_sql.py

@@ -1,14 +1,9 @@
 import inspect
-import sys
 from datetime import datetime
 from enum import Enum
+from types import UnionType
 from typing import Annotated, Literal, Union, get_args, get_origin
 
-if sys.version_info >= (3, 10):
-    from types import UnionType
-else:
-    UnionType = None
-
 from pydantic import BaseModel
 from typing_extensions import Literal as LiteralEx
 
@@ -40,13 +35,6 @@ PYTHON_TO_SQL = {
 }
 
 
-def _is_union(orig) -> bool:
-    if orig == Union:
-        return True
-    # some code is unreachab in python<3.10
-    return UnionType is not None and orig is UnionType  # type: ignore[unreachable]
-
-
 def python_to_sql(typ):  # noqa: PLR0911
     if inspect.isclass(typ):
         if issubclass(typ, SQLType):
@@ -82,12 +70,12 @@ def python_to_sql(typ):  # noqa: PLR0911
     if inspect.isclass(orig) and issubclass(dict, orig):
         return JSON
 
-    if _is_union(orig):
+    if orig in (Union, UnionType):
         if len(args) == 2 and (type(None) in args):
             non_none_arg = args[0] if args[0] is not type(None) else args[1]
             return python_to_sql(non_none_arg)
 
-        if _is_union_str_literal(orig, args):
+        if all(arg is str or get_origin(arg) in (Literal, LiteralEx) for arg in args):
             return String
 
         if _is_json_inside_union(orig, args):
@@ -109,7 +97,7 @@ def list_of_args_to_type(args) -> SQLType:
 
 
 def _is_json_inside_union(orig, args) -> bool:
-    if _is_union(orig) and len(args) >= 2:
+    if orig in (Union, UnionType) and len(args) >= 2:
         # List in JSON: Union[dict, list[dict]]
         args_no_nones = [arg for arg in args if arg != type(None)]  # noqa: E721
         if len(args_no_nones) == 2:
@@ -123,9 +111,3 @@ def _is_json_inside_union(orig, args) -> bool:
         if any(inspect.isclass(arg) and issubclass(arg, BaseModel) for arg in args):
             return True
     return False
-
-
-def _is_union_str_literal(orig, args) -> bool:
-    if not _is_union(orig):
-        return False
-    return all(arg is str or get_origin(arg) in (Literal, LiteralEx) for arg in args)

datachain/lib/convert/values_to_tuples.py

@@ -1,13 +1,8 @@
 import itertools
 from collections.abc import Sequence
-from typing import Any, Union
+from typing import Any
 
-from datachain.lib.data_model import (
-    DataType,
-    DataTypeNames,
-    DataValue,
-    is_chain_type,
-)
+from datachain.lib.data_model import DataType, DataTypeNames, DataValue, is_chain_type
 from datachain.lib.utils import DataChainParamsError
 
 
@@ -20,7 +15,7 @@ class ValuesToTupleError(DataChainParamsError):
 
 def values_to_tuples(  # noqa: C901, PLR0912
     ds_name: str = "",
-    output: Union[None, DataType, Sequence[str], dict[str, DataType]] = None,
+    output: DataType | Sequence[str] | dict[str, DataType] | None = None,
     **fr_map: Sequence[DataValue],
 ) -> tuple[Any, Any, Any]:
     if output:
@@ -111,7 +106,7 @@ def values_to_tuples(  # noqa: C901, PLR0912
     if len(output) > 1:  # type: ignore[arg-type]
         tuple_type = tuple(output_types)
         res_type = tuple[tuple_type]  # type: ignore[valid-type]
-        res_values: Sequence[Any] = list(zip(*fr_map.values()))
+        res_values: Sequence[Any] = list(zip(*fr_map.values(), strict=False))
     else:
         res_type = output_types[0]  # type: ignore[misc]
         res_values = next(iter(fr_map.values()))

datachain/lib/data_model.py

@@ -1,8 +1,9 @@
 import inspect
+import types
 import uuid
 from collections.abc import Sequence
 from datetime import datetime
-from typing import ClassVar, Optional, Union, get_args, get_origin
+from typing import ClassVar, Union, get_args, get_origin
 
 from pydantic import AliasChoices, BaseModel, Field, create_model
 from pydantic.fields import FieldInfo
@@ -10,19 +11,19 @@ from pydantic.fields import FieldInfo
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import normalize_col_names
 
-StandardType = Union[
-    type[int],
-    type[str],
-    type[float],
-    type[bool],
-    type[list],
-    type[dict],
-    type[bytes],
-    type[datetime],
-]
-DataType = Union[type[BaseModel], StandardType]
+StandardType = (
+    type[int]
+    | type[str]
+    | type[float]
+    | type[bool]
+    | type[list]
+    | type[dict]
+    | type[bytes]
+    | type[datetime]
+)
+DataType = type[BaseModel] | StandardType
 DataTypeNames = "BaseModel, int, str, float, bool, list, dict, bytes, datetime"
-DataValue = Union[BaseModel, int, str, float, bool, list, dict, bytes, datetime]
+DataValue = BaseModel | int | str | float | bool | list | dict | bytes | datetime
 
 
 class DataModel(BaseModel):
@@ -37,7 +38,7 @@ class DataModel(BaseModel):
         ModelStore.register(cls)
 
     @staticmethod
-    def register(models: Union[DataType, Sequence[DataType]]):
+    def register(models: DataType | Sequence[DataType]):
         """For registering classes manually. It accepts a single class or a sequence of
         classes."""
         if not isinstance(models, Sequence):
@@ -63,8 +64,8 @@ def is_chain_type(t: type) -> bool:
     if orig is list and len(args) == 1:
         return is_chain_type(get_args(t)[0])
 
-    if orig is Union and len(args) == 2 and (type(None) in args):
-        return is_chain_type(args[0])
+    if orig in (Union, types.UnionType) and len(args) == 2 and (type(None) in args):
+        return is_chain_type(args[0] if args[1] is type(None) else args[1])
 
     return False
 
@@ -72,19 +73,19 @@ def is_chain_type(t: type) -> bool:
 def dict_to_data_model(
     name: str,
     data_dict: dict[str, DataType],
-    original_names: Optional[list[str]] = None,
+    original_names: list[str] | None = None,
 ) -> type[BaseModel]:
     if not original_names:
         # Gets a map of a normalized_name -> original_name
         columns = normalize_col_names(list(data_dict))
-        data_dict = dict(zip(columns.keys(), data_dict.values()))
+        data_dict = dict(zip(columns.keys(), data_dict.values(), strict=False))
         original_names = list(columns.values())
 
     fields = {
         name: (
             anno
             if inspect.isclass(anno) and issubclass(anno, BaseModel)
-            else Optional[anno],
+            else anno | None,
             Field(
                 validation_alias=AliasChoices(name, original_names[idx] or name),
                 default=None,

datachain/lib/dataset_info.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
 from pydantic import Field, field_validator
@@ -28,9 +28,9 @@ class DatasetInfo(DataModel):
     version: str = Field(default=DEFAULT_DATASET_VERSION)
     status: int = Field(default=DatasetStatus.CREATED)
     created_at: datetime = Field(default=TIME_ZERO)
-    finished_at: Optional[datetime] = Field(default=None)
-    num_objects: Optional[int] = Field(default=None)
-    size: Optional[int] = Field(default=None)
+    finished_at: datetime | None = Field(default=None)
+    num_objects: int | None = Field(default=None)
+    size: int | None = Field(default=None)
     params: dict[str, str] = Field(default={})
     metrics: dict[str, Any] = Field(default={})
     error_message: str = Field(default="")
@@ -59,7 +59,7 @@ class DatasetInfo(DataModel):
 
     @staticmethod
     def _validate_dict(
-        v: Optional[Union[str, dict]],
+        v: str | dict | None,
     ) -> dict:
         if v is None or v == "":
             return {}
@@ -88,7 +88,7 @@ class DatasetInfo(DataModel):
         cls,
         dataset: DatasetListRecord,
         version: DatasetListVersion,
-        job: Optional[Job],
+        job: Job | None,
     ) -> "Self":
         return cls(
             uuid=version.uuid,