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

datachain/query/dataset.py

@@ -8,19 +8,11 @@ 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
@@ -67,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.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
@@ -83,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")
@@ -411,14 +401,14 @@ def get_generated_callback(is_generator: bool = False) -> Callback:
 class UDFStep(Step, ABC):
     udf: "UDFAdapter"
     catalog: "Catalog"
-    partition_by: Optional[PartitionByType] = None
+    partition_by: PartitionByType | None = None
     is_generator = False
     # Parameters from Settings
     cache: bool = False
-    parallel: Optional[int] = None
-    workers: Union[bool, int] = False
-    min_task_size: Optional[int] = None
-    batch_size: 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 [])
@@ -624,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,
@@ -681,14 +671,14 @@ class UDFStep(Step, ABC):
 class UDFSignal(UDFStep):
     udf: "UDFAdapter"
     catalog: "Catalog"
-    partition_by: Optional[PartitionByType] = None
+    partition_by: PartitionByType | None = None
     is_generator = False
     # Parameters from Settings
     cache: bool = False
-    parallel: Optional[int] = None
-    workers: Union[bool, int] = False
-    min_task_size: Optional[int] = None
-    batch_size: Optional[int] = None
+    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]] = [
@@ -760,14 +750,14 @@ class RowGenerator(UDFStep):
 
     udf: "UDFAdapter"
     catalog: "Catalog"
-    partition_by: Optional[PartitionByType] = None
+    partition_by: PartitionByType | None = None
     is_generator = True
     # Parameters from Settings
     cache: bool = False
-    parallel: Optional[int] = None
-    workers: Union[bool, int] = False
-    min_task_size: Optional[int] = None
-    batch_size: Optional[int] = None
+    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
@@ -814,7 +804,7 @@ 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)
 
@@ -825,7 +815,7 @@ class SQLClause(Step, ABC):
 
 @frozen
 class SQLSelect(SQLClause):
-    args: tuple[Union[Function, ColumnElement], ...]
+    args: tuple[Function | ColumnElement, ...]
 
     def hash_inputs(self) -> str:
         return hash_column_elements(self.args)
@@ -844,7 +834,7 @@ 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)
@@ -890,7 +880,7 @@ 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)
@@ -906,7 +896,7 @@ 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)
@@ -1011,7 +1001,7 @@ class SQLJoin(Step):
     catalog: "Catalog"
     query1: "DatasetQuery"
     query2: "DatasetQuery"
-    predicates: Union[JoinPredicateType, tuple[JoinPredicateType, ...]]
+    predicates: JoinPredicateType | tuple[JoinPredicateType, ...]
     inner: bool
     full: bool
     rname: str
@@ -1150,8 +1140,8 @@ 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(
@@ -1211,6 +1201,7 @@ def _validate_columns(
                 missing_left,
             ],
             ["left", "right"],
+            strict=False,
         )
         if missing_columns
     ]
@@ -1243,32 +1234,32 @@ 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
@@ -1352,7 +1343,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)
@@ -1447,6 +1438,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:
@@ -1461,7 +1453,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]:
@@ -1500,7 +1492,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(
@@ -1540,6 +1532,7 @@ class DatasetQuery:
         obj.steps = obj.steps.copy()
         if new_table:
             obj.table = self.get_table()
+        obj.temp_table_names = []
         return obj
 
     @detach
@@ -1720,7 +1713,7 @@ class DatasetQuery:
     def join(
         self,
         dataset_query: "DatasetQuery",
-        predicates: Union[JoinPredicateType, Sequence[JoinPredicateType]],
+        predicates: JoinPredicateType | Sequence[JoinPredicateType],
         inner=False,
         full=False,
         rname="{name}_right",
@@ -1762,17 +1755,17 @@ class DatasetQuery:
     def add_signals(
         self,
         udf: "UDFAdapter",
-        partition_by: Optional[PartitionByType] = None,
+        partition_by: PartitionByType | None = None,
         # Parameters from Settings
         cache: bool = False,
-        parallel: Optional[int] = None,
-        workers: Union[bool, int] = False,
-        min_task_size: Optional[int] = None,
-        batch_size: 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: Optional[int] = None,
-        namespace: Optional[str] = None,
-        project: Optional[str] = None,
+        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.
@@ -1813,17 +1806,17 @@ class DatasetQuery:
     def generate(
         self,
         udf: "UDFAdapter",
-        partition_by: Optional[PartitionByType] = None,
+        partition_by: PartitionByType | None = None,
         # Parameters from Settings
         cache: bool = False,
-        parallel: Optional[int] = None,
-        workers: Union[bool, int] = False,
-        min_task_size: Optional[int] = None,
-        batch_size: 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: Optional[int] = None,
-        namespace: Optional[str] = None,
-        project: Optional[str] = None,
+        prefetch: int | None = None,
+        namespace: str | None = None,
+        project: str | None = None,
     ) -> "Self":
         query = self.clone()
         steps = query.steps
@@ -1879,23 +1872,23 @@ 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."""
@@ -1989,5 +1982,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

datachain/query/dispatch.py

@@ -3,9 +3,8 @@ from collections.abc import Iterable, Sequence
 from itertools import chain
 from multiprocessing import cpu_count
 from sys import stdin
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, Literal
 
-import multiprocess
 from cloudpickle import load, loads
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from multiprocess import get_context
@@ -27,6 +26,7 @@ from datachain.query.utils import get_query_id_column
 from datachain.utils import batched, flatten, safe_closing
 
 if TYPE_CHECKING:
+    import multiprocess
     from sqlalchemy import Select, Table
 
     from datachain.data_storage import AbstractMetastore, AbstractWarehouse
@@ -41,7 +41,7 @@ FAILED_STATUS = "FAILED"
 NOTIFY_STATUS = "NOTIFY"
 
 
-def get_n_workers_from_arg(n_workers: Optional[int] = None) -> int:
+def get_n_workers_from_arg(n_workers: int | None = None) -> int:
     if not n_workers:
         return cpu_count()
     if n_workers < 1:
@@ -86,7 +86,7 @@ def udf_entrypoint() -> int:
     return 0
 
 
-def udf_worker_entrypoint(fd: Optional[int] = None) -> int:
+def udf_worker_entrypoint(fd: int | None = None) -> int:
     if not (udf_distributor_class := get_udf_distributor_class()):
         raise RuntimeError(
             f"{DISTRIBUTED_IMPORT_PATH} import path is required "
@@ -97,9 +97,9 @@ def udf_worker_entrypoint(fd: Optional[int] = None) -> int:
 
 
 class UDFDispatcher:
-    _catalog: Optional[Catalog] = None
-    task_queue: Optional[multiprocess.Queue] = None
-    done_queue: Optional[multiprocess.Queue] = None
+    _catalog: Catalog | None = None
+    task_queue: "multiprocess.Queue | None" = None
+    done_queue: "multiprocess.Queue | None" = None
 
     def __init__(self, udf_info: UdfInfo, buffer_size: int = DEFAULT_BATCH_SIZE):
         self.udf_data = udf_info["udf_data"]

datachain/query/metrics.py

@@ -1,10 +1,9 @@
 import os
-from typing import Optional, Union
 
-metrics: dict[str, Union[str, int, float, bool, None]] = {}
+metrics: dict[str, str | int | float | bool | None] = {}
 
 
-def set(key: str, value: Union[str, int, float, bool, None]) -> None:  # noqa: PYI041
+def set(key: str, value: str | int | float | bool | None) -> None:  # noqa: PYI041
     """Set a metric value."""
     if not isinstance(key, str):
         raise TypeError("Key must be a string")
@@ -21,6 +20,6 @@ def set(key: str, value: Union[str, int, float, bool, None]) -> None:  # noqa: P
         metastore.update_job(job_id, metrics=metrics)
 
 
-def get(key: str) -> Optional[Union[str, int, float, bool]]:
+def get(key: str) -> str | int | float | bool | None:
     """Get a metric value."""
     return metrics[key]

datachain/query/params.py

@@ -1,11 +1,10 @@
 import json
 import os
-from typing import Optional
 
-params_cache: Optional[dict[str, str]] = None
+params_cache: dict[str, str] | None = None
 
 
-def param(key: str, default: Optional[str] = None) -> Optional[str]:
+def param(key: str, default: str | None = None) -> str | None:
     """Get query parameter."""
     if not isinstance(key, str):
         raise TypeError("Param key must be a string")

datachain/query/schema.py

@@ -1,7 +1,8 @@
 import functools
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from fnmatch import fnmatch
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any
 
 import attrs
 import sqlalchemy as sa
@@ -42,7 +43,7 @@ class ColumnMeta(type):
 
 
 class Column(sa.ColumnClause, metaclass=ColumnMeta):
-    inherit_cache: Optional[bool] = True
+    inherit_cache: bool | None = True
 
     def __init__(self, text, type_=None, is_literal=False, _selectable=None):
         """Dataset column."""
@@ -177,7 +178,7 @@ class LocalFilename(UDFParameter):
     otherwise None will be returned.
     """
 
-    glob: Optional[str] = None
+    glob: str | None = None
 
     def get_value(
         self,
@@ -186,7 +187,7 @@ class LocalFilename(UDFParameter):
         *,
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
-    ) -> Optional[str]:
+    ) -> str | None:
         if self.glob and not fnmatch(row["name"], self.glob):  # type: ignore[type-var]
             # If the glob pattern is specified and the row filename
             # does not match it, then return None
@@ -205,7 +206,7 @@ class LocalFilename(UDFParameter):
         cache: bool = False,
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
-    ) -> Optional[str]:
+    ) -> str | None:
         if self.glob and not fnmatch(row["name"], self.glob):  # type: ignore[type-var]
             # If the glob pattern is specified and the row filename
             # does not match it, then return None
@@ -216,7 +217,7 @@ class LocalFilename(UDFParameter):
         return client.cache.get_path(file)
 
 
-UDFParamSpec = Union[str, Column, UDFParameter]
+UDFParamSpec = str | Column | UDFParameter
 
 
 def normalize_param(param: UDFParamSpec) -> UDFParameter:

datachain/query/session.py

@@ -3,7 +3,7 @@ import gc
 import logging
 import re
 import sys
-from typing import TYPE_CHECKING, ClassVar, Optional
+from typing import TYPE_CHECKING, ClassVar
 from uuid import uuid4
 
 from datachain.catalog import get_catalog
@@ -39,7 +39,7 @@ class Session:
     catalog (Catalog): Catalog object.
     """
 
-    GLOBAL_SESSION_CTX: Optional["Session"] = None
+    GLOBAL_SESSION_CTX: "Session | None" = None
     SESSION_CONTEXTS: ClassVar[list["Session"]] = []
     ORIGINAL_EXCEPT_HOOK = None
 
@@ -51,8 +51,8 @@ class Session:
     def __init__(
         self,
         name="",
-        catalog: Optional["Catalog"] = None,
-        client_config: Optional[dict] = None,
+        catalog: "Catalog | None" = None,
+        client_config: dict | None = None,
         in_memory: bool = False,
     ):
         if re.match(r"^[0-9a-zA-Z]*$", name) is None:
@@ -126,9 +126,9 @@ class Session:
     @classmethod
     def get(
         cls,
-        session: Optional["Session"] = None,
-        catalog: Optional["Catalog"] = None,
-        client_config: Optional[dict] = None,
+        session: "Session | None" = None,
+        catalog: "Catalog | None" = None,
+        client_config: dict | None = None,
         in_memory: bool = False,
     ) -> "Session":
         """Creates a Session() object from a catalog.

datachain/query/udf.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict, Union
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypedDict
 
 if TYPE_CHECKING:
     from sqlalchemy import Select, Table
@@ -17,7 +18,7 @@ class UdfInfo(TypedDict):
     query: "Select"
     udf_fields: list[str]
     batching: "BatchingStrategy"
-    processes: Optional[int]
+    processes: int | None
     is_generator: bool
     cache: bool
     rows_total: int
@@ -33,14 +34,14 @@ class AbstractUDFDistributor(ABC):
         query: "Select",
         udf_data: bytes,
         batching: "BatchingStrategy",
-        workers: Union[bool, int],
-        processes: Union[bool, int],
+        workers: bool | int,
+        processes: bool | int,
         udf_fields: list[str],
         rows_total: int,
         use_cache: bool,
         is_generator: bool = False,
-        min_task_size: Optional[Union[str, int]] = None,
-        batch_size: Optional[int] = None,
+        min_task_size: str | int | None = None,
+        batch_size: int | None = None,
     ) -> None: ...
 
     @abstractmethod
@@ -48,4 +49,4 @@ class AbstractUDFDistributor(ABC):
 
     @staticmethod
     @abstractmethod
-    def run_udf(fd: Optional[int] = None) -> int: ...
+    def run_udf(fd: int | None = None) -> int: ...

datachain/query/utils.py

@@ -1,8 +1,6 @@
-from typing import Optional, Union
-
 import sqlalchemy as sa
 
-ColT = Union[sa.ColumnClause, sa.Column, sa.ColumnElement, sa.TextClause, sa.Label]
+ColT = sa.ColumnClause | sa.Column | sa.ColumnElement | sa.TextClause | sa.Label
 
 
 def column_name(col: ColT) -> str:
@@ -14,12 +12,12 @@ def column_name(col: ColT) -> str:
     )
 
 
-def get_query_column(query: sa.Select, name: str) -> Optional[ColT]:
+def get_query_column(query: sa.Select, name: str) -> ColT | None:
     """Returns column element from query by name or None if column not found."""
     return next((col for col in query.inner_columns if column_name(col) == name), None)
 
 
-def get_query_id_column(query: sa.Select) -> Optional[sa.ColumnElement]:
+def get_query_id_column(query: sa.Select) -> sa.ColumnElement | None:
     """Returns ID column element from query or None if column not found."""
     col = get_query_column(query, "sys__id")
     return col if col is not None and isinstance(col, sa.ColumnElement) else None