datachain 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

datachain/catalog/catalog.py

@@ -156,8 +156,6 @@ class QueryResult(NamedTuple):
     dataset: Optional[DatasetRecord]
     version: Optional[int]
     output: str
-    preview: Optional[list[dict]]
-    metrics: dict[str, Any]
 
 
 class DatasetRowsFetcher(NodesThreadPool):
@@ -1020,20 +1018,6 @@ class Catalog:
 
         return node_groups
 
-    def unlist_source(self, uri: StorageURI) -> None:
-        self.metastore.clone(uri=uri).mark_storage_not_indexed(uri)
-
-    def storage_stats(self, uri: StorageURI) -> Optional[DatasetStats]:
-        """
-        Returns tuple with storage stats: total number of rows and total dataset size.
-        """
-        partial_path = self.metastore.get_last_partial_path(uri)
-        if partial_path is None:
-            return None
-        dataset = self.get_dataset(Storage.dataset_name(uri, partial_path))
-
-        return self.dataset_stats(dataset.name, dataset.latest_version)
-
     def create_dataset(
         self,
         name: str,
@@ -1297,19 +1281,6 @@ class Catalog:
 
         return self.get_dataset(name)
 
-    def register_new_dataset(
-        self,
-        source_dataset: DatasetRecord,
-        source_version: int,
-        target_name: str,
-    ) -> DatasetRecord:
-        target_dataset = self.metastore.create_dataset(
-            target_name,
-            query_script=source_dataset.query_script,
-            schema=source_dataset.serialized_schema,
-        )
-        return self.register_dataset(source_dataset, source_version, target_dataset, 1)
-
     def register_dataset(
         self,
         dataset: DatasetRecord,
@@ -1422,17 +1393,18 @@ class Catalog:
 
         return direct_dependencies
 
-    def ls_datasets(self) -> Iterator[DatasetRecord]:
+    def ls_datasets(self, include_listing: bool = False) -> Iterator[DatasetRecord]:
         datasets = self.metastore.list_datasets()
         for d in datasets:
-            if not d.is_bucket_listing:
+            if not d.is_bucket_listing or include_listing:
                 yield d
 
     def list_datasets_versions(
         self,
+        include_listing: bool = False,
     ) -> Iterator[tuple[DatasetRecord, "DatasetVersion", Optional["Job"]]]:
         """Iterate over all dataset versions with related jobs."""
-        datasets = list(self.ls_datasets())
+        datasets = list(self.ls_datasets(include_listing=include_listing))
 
         # preselect dataset versions jobs from db to avoid multiple queries
         jobs_ids: set[str] = {
@@ -1560,17 +1532,8 @@ class Catalog:
         version = self.get_dataset(dataset_name).get_version(dataset_version)
 
         file_signals_values = {}
-        file_schemas = {}
-        # TODO: To remove after we properly fix deserialization
-        for signal, type_name in version.feature_schema.items():
-            from datachain.lib.model_store import ModelStore
-
-            type_name_parsed, v = ModelStore.parse_name_version(type_name)
-            fr = ModelStore.get(type_name_parsed, v)
-            if fr and issubclass(fr, File):
-                file_schemas[signal] = type_name
 
-        schema = SignalSchema.deserialize(file_schemas)
+        schema = SignalSchema.deserialize(version.feature_schema)
         for file_signals in schema.get_signals(File):
             prefix = file_signals.replace(".", DEFAULT_DELIMITER) + DEFAULT_DELIMITER
             file_signals_values[file_signals] = {
@@ -1641,15 +1604,6 @@ class Catalog:
         for source in data_sources:  # type: ignore [union-attr]
             yield source, source.ls(fields)
 
-    def ls_storage_uris(self) -> Iterator[str]:
-        yield from self.metastore.get_all_storage_uris()
-
-    def get_storage(self, uri: StorageURI) -> Storage:
-        return self.metastore.get_storage(uri)
-
-    def ls_storages(self) -> list[Storage]:
-        return self.metastore.list_storages()
-
     def pull_dataset(
         self,
         dataset_uri: str,
@@ -1883,10 +1837,6 @@ class Catalog:
         envs: Optional[Mapping[str, str]] = None,
         python_executable: Optional[str] = None,
         save: bool = False,
-        save_as: Optional[str] = None,
-        preview_limit: int = 10,
-        preview_offset: int = 0,
-        preview_columns: Optional[list[str]] = None,
         capture_output: bool = True,
         output_hook: Callable[[str], None] = noop,
         params: Optional[dict[str, str]] = None,
@@ -1914,9 +1864,8 @@ class Catalog:
                 C.size > 1000
             )
         """
-        from datachain.query.dataset import ExecutionResult
 
-        feature_file = tempfile.NamedTemporaryFile(
+        feature_file = tempfile.NamedTemporaryFile(  # noqa: SIM115
             dir=os.getcwd(), suffix=".py", delete=False
         )
         _, feature_module = os.path.split(feature_file.name)
@@ -1931,11 +1880,7 @@ class Catalog:
                 feature_module,
                 output_hook,
                 params,
-                preview_columns,
-                preview_limit,
-                preview_offset,
                 save,
-                save_as,
                 job_id,
             )
         finally:
@@ -1964,25 +1909,18 @@ class Catalog:
             )
 
         try:
-            response = json.loads(response_text)
+            result = json.loads(response_text)
         except ValueError:
-            response = {}
-        exec_result = ExecutionResult(**response)
+            result = None
 
         dataset: Optional[DatasetRecord] = None
         version: Optional[int] = None
-        if save or save_as:
+        if save:
             dataset, version = self.save_result(
-                query_script, exec_result, output, version, job_id
+                query_script, result, output, version, job_id
             )
 
-        return QueryResult(
-            dataset=dataset,
-            version=version,
-            output=output,
-            preview=exec_result.preview,
-            metrics=exec_result.metrics,
-        )
+        return QueryResult(dataset=dataset, version=version, output=output)
 
     def run_query(
         self,
@@ -1994,11 +1932,7 @@ class Catalog:
         feature_module: str,
         output_hook: Callable[[str], None],
         params: Optional[dict[str, str]],
-        preview_columns: Optional[list[str]],
-        preview_limit: int,
-        preview_offset: int,
         save: bool,
-        save_as: Optional[str],
         job_id: Optional[str],
     ) -> tuple[list[str], subprocess.Popen, str]:
         try:
@@ -2013,10 +1947,6 @@ class Catalog:
             raise QueryScriptCompileError(
                 f"Query script failed to compile, reason: {exc}"
             ) from exc
-        if save_as and save_as.startswith(QUERY_DATASET_PREFIX):
-            raise ValueError(
-                f"Cannot use {QUERY_DATASET_PREFIX} prefix for dataset name"
-            )
         r, w = os.pipe()
         if os.name == "nt":
             import msvcrt
@@ -2039,15 +1969,7 @@ class Catalog:
             {
                 "DATACHAIN_QUERY_PARAMS": json.dumps(params or {}),
                 "PYTHONPATH": os.getcwd(),  # For local imports
-                "DATACHAIN_QUERY_PREVIEW_ARGS": json.dumps(
-                    {
-                        "limit": preview_limit,
-                        "offset": preview_offset,
-                        "columns": preview_columns,
-                    }
-                ),
                 "DATACHAIN_QUERY_SAVE": "1" if save else "",
-                "DATACHAIN_QUERY_SAVE_AS": save_as or "",
                 "PYTHONUNBUFFERED": "1",
                 "DATACHAIN_OUTPUT_FD": str(handle),
                 "DATACHAIN_JOB_ID": job_id or "",
@@ -2077,12 +1999,12 @@ class Catalog:
         return lines, proc, response_text
 
     def save_result(self, query_script, exec_result, output, version, job_id):
-        if not exec_result.dataset:
+        if not exec_result:
             raise QueryScriptDatasetNotFound(
                 "No dataset found after running Query script",
                 output=output,
             )
-        name, version = exec_result.dataset
+        name, version = exec_result
         # finding returning dataset
         try:
             dataset = self.get_dataset(name)

datachain/cli.py

@@ -14,6 +14,7 @@ import shtab
 
 from datachain import utils
 from datachain.cli_utils import BooleanOptionalAction, CommaSeparatedArgs, KeyValueArgs
+from datachain.lib.dc import DataChain
 from datachain.utils import DataChainDir
 
 if TYPE_CHECKING:
@@ -472,9 +473,6 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
     query_parser.add_argument(
         "script", metavar="<script.py>", type=str, help="Filepath for script"
     )
-    query_parser.add_argument(
-        "dataset_name", nargs="?", type=str, help="Save result dataset as"
-    )
     query_parser.add_argument(
         "--parallel",
         nargs="?",
@@ -487,7 +485,6 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
             "N defaults to the CPU count."
         ),
     )
-    add_show_args(query_parser)
     query_parser.add_argument(
         "-p",
         "--param",
@@ -619,18 +616,6 @@ def _ls_urls_flat(
                 raise FileNotFoundError(f"No such file or directory: {source}")
 
 
-def ls_indexed_storages(catalog: "Catalog", long: bool = False) -> Iterator[str]:
-    from datachain.node import long_line_str
-
-    storage_uris = catalog.ls_storage_uris()
-    if long:
-        for uri in storage_uris:
-            # TODO: add Storage.created so it can be used here
-            yield long_line_str(uri, None, "")
-    else:
-        yield from storage_uris
-
-
 def ls_local(
     sources,
     long: bool = False,
@@ -661,8 +646,9 @@ def ls_local(
                 for entry in entries:
                     print(format_ls_entry(entry))
     else:
-        for entry in ls_indexed_storages(catalog, long=long):
-            print(format_ls_entry(entry))
+        chain = DataChain.listings()
+        for ls in chain.collect("listing"):
+            print(format_ls_entry(f"{ls.uri}@v{ls.version}"))  # type: ignore[union-attr]
 
 
 def format_ls_entry(entry: str) -> str:
@@ -813,16 +799,10 @@ def show(
 def query(
     catalog: "Catalog",
     script: str,
-    dataset_name: Optional[str] = None,
     parallel: Optional[int] = None,
-    limit: int = 10,
-    offset: int = 0,
-    columns: Optional[list[str]] = None,
-    no_collapse: bool = False,
     params: Optional[dict[str, str]] = None,
 ) -> None:
     from datachain.data_storage import JobQueryType, JobStatus
-    from datachain.utils import show_records
 
     with open(script, encoding="utf-8") as f:
         script_content = f.read()
@@ -843,13 +823,9 @@ def query(
     )
 
     try:
-        result = catalog.query(
+        catalog.query(
             script_content,
             python_executable=python_executable,
-            save_as=dataset_name,
-            preview_limit=limit,
-            preview_offset=offset,
-            preview_columns=columns,
             capture_output=False,
             params=params,
             job_id=job_id,
@@ -864,10 +840,7 @@ def query(
             error_stack=error_stack,
         )
         raise
-
-    catalog.metastore.set_job_status(job_id, JobStatus.COMPLETE, metrics=result.metrics)
-
-    show_records(result.preview, collapse_columns=not no_collapse)
+    catalog.metastore.set_job_status(job_id, JobStatus.COMPLETE)
 
 
 def clear_cache(catalog: "Catalog"):
@@ -1042,12 +1015,7 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
             query(
                 catalog,
                 args.script,
-                dataset_name=args.dataset_name,
                 parallel=args.parallel,
-                limit=args.limit,
-                offset=args.offset,
-                columns=args.columns,
-                no_collapse=args.no_collapse,
                 params=args.param,
             )
         elif args.command == "apply-udf":

datachain/client/fsspec.py

@@ -87,6 +87,7 @@ class Client(ABC):
     def get_implementation(url: str) -> type["Client"]:
         from .azure import AzureClient
         from .gcs import GCSClient
+        from .hf import HfClient
         from .local import FileClient
         from .s3 import ClientS3
 
@@ -104,6 +105,8 @@ class Client(ABC):
             return AzureClient
         if protocol == FileClient.protocol:
             return FileClient
+        if protocol == HfClient.protocol:
+            return HfClient
 
         raise NotImplementedError(f"Unsupported protocol: {protocol}")
 

datachain/client/hf.py

@@ -0,0 +1,47 @@
+import os
+import posixpath
+from typing import Any, cast
+
+from huggingface_hub import HfFileSystem
+
+from datachain.lib.file import File
+from datachain.node import Entry
+
+from .fsspec import Client
+
+
+class HfClient(Client):
+    FS_CLASS = HfFileSystem
+    PREFIX = "hf://"
+    protocol = "hf"
+
+    @classmethod
+    def create_fs(cls, **kwargs) -> HfFileSystem:
+        if os.environ.get("HF_TOKEN"):
+            kwargs["token"] = os.environ["HF_TOKEN"]
+
+        return cast(HfFileSystem, super().create_fs(**kwargs))
+
+    def convert_info(self, v: dict[str, Any], path: str) -> Entry:
+        return Entry.from_file(
+            path=path,
+            size=v["size"],
+            version=v["last_commit"].oid,
+            etag=v.get("blob_id", ""),
+            last_modified=v["last_commit"].date,
+        )
+
+    def info_to_file(self, v: dict[str, Any], path: str) -> File:
+        return File(
+            path=path,
+            size=v["size"],
+            version=v["last_commit"].oid,
+            etag=v.get("blob_id", ""),
+            last_modified=v["last_commit"].date,
+        )
+
+    async def ls_dir(self, path):
+        return self.fs.ls(path, detail=True)
+
+    def rel_path(self, path):
+        return posixpath.relpath(path, self.name)

datachain/data_storage/metastore.py

@@ -167,21 +167,10 @@ class AbstractMetastore(ABC, Serializable):
         This method should be called when index operation is finished.
         """
 
-    @abstractmethod
-    def mark_storage_not_indexed(self, uri: StorageURI) -> None:
-        """
-        Mark storage as not indexed.
-        This method should be called when storage index is deleted.
-        """
-
     @abstractmethod
     def update_last_inserted_at(self, uri: Optional[StorageURI] = None) -> None:
         """Updates last inserted datetime in bucket with current time."""
 
-    @abstractmethod
-    def get_all_storage_uris(self) -> Iterator[StorageURI]:
-        """Returns all storage uris."""
-
     @abstractmethod
     def get_storage(self, uri: StorageURI) -> Storage:
         """
@@ -189,10 +178,6 @@ class AbstractMetastore(ABC, Serializable):
         E.g. if s3 is used as storage this would be s3 bucket data.
         """
 
-    @abstractmethod
-    def list_storages(self) -> list[Storage]:
-        """Returns all storages."""
-
     @abstractmethod
     def mark_storage_pending(self, storage: Storage) -> Storage:
         """Marks storage as pending."""
@@ -324,7 +309,7 @@ class AbstractMetastore(ABC, Serializable):
             self.add_dataset_dependency(
                 source_dataset_name,
                 source_dataset_version,
-                dependency.name,
+                dependency.dataset_name,
                 int(dependency.version),
             )
         else:
@@ -906,11 +891,6 @@ class AbstractDBMetastore(AbstractMetastore):
             self._storages_update().where(s.c.uri == uri).values(**updates)  # type: ignore [attr-defined]
         )
 
-    def get_all_storage_uris(self) -> Iterator[StorageURI]:
-        """Returns all storage uris."""
-        s = self._storages
-        yield from (r[0] for r in self.db.execute(self._storages_select(s.c.uri)))
-
     def get_storage(self, uri: StorageURI, conn=None) -> Storage:
         """
         Gets storage representation from database.
@@ -926,13 +906,6 @@ class AbstractDBMetastore(AbstractMetastore):
 
         return self.storage_class._make(result)
 
-    def list_storages(self) -> list[Storage]:
-        result = self.db.execute(self._storages_select())
-        if not result:
-            return []
-
-        return [self.storage_class._make(r) for r in result]
-
     def mark_storage_pending(self, storage: Storage, conn=None) -> Storage:
         # Update status to pending and dates
         updates = {
@@ -1503,7 +1476,7 @@ class AbstractDBMetastore(AbstractMetastore):
         return self._jobs.update().where(*where)
 
     def _parse_job(self, rows) -> Job:
-        return Job.parse(*rows)
+        return self.job_class.parse(*rows)
 
     def _parse_jobs(self, rows) -> Iterator["Job"]:
         for _, g in groupby(rows, lambda r: r[0]):

datachain/data_storage/sqlite.py

@@ -143,7 +143,9 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             db.execute("PRAGMA synchronous = NORMAL")
             db.execute("PRAGMA case_sensitive_like = ON")
             if os.environ.get("DEBUG_SHOW_SQL_QUERIES"):
-                db.set_trace_callback(print)
+                import sys
+
+                db.set_trace_callback(sys.stderr.write)
 
             load_usearch_extension(db)
 
@@ -515,17 +517,6 @@ class SQLiteMetastore(AbstractDBMetastore):
     def _datasets_dependencies_insert(self) -> "Insert":
         return sqlite.insert(self._datasets_dependencies)
 
-    #
-    # Storages
-    #
-
-    def mark_storage_not_indexed(self, uri: StorageURI) -> None:
-        """
-        Mark storage as not indexed.
-        This method should be called when storage index is deleted.
-        """
-        self.db.execute(self._storages_delete().where(self._storages.c.uri == uri))
-
     #
     # Dataset dependencies
     #

datachain/data_storage/warehouse.py

@@ -218,35 +218,26 @@ class AbstractWarehouse(ABC, Serializable):
         results = None
         offset = 0
         num_yielded = 0
-        try:
-            while True:
-                if limit is not None:
-                    limit -= num_yielded
-                    if limit == 0:
-                        break
-                    if limit < page_size:
-                        paginated_query = paginated_query.limit(None).limit(limit)
-
-                results = self.dataset_rows_select(paginated_query.offset(offset))
-
-                processed = False
-                for row in results:
-                    processed = True
-                    yield row
-                    num_yielded += 1
-
-                if not processed:
-                    break  # no more results
-                offset += page_size
-        finally:
-            # https://www2.sqlite.org/cvstrac/wiki?p=DatabaseIsLocked (SELECT not
-            # finalized or reset) to prevent database table is locked error when an
-            # exception is raised in the middle of processing the results (e.g.
-            # https://github.com/iterative/dvcx/issues/924). Connections close
-            # apparently is not enough in some cases, at least on sqlite
-            # https://www.sqlite.org/c3ref/close.html
-            if results and hasattr(results, "close"):
-                results.close()
+
+        while True:
+            if limit is not None:
+                limit -= num_yielded
+                if limit == 0:
+                    break
+                if limit < page_size:
+                    paginated_query = paginated_query.limit(None).limit(limit)
+
+            results = self.dataset_rows_select(paginated_query.offset(offset))
+
+            processed = False
+            for row in results:
+                processed = True
+                yield row
+                num_yielded += 1
+
+            if not processed:
+                break  # no more results
+            offset += page_size
 
     #
     # Table Name Internal Functions

datachain/dataset.py

@@ -11,8 +11,6 @@ from typing import (
 )
 from urllib.parse import urlparse
 
-from dateutil.parser import isoparse
-
 from datachain.client import Client
 from datachain.sql.types import NAME_TYPES_MAPPING, SQLType
 
@@ -25,6 +23,7 @@ DD = TypeVar("DD", bound="DatasetDependency")
 
 DATASET_PREFIX = "ds://"
 QUERY_DATASET_PREFIX = "ds_query_"
+LISTING_PREFIX = "lst__"
 
 
 def parse_dataset_uri(uri: str) -> tuple[str, Optional[int]]:
@@ -72,11 +71,22 @@ class DatasetDependencyType:
 class DatasetDependency:
     id: int
     type: str
-    name: str  # when the type is STORAGE, this is actually StorageURI
-    version: str  # string until we'll have proper bucket listing versions
+    name: str
+    version: str  # TODO change to int
     created_at: datetime
     dependencies: list[Optional["DatasetDependency"]]
 
+    @property
+    def dataset_name(self) -> str:
+        """Returns clean dependency dataset name"""
+        from datachain.lib.listing import parse_listing_uri
+
+        if self.type == DatasetDependencyType.DATASET:
+            return self.name
+
+        list_dataset_name, _, _ = parse_listing_uri(self.name.strip("/"), None, {})
+        return list_dataset_name
+
     @classmethod
     def parse(
         cls: builtins.type[DD],
@@ -91,33 +101,31 @@ class DatasetDependency:
         dataset_version_created_at: Optional[datetime],
         bucket_uri: Optional["StorageURI"],
     ) -> Optional["DatasetDependency"]:
-        if dataset_id:
-            assert dataset_name is not None
-            return cls(
-                id,
-                DatasetDependencyType.DATASET,
-                dataset_name,
-                (
-                    str(dataset_version)  # type: ignore[arg-type]
-                    if dataset_version
-                    else None
-                ),
-                dataset_version_created_at or dataset_created_at,  # type: ignore[arg-type]
-                [],
-            )
-        if bucket_uri:
-            return cls(
-                id,
-                DatasetDependencyType.STORAGE,
-                bucket_uri,
-                bucket_version,  # type: ignore[arg-type]
-                isoparse(bucket_version),  # type: ignore[arg-type]
-                [],
-            )
-        # dependency has been removed
-        # TODO we should introduce flags for removed datasets, instead of
-        # removing them from tables so that we can still have references
-        return None
+        from datachain.lib.listing import is_listing_dataset, listing_uri_from_name
+
+        if not dataset_id:
+            return None
+
+        assert dataset_name is not None
+        dependency_type = DatasetDependencyType.DATASET
+        dependency_name = dataset_name
+
+        if is_listing_dataset(dataset_name):
+            dependency_type = DatasetDependencyType.STORAGE  # type: ignore[arg-type]
+            dependency_name = listing_uri_from_name(dataset_name)
+
+        return cls(
+            id,
+            dependency_type,
+            dependency_name,
+            (
+                str(dataset_version)  # type: ignore[arg-type]
+                if dataset_version
+                else None
+            ),
+            dataset_version_created_at or dataset_created_at,  # type: ignore[arg-type]
+            [],
+        )
 
     @property
     def is_dataset(self) -> bool:
@@ -443,7 +451,11 @@ class DatasetRecord:
         For bucket listing we implicitly create underlying dataset to hold data. This
         method is checking if this is one of those datasets.
         """
-        return Client.is_data_source_uri(self.name)
+        # TODO refactor and maybe remove method in
+        # https://github.com/iterative/datachain/issues/318
+        return Client.is_data_source_uri(self.name) or self.name.startswith(
+            LISTING_PREFIX
+        )
 
     @property
     def versions_values(self) -> list[int]: