datachain 0.34.5__py3-none-any.whl → 0.34.7__py3-none-any.whl

datachain/asyn.py

@@ -3,6 +3,7 @@ import threading
 from collections.abc import (
     AsyncIterable,
     Awaitable,
+    Callable,
     Coroutine,
     Generator,
     Iterable,
@@ -10,7 +11,7 @@ from collections.abc import (
 )
 from concurrent.futures import ThreadPoolExecutor, wait
 from heapq import heappop, heappush
-from typing import Any, Callable, Generic, Optional, TypeVar
+from typing import Any, Generic, TypeVar
 
 from fsspec.asyn import get_loop
 
@@ -49,7 +50,7 @@ class AsyncMapper(Generic[InputT, ResultT]):
         iterable: Iterable[InputT],
         *,
         workers: int = ASYNC_WORKERS,
-        loop: Optional[asyncio.AbstractEventLoop] = None,
+        loop: asyncio.AbstractEventLoop | None = None,
     ):
         self.func = func
         self.iterable = iterable
@@ -107,9 +108,7 @@ class AsyncMapper(Generic[InputT, ResultT]):
 
     async def init(self) -> None:
         self.work_queue = asyncio.Queue(2 * self.workers)
-        self.result_queue: asyncio.Queue[Optional[ResultT]] = asyncio.Queue(
-            self.workers
-        )
+        self.result_queue: asyncio.Queue[ResultT | None] = asyncio.Queue(self.workers)
 
     async def run(self) -> None:
         producer = self.start_task(self.produce())
@@ -149,10 +148,10 @@ class AsyncMapper(Generic[InputT, ResultT]):
             if exc:
                 raise exc
 
-    async def _pop_result(self) -> Optional[ResultT]:
+    async def _pop_result(self) -> ResultT | None:
         return await self.result_queue.get()
 
-    def next_result(self, timeout=None) -> Optional[ResultT]:
+    def next_result(self, timeout=None) -> ResultT | None:
         """
         Return the next available result.
 
@@ -212,17 +211,17 @@ class OrderedMapper(AsyncMapper[InputT, ResultT]):
         iterable: Iterable[InputT],
         *,
         workers: int = ASYNC_WORKERS,
-        loop: Optional[asyncio.AbstractEventLoop] = None,
+        loop: asyncio.AbstractEventLoop | None = None,
     ):
         super().__init__(func, iterable, workers=workers, loop=loop)
         self._waiters: dict[int, Any] = {}
-        self._getters: dict[int, asyncio.Future[Optional[ResultT]]] = {}
-        self.heap: list[tuple[int, Optional[ResultT]]] = []
+        self._getters: dict[int, asyncio.Future[ResultT | None]] = {}
+        self.heap: list[tuple[int, ResultT | None]] = []
         self._next_yield = 0
         self._items_seen = 0
         self._window = 2 * workers
 
-    def _push_result(self, i: int, result: Optional[ResultT]) -> None:
+    def _push_result(self, i: int, result: ResultT | None) -> None:
         if i in self._getters:
             future = self._getters.pop(i)
             future.set_result(result)
@@ -243,7 +242,7 @@ class OrderedMapper(AsyncMapper[InputT, ResultT]):
     async def init(self) -> None:
         self.work_queue = asyncio.Queue(2 * self.workers)
 
-    async def _pop_result(self) -> Optional[ResultT]:
+    async def _pop_result(self) -> ResultT | None:
         if self.heap and self.heap[0][0] == self._next_yield:
             _i, out = heappop(self.heap)
         else:

datachain/cache.py

@@ -2,7 +2,7 @@ import os
 from collections.abc import Iterator
 from contextlib import contextmanager
 from tempfile import mkdtemp
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from dvc_data.hashfile.db.local import LocalHashFileDB
 from dvc_objects.fs.local import LocalFileSystem
@@ -22,14 +22,14 @@ def try_scandir(path):
         pass
 
 
-def get_temp_cache(tmp_dir: str, prefix: Optional[str] = None) -> "Cache":
+def get_temp_cache(tmp_dir: str, prefix: str | None = None) -> "Cache":
     cache_dir = mkdtemp(prefix=prefix, dir=tmp_dir)
     return Cache(cache_dir, tmp_dir=tmp_dir)
 
 
 @contextmanager
 def temporary_cache(
-    tmp_dir: str, prefix: Optional[str] = None, delete: bool = True
+    tmp_dir: str, prefix: str | None = None, delete: bool = True
 ) -> Iterator["Cache"]:
     cache = get_temp_cache(tmp_dir, prefix=prefix)
     try:
@@ -58,7 +58,7 @@ class Cache:  # noqa: PLW1641
     def tmp_dir(self):
         return self.odb.tmp_dir
 
-    def get_path(self, file: "File") -> Optional[str]:
+    def get_path(self, file: "File") -> str | None:
         if self.contains(file):
             return self.path_from_checksum(file.get_hash())
         return None
@@ -74,7 +74,7 @@ class Cache:  # noqa: PLW1641
         self.odb.delete(file.get_hash())
 
     async def download(
-        self, file: "File", client: "Client", callback: Optional[Callback] = None
+        self, file: "File", client: "Client", callback: Callback | None = None
     ) -> None:
         from dvc_objects.fs.utils import tmp_fname
 

datachain/catalog/catalog.py

@@ -9,20 +9,12 @@ import subprocess
 import sys
 import time
 import traceback
-from collections.abc import Iterable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from copy import copy
 from dataclasses import dataclass
 from functools import cached_property, reduce
 from threading import Thread
-from typing import (
-    IO,
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    NoReturn,
-    Optional,
-    Union,
-)
+from typing import IO, TYPE_CHECKING, Any, NoReturn
 from uuid import uuid4
 
 import sqlalchemy as sa
@@ -64,10 +56,7 @@ from datachain.utils import DataChainDir
 from .datasource import DataSource
 
 if TYPE_CHECKING:
-    from datachain.data_storage import (
-        AbstractMetastore,
-        AbstractWarehouse,
-    )
+    from datachain.data_storage import AbstractMetastore, AbstractWarehouse
     from datachain.dataset import DatasetListVersion
     from datachain.job import Job
     from datachain.lib.listing_info import ListingInfo
@@ -120,8 +109,8 @@ def is_namespace_local(namespace_name) -> bool:
 
 def shutdown_process(
     proc: subprocess.Popen,
-    interrupt_timeout: Optional[int] = None,
-    terminate_timeout: Optional[int] = None,
+    interrupt_timeout: int | None = None,
+    terminate_timeout: int | None = None,
 ) -> int:
     """Shut down the process gracefully with SIGINT -> SIGTERM -> SIGKILL."""
 
@@ -168,7 +157,7 @@ class DatasetRowsFetcher(NodesThreadPool):
         remote_ds_version: str,
         local_ds: DatasetRecord,
         local_ds_version: str,
-        schema: dict[str, Union[SQLType, type[SQLType]]],
+        schema: dict[str, SQLType | type[SQLType]],
         max_threads: int = PULL_DATASET_MAX_THREADS,
         progress_bar=None,
     ):
@@ -183,7 +172,7 @@ class DatasetRowsFetcher(NodesThreadPool):
         self.local_ds = local_ds
         self.local_ds_version = local_ds_version
         self.schema = schema
-        self.last_status_check: Optional[float] = None
+        self.last_status_check: float | None = None
         self.studio_client = StudioClient()
         self.progress_bar = progress_bar
 
@@ -287,16 +276,16 @@ class DatasetRowsFetcher(NodesThreadPool):
 class NodeGroup:
     """Class for a group of nodes from the same source"""
 
-    listing: Optional["Listing"]
-    client: "Client"
+    listing: "Listing | None"
+    client: Client
     sources: list[DataSource]
 
     # The source path within the bucket
     # (not including the bucket name or s3:// prefix)
     source_path: str = ""
-    dataset_name: Optional[str] = None
-    dataset_version: Optional[str] = None
-    instantiated_nodes: Optional[list[NodeWithPath]] = None
+    dataset_name: str | None = None
+    dataset_version: str | None = None
+    instantiated_nodes: list[NodeWithPath] | None = None
 
     @property
     def is_dataset(self) -> bool:
@@ -323,7 +312,7 @@ def prepare_output_for_cp(
     output: str,
     force: bool = False,
     no_cp: bool = False,
-) -> tuple[bool, Optional[str]]:
+) -> tuple[bool, str | None]:
     total_node_count = 0
     for node_group in node_groups:
         if not node_group.sources:
@@ -372,7 +361,7 @@ def collect_nodes_for_cp(
 
     # Collect all sources to process
     for node_group in node_groups:
-        listing: Optional[Listing] = node_group.listing
+        listing: Listing | None = node_group.listing
         valid_sources: list[DataSource] = []
         for dsrc in node_group.sources:
             if dsrc.is_single_object():
@@ -416,7 +405,7 @@ def instantiate_node_groups(
     recursive: bool = False,
     virtual_only: bool = False,
     always_copy_dir_contents: bool = False,
-    copy_to_filename: Optional[str] = None,
+    copy_to_filename: str | None = None,
 ) -> None:
     instantiate_progress_bar = (
         None
@@ -444,7 +433,7 @@ def instantiate_node_groups(
     for node_group in node_groups:
         if not node_group.sources:
             continue
-        listing: Optional[Listing] = node_group.listing
+        listing: Listing | None = node_group.listing
         source_path: str = node_group.source_path
 
         copy_dir_contents = always_copy_dir_contents or source_path.endswith("/")
@@ -527,10 +516,8 @@ class Catalog:
         warehouse: "AbstractWarehouse",
         cache_dir=None,
         tmp_dir=None,
-        client_config: Optional[dict[str, Any]] = None,
-        warehouse_ready_callback: Optional[
-            Callable[["AbstractWarehouse"], None]
-        ] = None,
+        client_config: dict[str, Any] | None = None,
+        warehouse_ready_callback: Callable[["AbstractWarehouse"], None] | None = None,
         in_memory: bool = False,
     ):
         datachain_dir = DataChainDir(cache=cache_dir, tmp=tmp_dir)
@@ -592,7 +579,7 @@ class Catalog:
         client_config=None,
         column="file",
         skip_indexing=False,
-    ) -> tuple[Optional["Listing"], "Client", str]:
+    ) -> tuple["Listing | None", Client, str]:
         from datachain import read_storage
         from datachain.listing import Listing
 
@@ -633,7 +620,7 @@ class Catalog:
         skip_indexing=False,
         client_config=None,
         only_index=False,
-    ) -> Optional[list["DataSource"]]:
+    ) -> list["DataSource"] | None:
         enlisted_sources = []
         for src in sources:  # Opt: parallel
             listing, client, file_path = self.enlist_source(
@@ -679,7 +666,7 @@ class Catalog:
         enlisted_sources: list[tuple[bool, bool, Any]] = []
         client_config = client_config or self.client_config
         for src in sources:  # Opt: parallel
-            listing: Optional[Listing]
+            listing: Listing | None
             if src.startswith("ds://"):
                 ds_name, ds_version = parse_dataset_uri(src)
                 ds_namespace, ds_project, ds_name = parse_dataset_name(ds_name)
@@ -785,19 +772,19 @@ class Catalog:
     def create_dataset(
         self,
         name: str,
-        project: Optional[Project] = None,
-        version: Optional[str] = None,
+        project: Project | None = None,
+        version: str | None = None,
         *,
         columns: Sequence[Column],
-        feature_schema: Optional[dict] = None,
+        feature_schema: dict | None = None,
         query_script: str = "",
-        create_rows: Optional[bool] = True,
-        validate_version: Optional[bool] = True,
-        listing: Optional[bool] = False,
-        uuid: Optional[str] = None,
-        description: Optional[str] = None,
-        attrs: Optional[list[str]] = None,
-        update_version: Optional[str] = "patch",
+        create_rows: bool | None = True,
+        validate_version: bool | None = True,
+        listing: bool | None = False,
+        uuid: str | None = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
+        update_version: str | None = "patch",
     ) -> "DatasetRecord":
         """
         Creates new dataset of a specific version.
@@ -886,8 +873,8 @@ class Catalog:
         error_stack="",
         script_output="",
         create_rows_table=True,
-        job_id: Optional[str] = None,
-        uuid: Optional[str] = None,
+        job_id: str | None = None,
+        uuid: str | None = None,
     ) -> DatasetRecord:
         """
         Creates dataset version if it doesn't exist.
@@ -971,7 +958,7 @@ class Catalog:
         return dataset_updated
 
     def remove_dataset_version(
-        self, dataset: DatasetRecord, version: str, drop_rows: Optional[bool] = True
+        self, dataset: DatasetRecord, version: str, drop_rows: bool | None = True
     ) -> None:
         """
         Deletes one single dataset version.
@@ -999,7 +986,7 @@ class Catalog:
         self,
         name: str,
         sources: list[str],
-        project: Optional[Project] = None,
+        project: Project | None = None,
         client_config=None,
         recursive=False,
     ) -> DatasetRecord:
@@ -1068,8 +1055,8 @@ class Catalog:
     def get_full_dataset_name(
         self,
         name: str,
-        project_name: Optional[str] = None,
-        namespace_name: Optional[str] = None,
+        project_name: str | None = None,
+        namespace_name: str | None = None,
     ) -> tuple[str, str, str]:
         """
         Returns dataset name together with separated namespace and project name.
@@ -1101,8 +1088,8 @@ class Catalog:
     def get_dataset(
         self,
         name: str,
-        namespace_name: Optional[str] = None,
-        project_name: Optional[str] = None,
+        namespace_name: str | None = None,
+        project_name: str | None = None,
     ) -> DatasetRecord:
         from datachain.lib.listing import is_listing_dataset
 
@@ -1122,7 +1109,7 @@ class Catalog:
         name: str,
         namespace_name: str,
         project_name: str,
-        version: Optional[str] = None,
+        version: str | None = None,
         pull_dataset: bool = False,
         update: bool = False,
     ) -> DatasetRecord:
@@ -1213,10 +1200,10 @@ class Catalog:
         self,
         name: str,
         version: str,
-        namespace_name: Optional[str] = None,
-        project_name: Optional[str] = None,
+        namespace_name: str | None = None,
+        project_name: str | None = None,
         indirect=False,
-    ) -> list[Optional[DatasetDependency]]:
+    ) -> list[DatasetDependency | None]:
         dataset = self.get_dataset(
             name,
             namespace_name=namespace_name,
@@ -1248,10 +1235,10 @@ class Catalog:
 
     def ls_datasets(
         self,
-        prefix: Optional[str] = None,
+        prefix: str | None = None,
         include_listing: bool = False,
         studio: bool = False,
-        project: Optional[Project] = None,
+        project: Project | None = None,
     ) -> Iterator[DatasetListRecord]:
         from datachain.remote.studio import StudioClient
 
@@ -1283,12 +1270,12 @@ class Catalog:
 
     def list_datasets_versions(
         self,
-        prefix: Optional[str] = None,
+        prefix: str | None = None,
         include_listing: bool = False,
         with_job: bool = True,
         studio: bool = False,
-        project: Optional[Project] = None,
-    ) -> Iterator[tuple[DatasetListRecord, "DatasetListVersion", Optional["Job"]]]:
+        project: Project | None = None,
+    ) -> Iterator[tuple[DatasetListRecord, "DatasetListVersion", "Job | None"]]:
         """Iterate over all dataset versions with related jobs."""
         datasets = list(
             self.ls_datasets(
@@ -1316,7 +1303,7 @@ class Catalog:
                 for v in d.versions
             )
 
-    def listings(self, prefix: Optional[str] = None) -> list["ListingInfo"]:
+    def listings(self, prefix: str | None = None) -> list["ListingInfo"]:
         """
         Returns list of ListingInfo objects which are representing specific
         storage listing datasets
@@ -1367,9 +1354,9 @@ class Catalog:
         self,
         source: str,
         path: str,
-        version_id: Optional[str] = None,
+        version_id: str | None = None,
         client_config=None,
-        content_disposition: Optional[str] = None,
+        content_disposition: str | None = None,
         **kwargs,
     ) -> str:
         client_config = client_config or self.client_config
@@ -1388,7 +1375,7 @@ class Catalog:
         bucket_uri: str,
         name: str,
         version: str,
-        project: Optional[Project] = None,
+        project: Project | None = None,
         client_config=None,
     ) -> list[str]:
         dataset = self.get_dataset(
@@ -1402,7 +1389,7 @@ class Catalog:
         )
 
     def dataset_table_export_file_names(
-        self, name: str, version: str, project: Optional[Project] = None
+        self, name: str, version: str, project: Project | None = None
     ) -> list[str]:
         dataset = self.get_dataset(
             name,
@@ -1414,9 +1401,9 @@ class Catalog:
     def remove_dataset(
         self,
         name: str,
-        project: Optional[Project] = None,
-        version: Optional[str] = None,
-        force: Optional[bool] = False,
+        project: Project | None = None,
+        version: str | None = None,
+        force: bool | None = False,
     ):
         dataset = self.get_dataset(
             name,
@@ -1444,10 +1431,10 @@ class Catalog:
     def edit_dataset(
         self,
         name: str,
-        project: Optional[Project] = None,
-        new_name: Optional[str] = None,
-        description: Optional[str] = None,
-        attrs: Optional[list[str]] = None,
+        project: Project | None = None,
+        new_name: str | None = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
     ) -> DatasetRecord:
         update_data = {}
         if new_name:
@@ -1487,9 +1474,9 @@ class Catalog:
     def pull_dataset(  # noqa: C901, PLR0915
         self,
         remote_ds_uri: str,
-        output: Optional[str] = None,
-        local_ds_name: Optional[str] = None,
-        local_ds_version: Optional[str] = None,
+        output: str | None = None,
+        local_ds_name: str | None = None,
+        local_ds_version: str | None = None,
         cp: bool = False,
         force: bool = False,
         *,
@@ -1763,21 +1750,26 @@ class Catalog:
     def query(
         self,
         query_script: str,
-        env: Optional[Mapping[str, str]] = None,
+        env: Mapping[str, str] | None = None,
         python_executable: str = sys.executable,
         capture_output: bool = False,
         output_hook: Callable[[str], None] = noop,
-        params: Optional[dict[str, str]] = None,
-        job_id: Optional[str] = None,
-        interrupt_timeout: Optional[int] = None,
-        terminate_timeout: Optional[int] = None,
+        params: dict[str, str] | None = None,
+        job_id: str | None = None,
+        reset: bool = False,
+        interrupt_timeout: int | None = None,
+        terminate_timeout: int | None = None,
     ) -> None:
+        if not isinstance(reset, bool):
+            raise TypeError(f"reset must be a bool, got {type(reset).__name__}")
+
         cmd = [python_executable, "-c", query_script]
         env = dict(env or os.environ)
         env.update(
             {
                 "DATACHAIN_QUERY_PARAMS": json.dumps(params or {}),
                 "DATACHAIN_JOB_ID": job_id or "",
+                "DATACHAIN_CHECKPOINTS_RESET": str(reset),
             },
         )
         popen_kwargs: dict[str, Any] = {}
@@ -1787,7 +1779,7 @@ class Catalog:
         def raise_termination_signal(sig: int, _: Any) -> NoReturn:
             raise TerminationSignal(sig)
 
-        thread: Optional[Thread] = None
+        thread: Thread | None = None
         with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
             logger.info("Starting process %s", proc.pid)
 
@@ -1850,7 +1842,7 @@ class Catalog:
         no_cp: bool = False,
         no_glob: bool = False,
         *,
-        client_config: Optional["dict"] = None,
+        client_config: dict | None = None,
     ) -> None:
         """
         This function copies files from cloud sources to local destination directory

datachain/catalog/loader.py

@@ -1,7 +1,7 @@
 import os
 import sys
 from importlib import import_module
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 from datachain.plugins import ensure_plugins_loaded
 from datachain.utils import get_envs_by_prefix
@@ -108,7 +108,7 @@ def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
     return warehouse_class(**warehouse_args)
 
 
-def get_udf_distributor_class() -> Optional[type["AbstractUDFDistributor"]]:
+def get_udf_distributor_class() -> type["AbstractUDFDistributor"] | None:
     if os.environ.get(DISTRIBUTED_DISABLED) == "True":
         return None
 
@@ -132,7 +132,7 @@ def get_udf_distributor_class() -> Optional[type["AbstractUDFDistributor"]]:
 
 
 def get_catalog(
-    client_config: Optional[dict[str, Any]] = None,
+    client_config: dict[str, Any] | None = None,
     in_memory: bool = False,
 ) -> "Catalog":
     """

datachain/checkpoint.py

@@ -1,7 +1,6 @@
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Union
 
 
 @dataclass
@@ -29,7 +28,7 @@ class Checkpoint:
     @classmethod
     def parse(
         cls,
-        id: Union[str, uuid.UUID],
+        id: str | uuid.UUID,
         job_id: str,
         _hash: str,
         partial: bool,

datachain/cli/__init__.py

@@ -3,10 +3,8 @@ import os
 import sys
 import traceback
 from multiprocessing import freeze_support
-from typing import Optional
 
 from datachain.cli.utils import get_logging_level
-from datachain.error import DataChainError as DataChainError
 
 from .commands import (
     clear_cache,
@@ -26,7 +24,7 @@ from .parser import get_parser
 logger = logging.getLogger("datachain")
 
 
-def main(argv: Optional[list[str]] = None) -> int:
+def main(argv: list[str] | None = None) -> int:
     from datachain.catalog import get_catalog
 
     # Required for Windows multiprocessing support
@@ -307,7 +305,7 @@ def handle_udf() -> int:
     return udf_entrypoint()
 
 
-def handle_udf_runner(fd: Optional[int] = None) -> int:
+def handle_udf_runner(fd: int | None = None) -> int:
     from datachain.query.dispatch import udf_worker_entrypoint
 
     return udf_worker_entrypoint(fd)

datachain/cli/commands/datasets.py

@@ -1,6 +1,6 @@
 import sys
 from collections.abc import Iterable, Iterator
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from tabulate import tabulate
 
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 
 def group_dataset_versions(
     datasets: Iterable[tuple[str, str]], latest_only=True
-) -> dict[str, Union[str, list[str]]]:
+) -> dict[str, str | list[str]]:
     grouped: dict[str, list[tuple[int, int, int]]] = {}
 
     # Sort to ensure groupby works as expected
@@ -43,9 +43,9 @@ def list_datasets(
     studio: bool = False,
     local: bool = False,
     all: bool = True,
-    team: Optional[str] = None,
+    team: str | None = None,
     latest_only: bool = True,
-    name: Optional[str] = None,
+    name: str | None = None,
 ) -> None:
     token = Config().read().get("studio", {}).get("token")
     all, local, studio = determine_flavors(studio, local, all, token)
@@ -107,7 +107,7 @@ def list_datasets(
 
 
 def list_datasets_local(
-    catalog: "Catalog", name: Optional[str] = None
+    catalog: "Catalog", name: str | None = None
 ) -> Iterator[tuple[str, str]]:
     if name:
         yield from list_datasets_local_versions(catalog, name)
@@ -147,10 +147,10 @@ def _datasets_tabulate_row(name, both, local_version, studio_version) -> dict[st
 def rm_dataset(
     catalog: "Catalog",
     name: str,
-    version: Optional[str] = None,
-    force: Optional[bool] = False,
-    studio: Optional[bool] = False,
-    team: Optional[str] = None,
+    version: str | None = None,
+    force: bool | None = False,
+    studio: bool | None = False,
+    team: str | None = None,
 ) -> None:
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
@@ -177,10 +177,10 @@ def rm_dataset(
 def edit_dataset(
     catalog: "Catalog",
     name: str,
-    new_name: Optional[str] = None,
-    description: Optional[str] = None,
-    attrs: Optional[list[str]] = None,
-    team: Optional[str] = None,
+    new_name: str | None = None,
+    description: str | None = None,
+    attrs: list[str] | None = None,
+    team: str | None = None,
 ) -> None:
     from datachain.lib.dc.utils import is_studio