datachain 0.7.11__tar.gz → 0.8.0__tar.gz

{datachain-0.7.11 → datachain-0.8.0}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.8.2'
+    rev: 'v0.8.3'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.7.11/src/datachain.egg-info → datachain-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.7.11
+Version: 0.8.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -46,6 +46,7 @@ Requires-Dist: iterative-telemetry>=0.0.9
 Requires-Dist: platformdirs
 Requires-Dist: dvc-studio-client<1,>=0.21
 Requires-Dist: tabulate
+Requires-Dist: websockets
 Provides-Extra: docs
 Requires-Dist: mkdocs>=1.5.2; extra == "docs"
 Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
@@ -98,7 +99,7 @@ Requires-Dist: unstructured[pdf]; extra == "examples"
 Requires-Dist: pdfplumber==0.11.4; extra == "examples"
 Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
 Requires-Dist: onnx==1.16.1; extra == "examples"
-Requires-Dist: ultralytics==8.3.48; extra == "examples"
+Requires-Dist: ultralytics==8.3.50; extra == "examples"
 
 ================
 |logo| DataChain

datachain-0.8.0/examples/get_started/json-csv-reader.py

@@ -0,0 +1,82 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from datachain import C, DataChain
+from datachain.lib.data_model import ModelStore
+from datachain.lib.meta_formats import gen_datamodel_code
+
+
+# Sample model for static JSON model
+class LicenseModel(BaseModel):
+    url: str
+    id: int
+    name: str
+
+
+LicenseFeature = ModelStore.register(LicenseModel)
+
+
+# Sample model for static CSV model
+class ChatDialog(BaseModel):
+    id: Optional[int] = None
+    count: Optional[int] = None
+    sender: Optional[str] = None
+    text: Optional[str] = None
+
+
+ChatFeature = ModelStore.register(ChatDialog)
+
+
+def main():
+    # Dynamic JSONl schema from 2 objects
+    uri = "gs://datachain-demo/jsonl/object.jsonl"
+    jsonl_ds = DataChain.from_json(uri, format="jsonl", anon="True")
+    jsonl_ds.show()
+
+    # Dynamic JSON schema from 200 OpenImage json-pairs with validation errors
+    uri = "gs://datachain-demo/openimages-v6-test-jsonpairs/*json"
+    schema_uri = (
+        "gs://datachain-demo/openimages-v6-test-jsonpairs/08392c290ecc9d2a.json"
+    )
+    json_pairs_ds = DataChain.from_json(
+        uri, schema_from=schema_uri, jmespath="@", model_name="OpenImage", anon="True"
+    )
+    json_pairs_ds.show()
+
+    uri = "gs://datachain-demo/coco2017/annotations_captions/"
+
+    # Print JSON schema in Pydantic format from main COCO annotation
+    chain = DataChain.from_storage(uri, anon="True").filter(
+        C("file.path").glob("*.json")
+    )
+    file = next(chain.limit(1).collect("file"))
+    print(gen_datamodel_code(file, jmespath="@", model_name="Coco"))
+
+    # Static JSON schema test parsing 3/7 objects
+    static_json_ds = DataChain.from_json(
+        uri, jmespath="licenses", spec=LicenseFeature, nrows=3, anon="True"
+    )
+    static_json_ds.show()
+
+    # Dynamic JSON schema test parsing 5K objects
+    dynamic_json_ds = DataChain.from_json(uri, jmespath="images", anon="True")
+    print(dynamic_json_ds.to_pandas())
+
+    # Static CSV with header schema test parsing 3.5K objects
+    uri = "gs://datachain-demo/chatbot-csv/"
+    static_csv_ds = DataChain.from_csv(
+        uri, output=ChatDialog, object_name="chat", anon="True"
+    )
+    static_csv_ds.print_schema()
+    static_csv_ds.show()
+
+    # Dynamic CSV with header schema test parsing 3/3M objects
+    uri = "gs://datachain-demo/laion-aesthetics-csv/laion_aesthetics_1024_33M_1.csv"
+    dynamic_csv_ds = DataChain.from_csv(uri, object_name="laion", nrows=3, anon="True")
+    dynamic_csv_ds.print_schema()
+    dynamic_csv_ds.show()
+
+
+if __name__ == "__main__":
+    main()

{datachain-0.7.11 → datachain-0.8.0}/pyproject.toml

@@ -48,7 +48,8 @@ dependencies = [
   "iterative-telemetry>=0.0.9",
   "platformdirs",
   "dvc-studio-client>=0.21,<1",
-  "tabulate"
+  "tabulate",
+  "websockets"
 ]
 
 [project.optional-dependencies]
@@ -111,7 +112,7 @@ examples = [
   "pdfplumber==0.11.4",
   "huggingface_hub[hf_transfer]",
   "onnx==1.16.1",
-  "ultralytics==8.3.48"
+  "ultralytics==8.3.50"
 ]
 
 [project.urls]

{datachain-0.7.11 → datachain-0.8.0}/src/datachain/catalog/catalog.py

@@ -1,7 +1,6 @@
 import io
 import json
 import logging
-import math
 import os
 import os.path
 import posixpath
@@ -13,7 +12,6 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence
 from copy import copy
 from dataclasses import dataclass
 from functools import cached_property, reduce
-from random import shuffle
 from threading import Thread
 from typing import (
     IO,
@@ -58,11 +56,7 @@ from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.remote.studio import StudioClient
 from datachain.sql.types import DateTime, SQLType
-from datachain.utils import (
-    DataChainDir,
-    batched,
-    datachain_paths_join,
-)
+from datachain.utils import DataChainDir, datachain_paths_join
 
 from .datasource import DataSource
 
@@ -90,7 +84,7 @@ QUERY_SCRIPT_INVALID_LAST_STATEMENT_EXIT_CODE = 10
 QUERY_SCRIPT_CANCELED_EXIT_CODE = 11
 
 # dataset pull
-PULL_DATASET_MAX_THREADS = 10
+PULL_DATASET_MAX_THREADS = 5
 PULL_DATASET_CHUNK_TIMEOUT = 3600
 PULL_DATASET_SLEEP_INTERVAL = 0.1  # sleep time while waiting for chunk to be available
 PULL_DATASET_CHECK_STATUS_INTERVAL = 20  # interval to check export status in Studio
@@ -130,6 +124,7 @@ class DatasetRowsFetcher(NodesThreadPool):
         local_ds_version: int,
         schema: dict[str, Union[SQLType, type[SQLType]]],
         max_threads: int = PULL_DATASET_MAX_THREADS,
+        progress_bar=None,
     ):
         super().__init__(max_threads)
         self._check_dependencies()
@@ -142,6 +137,7 @@ class DatasetRowsFetcher(NodesThreadPool):
         self.schema = schema
         self.last_status_check: Optional[float] = None
         self.studio_client = StudioClient()
+        self.progress_bar = progress_bar
 
     def done_task(self, done):
         for task in done:
@@ -198,6 +194,20 @@ class DatasetRowsFetcher(NodesThreadPool):
         for c in [c for c, t in self.schema.items() if t == DateTime]:
             df[c] = pd.to_datetime(df[c], unit="s")
 
+        # id will be autogenerated in DB
+        return df.drop("sys__id", axis=1)
+
+    def get_parquet_content(self, url: str):
+        while True:
+            if self.should_check_for_status():
+                self.check_for_status()
+            r = requests.get(url, timeout=PULL_DATASET_CHUNK_TIMEOUT)
+            if r.status_code == 404:
+                time.sleep(PULL_DATASET_SLEEP_INTERVAL)
+                continue
+            r.raise_for_status()
+            return r.content
+
     def do_task(self, urls):
         import lz4.frame
         import pandas as pd
@@ -207,31 +217,22 @@ class DatasetRowsFetcher(NodesThreadPool):
             local_ds = metastore.get_dataset(self.local_ds_name)
 
             urls = list(urls)
-            while urls:
-                for url in urls:
-                    if self.should_check_for_status():
-                        self.check_for_status()
-
-                    r = requests.get(url, timeout=PULL_DATASET_CHUNK_TIMEOUT)
-                    if r.status_code == 404:
-                        time.sleep(PULL_DATASET_SLEEP_INTERVAL)
-                        # moving to the next url
-                        continue
-
-                    r.raise_for_status()
-
-                    df = pd.read_parquet(io.BytesIO(lz4.frame.decompress(r.content)))
 
-                    self.fix_columns(df)
+            for url in urls:
+                if self.should_check_for_status():
+                    self.check_for_status()
 
-                    # id will be autogenerated in DB
-                    df = df.drop("sys__id", axis=1)
+                df = pd.read_parquet(
+                    io.BytesIO(lz4.frame.decompress(self.get_parquet_content(url)))
+                )
+                df = self.fix_columns(df)
 
-                    inserted = warehouse.insert_dataset_rows(
-                        df, local_ds, self.local_ds_version
-                    )
-                    self.increase_counter(inserted)  # type: ignore [arg-type]
-                    urls.remove(url)
+                inserted = warehouse.insert_dataset_rows(
+                    df, local_ds, self.local_ds_version
+                )
+                self.increase_counter(inserted)  # type: ignore [arg-type]
+                # sometimes progress bar doesn't get updated so manually updating it
+                self.update_progress_bar(self.progress_bar)
 
 
 @dataclass
@@ -1291,13 +1292,13 @@ class Catalog:
         for source in data_sources:  # type: ignore [union-attr]
             yield source, source.ls(fields)
 
-    def pull_dataset(  # noqa: PLR0915
+    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[int] = None,
-        no_cp: bool = False,
+        cp: bool = False,
         force: bool = False,
         edatachain: bool = False,
         edatachain_file: Optional[str] = None,
@@ -1305,7 +1306,7 @@ class Catalog:
         client_config=None,
     ) -> None:
         def _instantiate(ds_uri: str) -> None:
-            if no_cp:
+            if not cp:
                 return
             assert output
             self.cp(
@@ -1318,7 +1319,7 @@ class Catalog:
             )
             print(f"Dataset {ds_uri} instantiated locally to {output}")
 
-        if not output and not no_cp:
+        if cp and not output:
             raise ValueError("Please provide output directory for instantiation")
 
         studio_client = StudioClient()
@@ -1417,12 +1418,26 @@ class Catalog:
         signed_urls = export_response.data
 
         if signed_urls:
-            shuffle(signed_urls)
-
             with (
                 self.metastore.clone() as metastore,
                 self.warehouse.clone() as warehouse,
             ):
+
+                def batch(urls):
+                    """
+                    Batching urls in a way that fetching is most efficient as
+                    urls with lower id will be created first. Because that, we
+                    are making sure all threads are pulling most recent urls
+                    from beginning
+                    """
+                    res = [[] for i in range(PULL_DATASET_MAX_THREADS)]
+                    current_worker = 0
+                    for url in signed_urls:
+                        res[current_worker].append(url)
+                        current_worker = (current_worker + 1) % PULL_DATASET_MAX_THREADS
+
+                    return res
+
                 rows_fetcher = DatasetRowsFetcher(
                     metastore,
                     warehouse,
@@ -1431,14 +1446,11 @@ class Catalog:
                     local_ds_name,
                     local_ds_version,
                     schema,
+                    progress_bar=dataset_save_progress_bar,
                 )
                 try:
                     rows_fetcher.run(
-                        batched(
-                            signed_urls,
-                            math.ceil(len(signed_urls) / PULL_DATASET_MAX_THREADS),
-                        ),
-                        dataset_save_progress_bar,
+                        iter(batch(signed_urls)), dataset_save_progress_bar
                     )
                 except:
                     self.remove_dataset(local_ds_name, local_ds_version)

{datachain-0.7.11 → datachain-0.8.0}/src/datachain/cli.py

@@ -294,6 +294,28 @@ def add_studio_parser(subparsers, parent_parser) -> None:
         help="Python package requirement. Can be specified multiple times.",
     )
 
+    studio_cancel_help = "Cancel a job in Studio"
+    studio_cancel_description = "This command cancels a job in Studio."
+
+    studio_cancel_parser = studio_subparser.add_parser(
+        "cancel",
+        parents=[parent_parser],
+        description=studio_cancel_description,
+        help=studio_cancel_help,
+    )
+
+    studio_cancel_parser.add_argument(
+        "job_id",
+        action="store",
+        help="The job ID to cancel.",
+    )
+    studio_cancel_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="The team to cancel a job for. By default, it will use team from config.",
+    )
+
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
     try:
@@ -457,10 +479,10 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Copy directories recursively",
     )
     parse_pull.add_argument(
-        "--no-cp",
+        "--cp",
         default=False,
         action="store_true",
-        help="Do not copy files, just pull a remote dataset into local DB",
+        help="Copy actual files after pulling remote dataset into local DB",
     )
     parse_pull.add_argument(
         "--edatachain",
@@ -1300,7 +1322,7 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
                     args.output,
                     local_ds_name=args.local_name,
                     local_ds_version=args.local_version,
-                    no_cp=args.no_cp,
+                    cp=args.cp,
                     force=bool(args.force),
                     edatachain=args.edatachain,
                     edatachain_file=args.edatachain_file,

{datachain-0.7.11 → datachain-0.8.0}/src/datachain/data_storage/sqlite.py

@@ -209,10 +209,12 @@ class SQLiteDatabaseEngine(DatabaseEngine):
 
     @retry_sqlite_locks
     def executemany(
-        self, query, params, cursor: Optional[sqlite3.Cursor] = None
+        self, query, params, cursor: Optional[sqlite3.Cursor] = None, conn=None
     ) -> sqlite3.Cursor:
         if cursor:
             return cursor.executemany(self.compile(query).string, params)
+        if conn:
+            return conn.executemany(self.compile(query).string, params)
         return self.db.executemany(self.compile(query).string, params)
 
     @retry_sqlite_locks
@@ -222,7 +224,14 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         return self.db.execute(sql, parameters)
 
     def insert_dataframe(self, table_name: str, df) -> int:
-        return df.to_sql(table_name, self.db, if_exists="append", index=False)
+        return df.to_sql(
+            table_name,
+            self.db,
+            if_exists="append",
+            index=False,
+            method="multi",
+            chunksize=1000,
+        )
 
     def cursor(self, factory=None):
         if factory is None:
@@ -545,10 +554,15 @@ class SQLiteWarehouse(AbstractWarehouse):
         rows = list(rows)
         if not rows:
             return
-        self.db.executemany(
-            table.insert().values({f: bindparam(f) for f in rows[0]}),
-            rows,
-        )
+
+        with self.db.transaction() as conn:
+            # transactions speeds up inserts significantly as there is no separate
+            # transaction created for each insert row
+            self.db.executemany(
+                table.insert().values({f: bindparam(f) for f in rows[0]}),
+                rows,
+                conn=conn,
+            )
 
     def insert_dataset_rows(self, df, dataset: DatasetRecord, version: int) -> int:
         dr = self.dataset_rows(dataset, version)