datachain 0.3.17__py3-none-any.whl → 0.3.19__py3-none-any.whl

datachain/lib/dc.py

@@ -26,8 +26,8 @@ from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.values_to_tuples import values_to_tuples
 from datachain.lib.data_model import DataModel, DataType, dict_to_data_model
 from datachain.lib.dataset_info import DatasetInfo
+from datachain.lib.file import ArrowRow, File, get_file_type
 from datachain.lib.file import ExportPlacement as FileExportPlacement
-from datachain.lib.file import File, IndexedFile, get_file_type
 from datachain.lib.listing import (
     is_listing_dataset,
     is_listing_expired,
@@ -58,6 +58,7 @@ from datachain.query.dataset import (
 )
 from datachain.query.schema import DEFAULT_DELIMITER, Column, DatasetRow
 from datachain.sql.functions import path as pathfunc
+from datachain.telemetry import telemetry
 from datachain.utils import inside_notebook
 
 if TYPE_CHECKING:
@@ -246,6 +247,9 @@ class DataChain(DatasetQuery):
             **kwargs,
             indexing_column_types=File._datachain_column_types,
         )
+
+        telemetry.send_event_once("class", "datachain_init", **kwargs)
+
         if settings:
             self._settings = Settings(**settings)
         else:
@@ -1610,7 +1614,7 @@ class DataChain(DatasetQuery):
                 for name, info in output.model_fields.items()
             }
         if source:
-            output = {"source": IndexedFile} | output  # type: ignore[assignment,operator]
+            output = {"source": ArrowRow} | output  # type: ignore[assignment,operator]
         return self.gen(
             ArrowGenerator(schema, model, source, nrows, **kwargs), output=output
         )

datachain/lib/file.py

@@ -1,11 +1,14 @@
+import hashlib
 import io
 import json
 import logging
 import os
 import posixpath
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
 from contextlib import contextmanager
 from datetime import datetime
+from functools import partial
 from io import BytesIO
 from pathlib import Path, PurePosixPath
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, Union
@@ -14,12 +17,12 @@ from urllib.request import url2pathname
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from PIL import Image
+from pyarrow.dataset import dataset
 from pydantic import Field, field_validator
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-from datachain.cache import UniqueId
 from datachain.client.fileslice import FileSlice
 from datachain.lib.data_model import DataModel
 from datachain.lib.utils import DataChainError
@@ -27,7 +30,13 @@ from datachain.sql.types import JSON, Boolean, DateTime, Int, String
 from datachain.utils import TIME_ZERO
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from datachain.catalog import Catalog
+    from datachain.client.fsspec import Client
+    from datachain.dataset import RowDict
+
+sha256 = partial(hashlib.sha256, usedforsecurity=False)
 
 logger = logging.getLogger("datachain")
 
@@ -38,7 +47,7 @@ ExportPlacement = Literal["filename", "etag", "fullpath", "checksum"]
 class VFileError(DataChainError):
     def __init__(self, file: "File", message: str, vtype: str = ""):
         type_ = f" of vtype '{vtype}'" if vtype else ""
-        super().__init__(f"Error in v-file '{file.get_uid().path}'{type_}: {message}")
+        super().__init__(f"Error in v-file '{file.path}'{type_}: {message}")
 
 
 class FileError(DataChainError):
@@ -85,9 +94,8 @@ class TarVFile(VFile):
         tar_file = File(**parent)
         tar_file._set_stream(file._catalog)
 
-        tar_file_uid = tar_file.get_uid()
-        client = file._catalog.get_client(tar_file_uid.storage)
-        fd = client.open_object(tar_file_uid, use_cache=file._caching_enabled)
+        client = file._catalog.get_client(tar_file.source)
+        fd = client.open_object(tar_file, use_cache=file._caching_enabled)
         return FileSlice(fd, offset, size, file.name)
 
 
@@ -181,7 +189,11 @@ class File(DataModel):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self._catalog = None
-        self._caching_enabled = False
+        self._caching_enabled: bool = False
+
+    @classmethod
+    def _from_row(cls, row: "RowDict") -> "Self":
+        return cls(**{key: row[key] for key in cls._datachain_column_types})
 
     @property
     def name(self):
@@ -192,19 +204,18 @@ class File(DataModel):
         return str(PurePosixPath(self.path).parent)
 
     @contextmanager
-    def open(self, mode: Literal["rb", "r"] = "rb"):
+    def open(self, mode: Literal["rb", "r"] = "rb") -> Iterator[Any]:
         """Open the file and return a file object."""
         if self.location:
             with VFileRegistry.resolve(self, self.location) as f:  # type: ignore[arg-type]
                 yield f
 
         else:
-            uid = self.get_uid()
-            client = self._catalog.get_client(self.source)
             if self._caching_enabled:
-                client.download(uid, callback=self._download_cb)
+                self.ensure_cached()
+            client: Client = self._catalog.get_client(self.source)
             with client.open_object(
-                uid, use_cache=self._caching_enabled, cb=self._download_cb
+                self, use_cache=self._caching_enabled, cb=self._download_cb
             ) as f:
                 yield io.TextIOWrapper(f) if mode == "r" else f
 
@@ -252,23 +263,25 @@ class File(DataModel):
         self._caching_enabled = caching_enabled
         self._download_cb = download_cb
 
-    def get_uid(self) -> UniqueId:
-        """Returns unique ID for file."""
-        dump = self.model_dump()
-        return UniqueId(*(dump[k] for k in self._unique_id_keys))
+    def ensure_cached(self) -> None:
+        if self._catalog is None:
+            raise RuntimeError(
+                "cannot download file to cache because catalog is not setup"
+            )
+        client = self._catalog.get_client(self.source)
+        client.download(self, callback=self._download_cb)
 
-    def get_local_path(self, download: bool = False) -> Optional[str]:
-        """Returns path to a file in a local cache.
-        Return None if file is not cached. Throws an exception if cache is not setup."""
+    def get_local_path(self) -> Optional[str]:
+        """Return path to a file in a local cache.
+
+        Returns None if file is not cached.
+        Raises an exception if cache is not setup.
+        """
         if self._catalog is None:
             raise RuntimeError(
                 "cannot resolve local file path because catalog is not setup"
             )
-        uid = self.get_uid()
-        if download:
-            client = self._catalog.get_client(self.source)
-            client.download(uid, callback=self._download_cb)
-        return self._catalog.cache.get_path(uid)
+        return self._catalog.cache.get_path(self)
 
     def get_file_suffix(self):
         """Returns last part of file name with `.`."""
@@ -323,6 +336,12 @@ class File(DataModel):
         """Returns `fsspec` filesystem for the file."""
         return self._catalog.get_client(self.source).fs
 
+    def get_hash(self) -> str:
+        fingerprint = f"{self.source}/{self.path}/{self.version}/{self.etag}"
+        if self.location:
+            fingerprint += f"/{self.location}"
+        return sha256(fingerprint.encode()).hexdigest()
+
     def resolve(self) -> "Self":
         """
         Resolve a File object by checking its existence and updating its metadata.
@@ -421,14 +440,31 @@ class ImageFile(File):
         self.read().save(destination)
 
 
-class IndexedFile(DataModel):
-    """Metadata indexed from tabular files.
-
-    Includes `file` and `index` signals.
-    """
+class ArrowRow(DataModel):
+    """`DataModel` for reading row from Arrow-supported file."""
 
     file: File
     index: int
+    kwargs: dict
+
+    @contextmanager
+    def open(self):
+        """Stream row contents from indexed file."""
+        if self.file._caching_enabled:
+            self.file.ensure_cached()
+            path = self.file.get_local_path()
+            ds = dataset(path, **self.kwargs)
+
+        else:
+            path = self.file.get_path()
+            ds = dataset(path, filesystem=self.file.get_fs(), **self.kwargs)
+
+        return ds.take([self.index]).to_reader()
+
+    def read(self):
+        """Returns row contents as dict."""
+        with self.open() as record_batch:
+            return record_batch.to_pylist()[0]
 
 
 def get_file_type(type_: Literal["binary", "text", "image"] = "binary") -> type[File]:

datachain/lib/listing.py

@@ -11,6 +11,7 @@ from datachain.client import Client
 from datachain.lib.file import File
 from datachain.query.schema import Column
 from datachain.sql.functions import path as pathfunc
+from datachain.telemetry import telemetry
 from datachain.utils import uses_glob
 
 if TYPE_CHECKING:
@@ -80,6 +81,7 @@ def parse_listing_uri(uri: str, cache, client_config) -> tuple[str, str, str]:
     client_config = client_config or {}
     client = Client.get_client(uri, cache, **client_config)
     storage_uri, path = Client.parse_url(uri)
+    telemetry.log_param("client", client.PREFIX)
 
     # clean path without globs
     lst_uri_path = (

datachain/listing.py

@@ -156,12 +156,12 @@ class Listing:
 
     def instantiate_nodes(
         self,
-        all_nodes,
+        all_nodes: Iterable[NodeWithPath],
         output,
         total_files=None,
         force=False,
         shared_progress_bar=None,
-    ):
+    ) -> None:
         progress_bar = shared_progress_bar or tqdm(
             desc=f"Instantiating '{output}'",
             unit=" files",
@@ -175,8 +175,8 @@ class Listing:
             dst = os.path.join(output, *node.path)
             dst_dir = os.path.dirname(dst)
             os.makedirs(dst_dir, exist_ok=True)
-            uid = node.n.as_uid(self.client.uri)
-            self.client.instantiate_object(uid, dst, progress_bar, force)
+            file = node.n.to_file(self.client.uri)
+            self.client.instantiate_object(file, dst, progress_bar, force)
             counter += 1
             if counter > 1000:
                 progress_bar.update(counter)

datachain/node.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Optional
 
 import attrs
 
-from datachain.cache import UniqueId
+from datachain.lib.file import File
 from datachain.storage import StorageURI
 from datachain.utils import TIME_ZERO, time_to_str
 
@@ -99,11 +99,11 @@ class Node:
             return self.path + "/"
         return self.path
 
-    def as_uid(self, storage: Optional[StorageURI] = None) -> UniqueId:
-        if storage is None:
-            storage = self.source
-        return UniqueId(
-            storage=storage,
+    def to_file(self, source: Optional[StorageURI] = None) -> File:
+        if source is None:
+            source = self.source
+        return File(
+            source=source,
             path=self.path,
             size=self.size,
             version=self.version or "",

datachain/nodes_fetcher.py

@@ -1,12 +1,19 @@
 import logging
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
 
+from datachain.node import Node
 from datachain.nodes_thread_pool import NodesThreadPool
 
+if TYPE_CHECKING:
+    from datachain.cache import DataChainCache
+    from datachain.client.fsspec import Client
+
 logger = logging.getLogger("datachain")
 
 
 class NodesFetcher(NodesThreadPool):
-    def __init__(self, client, max_threads, cache):
+    def __init__(self, client: "Client", max_threads: int, cache: "DataChainCache"):
         super().__init__(max_threads)
         self.client = client
         self.cache = cache
@@ -15,7 +22,7 @@ class NodesFetcher(NodesThreadPool):
         for task in done:
             task.result()
 
-    def do_task(self, chunk):
+    def do_task(self, chunk: Iterable[Node]) -> None:
         from fsspec import Callback
 
         class _CB(Callback):
@@ -23,8 +30,8 @@ class NodesFetcher(NodesThreadPool):
                 self.increase_counter(inc)
 
         for node in chunk:
-            uid = node.as_uid(self.client.uri)
-            if self.cache.contains(uid):
+            file = node.to_file(self.client.uri)
+            if self.cache.contains(file):
                 self.increase_counter(node.size)
             else:
-                self.client.put_in_cache(uid, callback=_CB())
+                self.client.put_in_cache(file, callback=_CB())

datachain/nodes_thread_pool.py

@@ -20,7 +20,7 @@ class NodeChunk:
     def next_downloadable(self):
         node = next(self.nodes, None)
         while node and (
-            not node.is_downloadable or self.cache.contains(node.as_uid(self.storage))
+            not node.is_downloadable or self.cache.contains(node.to_file(self.storage))
         ):
             node = next(self.nodes, None)
         return node

datachain/progress.py

@@ -1,8 +1,6 @@
 """Manages progress bars."""
 
 import logging
-import os
-import re
 import sys
 from threading import RLock
 from typing import Any, ClassVar
@@ -10,20 +8,12 @@ from typing import Any, ClassVar
 from fsspec.callbacks import TqdmCallback
 from tqdm import tqdm
 
+from datachain.utils import env2bool
+
 logger = logging.getLogger(__name__)
 tqdm.set_lock(RLock())
 
 
-def env2bool(var, undefined=False):
-    """
-    undefined: return value if env var is unset
-    """
-    var = os.getenv(var, None)
-    if var is None:
-        return undefined
-    return bool(re.search("1|y|yes|true", var, flags=re.IGNORECASE))
-
-
 class Tqdm(tqdm):
     """
     maximum-compatibility tqdm-based progressbars

datachain/query/dataset.py

@@ -53,7 +53,7 @@ from datachain.utils import (
 
 from .schema import C, UDFParamSpec, normalize_param
 from .session import Session
-from .udf import UDFBase, UDFClassWrapper, UDFFactory, UDFType
+from .udf import UDFBase
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ClauseElement
@@ -364,7 +364,7 @@ def get_generated_callback(is_generator: bool = False) -> Callback:
 
 @frozen
 class UDFStep(Step, ABC):
-    udf: UDFType
+    udf: UDFBase
     catalog: "Catalog"
     partition_by: Optional[PartitionByType] = None
     parallel: Optional[int] = None
@@ -470,12 +470,6 @@ class UDFStep(Step, ABC):
 
             else:
                 # Otherwise process single-threaded (faster for smaller UDFs)
-                # Optionally instantiate the UDF instance if a class is provided.
-                if isinstance(self.udf, UDFFactory):
-                    udf: UDFBase = self.udf()
-                else:
-                    udf = self.udf
-
                 warehouse = self.catalog.warehouse
 
                 with contextlib.closing(
@@ -485,7 +479,7 @@ class UDFStep(Step, ABC):
                     processed_cb = get_processed_callback()
                     generated_cb = get_generated_callback(self.is_generator)
                     try:
-                        udf_results = udf.run(
+                        udf_results = self.udf.run(
                             udf_fields,
                             udf_inputs,
                             self.catalog,
@@ -498,7 +492,7 @@ class UDFStep(Step, ABC):
                             warehouse,
                             udf_table,
                             udf_results,
-                            udf,
+                            self.udf,
                             cb=generated_cb,
                         )
                     finally:
@@ -1471,7 +1465,7 @@ class DatasetQuery:
     @detach
     def add_signals(
         self,
-        udf: UDFType,
+        udf: UDFBase,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
@@ -1492,9 +1486,6 @@ class DatasetQuery:
         at least that minimum number of rows to each distributed worker, mostly useful
         if there are a very large number of small tasks to process.
         """
-        if isinstance(udf, UDFClassWrapper):  # type: ignore[unreachable]
-            # This is a bare decorated class, "instantiate" it now.
-            udf = udf()  # type: ignore[unreachable]
         query = self.clone()
         query.steps.append(
             UDFSignal(
@@ -1518,16 +1509,13 @@ class DatasetQuery:
     @detach
     def generate(
         self,
-        udf: UDFType,
+        udf: UDFBase,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
         partition_by: Optional[PartitionByType] = None,
         cache: bool = False,
     ) -> "Self":
-        if isinstance(udf, UDFClassWrapper):  # type: ignore[unreachable]
-            # This is a bare decorated class, "instantiate" it now.
-            udf = udf()  # type: ignore[unreachable]
         query = self.clone()
         steps = query.steps
         steps.append(
@@ -1616,25 +1604,3 @@ class DatasetQuery:
         finally:
             self.cleanup()
         return self.__class__(name=name, version=version, catalog=self.catalog)
-
-
-def query_wrapper(dataset_query: Any) -> Any:
-    """
-    Wrapper function that wraps the last statement of user query script.
-    Last statement MUST be instance of DatasetQuery, otherwise script exits with
-    error code 10
-    """
-    if not isinstance(dataset_query, DatasetQuery):
-        return dataset_query
-
-    catalog = dataset_query.catalog
-    save = bool(os.getenv("DATACHAIN_QUERY_SAVE"))
-
-    is_session_temp_dataset = dataset_query.name and dataset_query.name.startswith(
-        dataset_query.session.get_temp_prefix()
-    )
-
-    if save and (is_session_temp_dataset or not dataset_query.attached):
-        name = catalog.generate_query_dataset_name()
-        dataset_query = dataset_query.save(name)
-    return dataset_query

datachain/query/dispatch.py

@@ -27,7 +27,7 @@ from datachain.query.queue import (
     put_into_queue,
     unmarshal,
 )
-from datachain.query.udf import UDFBase, UDFFactory, UDFResult
+from datachain.query.udf import UDFBase, UDFResult
 from datachain.utils import batched_it
 
 DEFAULT_BATCH_SIZE = 10000
@@ -156,8 +156,6 @@ class UDFDispatcher:
 
     @property
     def batch_size(self):
-        if not self.udf:
-            self.udf = self.udf_factory()
         if self._batch_size is None:
             if hasattr(self.udf, "properties") and hasattr(
                 self.udf.properties, "batch"
@@ -181,18 +179,7 @@ class UDFDispatcher:
             self.catalog = Catalog(
                 id_generator, metastore, warehouse, **self.catalog_init_params
             )
-        udf = loads(self.udf_data)
-        # isinstance cannot be used here, as cloudpickle packages the entire class
-        # definition, and so these two types are not considered exactly equal,
-        # even if they have the same import path.
-        if full_module_type_path(type(udf)) != full_module_type_path(UDFFactory):
-            self.udf = udf
-        else:
-            self.udf = None
-            self.udf_factory = udf
-        if not self.udf:
-            self.udf = self.udf_factory()
-
+        self.udf = loads(self.udf_data)
         return UDFWorker(
             self.catalog,
             self.udf,

datachain/query/schema.py

@@ -9,6 +9,7 @@ import attrs
 import sqlalchemy as sa
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 
+from datachain.lib.file import File
 from datachain.sql.types import JSON, Boolean, DateTime, Int64, SQLType, String
 
 if TYPE_CHECKING:
@@ -97,11 +98,11 @@ class Object(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["file__source"])
-        uid = catalog._get_row_uid(file_signals(row))
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            client.download(uid, callback=cb)
-        with client.open_object(uid, use_cache=cache, cb=cb) as f:
+            client.download(file, callback=cb)
+        with client.open_object(file, use_cache=cache, cb=cb) as f:
             return self.reader(f)
 
     async def get_value_async(
@@ -114,12 +115,12 @@ class Object(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["file__source"])
-        uid = catalog._get_row_uid(file_signals(row))
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            await client._download(uid, callback=cb)
+            await client._download(file, callback=cb)
         obj = await mapper.to_thread(
-            functools.partial(client.open_object, uid, use_cache=cache, cb=cb)
+            functools.partial(client.open_object, file, use_cache=cache, cb=cb)
         )
         with obj:
             return await mapper.to_thread(self.reader, obj)
@@ -140,11 +141,11 @@ class Stream(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["file__source"])
-        uid = catalog._get_row_uid(file_signals(row))
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            client.download(uid, callback=cb)
-        return client.open_object(uid, use_cache=cache, cb=cb)
+            client.download(file, callback=cb)
+        return client.open_object(file, use_cache=cache, cb=cb)
 
     async def get_value_async(
         self,
@@ -156,12 +157,12 @@ class Stream(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["file__source"])
-        uid = catalog._get_row_uid(file_signals(row))
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            await client._download(uid, callback=cb)
+            await client._download(file, callback=cb)
         return await mapper.to_thread(
-            functools.partial(client.open_object, uid, use_cache=cache, cb=cb)
+            functools.partial(client.open_object, file, use_cache=cache, cb=cb)
         )
 
 
@@ -189,10 +190,10 @@ class LocalFilename(UDFParameter):
             # If the glob pattern is specified and the row filename
             # does not match it, then return None
             return None
-        client = catalog.get_client(row["file__source"])
-        uid = catalog._get_row_uid(file_signals(row))
-        client.download(uid, callback=cb)
-        return client.cache.get_path(uid)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
+        client.download(file, callback=cb)
+        return client.cache.get_path(file)
 
     async def get_value_async(
         self,
@@ -208,10 +209,10 @@ class LocalFilename(UDFParameter):
             # If the glob pattern is specified and the row filename
             # does not match it, then return None
             return None
-        client = catalog.get_client(row["file__source"])
-        uid = catalog._get_row_uid(file_signals(row))
-        await client._download(uid, callback=cb)
-        return client.cache.get_path(uid)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
+        await client._download(file, callback=cb)
+        return client.cache.get_path(file)
 
 
 UDFParamSpec = Union[str, Column, UDFParameter]