datachain 0.30.5__py3-none-any.whl → 0.30.7__py3-none-any.whl

datachain/cli/commands/datasets.py

@@ -1,30 +1,41 @@
 import sys
-from typing import TYPE_CHECKING, Optional
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING, Optional, Union
 
 from tabulate import tabulate
 
-if TYPE_CHECKING:
-    from datachain.catalog import Catalog
-
+from datachain import semver
 from datachain.catalog import is_namespace_local
 from datachain.cli.utils import determine_flavors
 from datachain.config import Config
 from datachain.error import DataChainError, DatasetNotFoundError
 from datachain.studio import list_datasets as list_datasets_studio
 
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def group_dataset_versions(
+    datasets: Iterable[tuple[str, str]], latest_only=True
+) -> dict[str, Union[str, list[str]]]:
+    grouped: dict[str, list[tuple[int, int, int]]] = {}
 
-def group_dataset_versions(datasets, latest_only=True):
-    grouped = {}
     # Sort to ensure groupby works as expected
     # (groupby expects consecutive items with the same key)
     for name, version in sorted(datasets):
-        grouped.setdefault(name, []).append(version)
+        grouped.setdefault(name, []).append(semver.parse(version))
 
     if latest_only:
         # For each dataset name, pick the highest version.
-        return {name: max(versions) for name, versions in grouped.items()}
+        return {
+            name: semver.create(*(max(versions))) for name, versions in grouped.items()
+        }
+
     # For each dataset name, return a sorted list of unique versions.
-    return {name: sorted(set(versions)) for name, versions in grouped.items()}
+    return {
+        name: [semver.create(*v) for v in sorted(set(versions))]
+        for name, versions in grouped.items()
+    }
 
 
 def list_datasets(
@@ -35,7 +46,7 @@ def list_datasets(
     team: Optional[str] = None,
     latest_only: bool = True,
     name: Optional[str] = None,
-):
+) -> None:
     token = Config().read().get("studio", {}).get("token")
     all, local, studio = determine_flavors(studio, local, all, token)
     if name:
@@ -95,27 +106,31 @@ def list_datasets(
     print(tabulate(rows, headers="keys"))
 
 
-def list_datasets_local(catalog: "Catalog", name: Optional[str] = None):
+def list_datasets_local(
+    catalog: "Catalog", name: Optional[str] = None
+) -> Iterator[tuple[str, str]]:
     if name:
         yield from list_datasets_local_versions(catalog, name)
         return
 
     for d in catalog.ls_datasets():
         for v in d.versions:
-            yield (d.full_name, v.version)
+            yield d.full_name, v.version
 
 
-def list_datasets_local_versions(catalog: "Catalog", name: str):
+def list_datasets_local_versions(
+    catalog: "Catalog", name: str
+) -> Iterator[tuple[str, str]]:
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
     ds = catalog.get_dataset(
         name, namespace_name=namespace_name, project_name=project_name
     )
     for v in ds.versions:
-        yield (name, v.version)
+        yield name, v.version
 
 
-def _datasets_tabulate_row(name, both, local_version, studio_version):
+def _datasets_tabulate_row(name, both, local_version, studio_version) -> dict[str, str]:
     row = {
         "Name": name,
     }
@@ -136,7 +151,7 @@ def rm_dataset(
     force: Optional[bool] = False,
     studio: Optional[bool] = False,
     team: Optional[str] = None,
-):
+) -> None:
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
     if studio:
@@ -166,7 +181,7 @@ def edit_dataset(
     description: Optional[str] = None,
     attrs: Optional[list[str]] = None,
     team: Optional[str] = None,
-):
+) -> None:
     from datachain.lib.dc.utils import is_studio
 
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)

datachain/data_storage/sqlite.py

@@ -37,6 +37,7 @@ from datachain import semver
 from datachain.data_storage import AbstractDBMetastore, AbstractWarehouse
 from datachain.data_storage.db_engine import DatabaseEngine
 from datachain.data_storage.schema import DefaultSchema
+from datachain.data_storage.warehouse import INSERT_BATCH_SIZE
 from datachain.dataset import DatasetRecord, StorageURI
 from datachain.error import DataChainError, OutdatedDatabaseSchemaError
 from datachain.namespace import Namespace
@@ -44,7 +45,7 @@ from datachain.project import Project
 from datachain.sql.sqlite import create_user_defined_sql_functions, sqlite_dialect
 from datachain.sql.sqlite.base import load_usearch_extension
 from datachain.sql.types import SQLType
-from datachain.utils import DataChainDir, batched_it
+from datachain.utils import DataChainDir, batched, batched_it
 
 if TYPE_CHECKING:
     from sqlalchemy.dialects.sqlite import Insert
@@ -712,19 +713,21 @@ class SQLiteWarehouse(AbstractWarehouse):
     def prepare_entries(self, entries: "Iterable[File]") -> Iterable[dict[str, Any]]:
         return (e.model_dump() for e in entries)
 
-    def insert_rows(self, table: Table, rows: Iterable[dict[str, Any]]) -> None:
-        rows = list(rows)
-        if not rows:
-            return
-
-        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_rows(
+        self,
+        table: Table,
+        rows: Iterable[dict[str, Any]],
+        batch_size: int = INSERT_BATCH_SIZE,
+    ) -> None:
+        for row_chunk in batched(rows, batch_size):
+            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 row_chunk[0]}),
+                    row_chunk,
+                    conn=conn,
+                )
 
     def insert_dataset_rows(self, df, dataset: DatasetRecord, version: str) -> int:
         dr = self.dataset_rows(dataset, version)
@@ -797,7 +800,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             .limit(None)
         )
 
-        for batch in batched_it(ids, 10_000):
+        for batch in batched_it(ids, INSERT_BATCH_SIZE):
             batch_ids = [row[0] for row in batch]
             select_q._where_criteria = (col_id.in_(batch_ids),)
             q = table.insert().from_select(list(select_q.selected_columns), select_q)

datachain/data_storage/warehouse.py

@@ -43,6 +43,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger("datachain")
 
 SELECT_BATCH_SIZE = 100_000  # number of rows to fetch at a time
+INSERT_BATCH_SIZE = 10_000  # number of rows to insert at a time
 
 
 class AbstractWarehouse(ABC, Serializable):
@@ -415,7 +416,12 @@ class AbstractWarehouse(ABC, Serializable):
         """Convert File entries so they can be passed on to `insert_rows()`"""
 
     @abstractmethod
-    def insert_rows(self, table: sa.Table, rows: Iterable[dict[str, Any]]) -> None:
+    def insert_rows(
+        self,
+        table: sa.Table,
+        rows: Iterable[dict[str, Any]],
+        batch_size: int = INSERT_BATCH_SIZE,
+    ) -> None:
         """Does batch inserts of any kind of rows into table"""
 
     def insert_rows_done(self, table: sa.Table) -> None:

datachain/delta.py

@@ -4,7 +4,7 @@ from functools import wraps
 from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
 
 import datachain
-from datachain.dataset import DatasetDependency
+from datachain.dataset import DatasetDependency, DatasetRecord
 from datachain.error import DatasetNotFoundError
 from datachain.project import Project
 
@@ -30,9 +30,10 @@ def delta_disabled(
 
     @wraps(method)
     def _inner(self: T, *args: "P.args", **kwargs: "P.kwargs") -> T:
-        if self.delta:
+        if self.delta and not self._delta_unsafe:
             raise NotImplementedError(
-                f"Delta update cannot be used with {method.__name__}"
+                f"Cannot use {method.__name__} with delta datasets - may cause"
+                " inconsistency. Use delta_unsafe flag to allow this operation."
             )
         return method(self, *args, **kwargs)
 
@@ -124,10 +125,19 @@ def _get_retry_chain(
     # Subtract also diff chain since some items might be picked
     # up by `delta=True` itself (e.g. records got modified AND are missing in the
     # result dataset atm)
-    return retry_chain.subtract(diff_chain, on=on) if retry_chain else None
+    on = [on] if isinstance(on, str) else on
+
+    return (
+        retry_chain.diff(
+            diff_chain, on=on, added=True, same=True, modified=False, deleted=False
+        ).distinct(*on)
+        if retry_chain
+        else None
+    )
 
 
 def _get_source_info(
+    source_ds: DatasetRecord,
     name: str,
     namespace_name: str,
     project_name: str,
@@ -154,25 +164,23 @@ def _get_source_info(
         indirect=False,
     )
 
-    dep = dependencies[0]
-    if not dep:
+    source_ds_dep = next((d for d in dependencies if d.name == source_ds.name), None)
+    if not source_ds_dep:
         # Starting dataset was removed, back off to normal dataset creation
         return None, None, None, None, None
 
-    source_ds_project = catalog.metastore.get_project(dep.project, dep.namespace)
-    source_ds_name = dep.name
-    source_ds_version = dep.version
-    source_ds_latest_version = catalog.get_dataset(
-        source_ds_name,
-        namespace_name=source_ds_project.namespace.name,
-        project_name=source_ds_project.name,
-    ).latest_version
+    # Refresh starting dataset to have new versions if they are created
+    source_ds = catalog.get_dataset(
+        source_ds.name,
+        namespace_name=source_ds.project.namespace.name,
+        project_name=source_ds.project.name,
+    )
 
     return (
-        source_ds_name,
-        source_ds_project,
-        source_ds_version,
-        source_ds_latest_version,
+        source_ds.name,
+        source_ds.project,
+        source_ds_dep.version,
+        source_ds.latest_version,
         dependencies,
     )
 
@@ -244,7 +252,14 @@ def delta_retry_update(
         source_ds_version,
         source_ds_latest_version,
         dependencies,
-    ) = _get_source_info(name, namespace_name, project_name, latest_version, catalog)
+    ) = _get_source_info(
+        dc._query.starting_step.dataset,  # type: ignore[union-attr]
+        name,
+        namespace_name,
+        project_name,
+        latest_version,
+        catalog,
+    )
 
     # If source_ds_name is None, starting dataset was removed
     if source_ds_name is None:
@@ -267,8 +282,9 @@ def delta_retry_update(
     if dependencies:
         dependencies = copy(dependencies)
         dependencies = [d for d in dependencies if d is not None]
+        source_ds_dep = next(d for d in dependencies if d.name == source_ds_name)
         # Update to latest version
-        dependencies[0].version = source_ds_latest_version  # type: ignore[union-attr]
+        source_ds_dep.version = source_ds_latest_version  # type: ignore[union-attr]
 
     # Handle retry functionality if enabled
     if delta_retry:

datachain/lib/dc/database.py

@@ -73,7 +73,7 @@ def to_database(
     table_name: str,
     connection: "ConnectionType",
     *,
-    batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
+    batch_size: int = DEFAULT_DATABASE_BATCH_SIZE,
     on_conflict: Optional[str] = None,
     conflict_columns: Optional[list[str]] = None,
     column_mapping: Optional[dict[str, Optional[str]]] = None,
@@ -124,7 +124,7 @@ def to_database(
                 table.create(conn, checkfirst=True)
 
                 rows_iter = chain._leaf_values()
-                for batch in batched(rows_iter, batch_rows):
+                for batch in batched(rows_iter, batch_size):
                     rows_affected = _process_batch(
                         conn,
                         table,

datachain/lib/dc/datachain.py

@@ -193,6 +193,7 @@ class DataChain:
         self._setup: dict = setup or {}
         self._sys = _sys
         self._delta = False
+        self._delta_unsafe = False
         self._delta_on: Optional[Union[str, Sequence[str]]] = None
         self._delta_result_on: Optional[Union[str, Sequence[str]]] = None
         self._delta_compare: Optional[Union[str, Sequence[str]]] = None
@@ -216,6 +217,7 @@ class DataChain:
         right_on: Optional[Union[str, Sequence[str]]] = None,
         compare: Optional[Union[str, Sequence[str]]] = None,
         delta_retry: Optional[Union[bool, str]] = None,
+        delta_unsafe: bool = False,
     ) -> "Self":
         """Marks this chain as delta, which means special delta process will be
         called on saving dataset for optimization"""
@@ -226,6 +228,7 @@ class DataChain:
         self._delta_result_on = right_on
         self._delta_compare = compare
         self._delta_retry = delta_retry
+        self._delta_unsafe = delta_unsafe
         return self
 
     @property
@@ -238,6 +241,10 @@ class DataChain:
         """Returns True if this chain is ran in "delta" update mode"""
         return self._delta
 
+    @property
+    def delta_unsafe(self) -> bool:
+        return self._delta_unsafe
+
     @property
     def schema(self) -> dict[str, DataType]:
         """Get schema of the chain."""
@@ -328,21 +335,22 @@ class DataChain:
                 right_on=self._delta_result_on,
                 compare=self._delta_compare,
                 delta_retry=self._delta_retry,
+                delta_unsafe=self._delta_unsafe,
             )
 
         return chain
 
     def settings(
         self,
-        cache=None,
-        parallel=None,
-        workers=None,
-        min_task_size=None,
-        prefetch: Optional[int] = None,
-        sys: Optional[bool] = None,
+        cache: Optional[bool] = None,
+        prefetch: Optional[Union[bool, int]] = None,
+        parallel: Optional[Union[bool, int]] = None,
+        workers: Optional[int] = None,
         namespace: Optional[str] = None,
         project: Optional[str] = None,
-        batch_rows: Optional[int] = None,
+        min_task_size: Optional[int] = None,
+        batch_size: Optional[int] = None,
+        sys: Optional[bool] = None,
     ) -> "Self":
         """Change settings for chain.
 
@@ -351,23 +359,23 @@ class DataChain:
 
         Parameters:
             cache : data caching. (default=False)
+            prefetch : number of workers to use for downloading files in advance.
+                      This is enabled by default and uses 2 workers.
+                      To disable prefetching, set it to 0 or False.
             parallel : number of thread for processors. True is a special value to
                 enable all available CPUs. (default=1)
             workers : number of distributed workers. Only for Studio mode. (default=1)
-            min_task_size : minimum number of tasks. (default=1)
-            prefetch : number of workers to use for downloading files in advance.
-                      This is enabled by default and uses 2 workers.
-                      To disable prefetching, set it to 0.
             namespace : namespace name.
             project : project name.
-            batch_rows : row limit per insert to balance speed and memory usage.
+            min_task_size : minimum number of tasks. (default=1)
+            batch_size : row limit per insert to balance speed and memory usage.
                       (default=2000)
 
         Example:
             ```py
             chain = (
                 chain
-                .settings(cache=True, parallel=8, batch_rows=300)
+                .settings(cache=True, parallel=8, batch_size=300)
                 .map(laion=process_webdataset(spec=WDSLaion), params="file")
             )
             ```
@@ -377,14 +385,14 @@ class DataChain:
         settings = copy.copy(self._settings)
         settings.add(
             Settings(
-                cache,
-                parallel,
-                workers,
-                min_task_size,
-                prefetch,
-                namespace,
-                project,
-                batch_rows,
+                cache=cache,
+                prefetch=prefetch,
+                parallel=parallel,
+                workers=workers,
+                namespace=namespace,
+                project=project,
+                min_task_size=min_task_size,
+                batch_size=batch_size,
             )
         )
         return self._evolve(settings=settings, _sys=sys)
@@ -737,7 +745,7 @@ class DataChain:
 
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows),
+                udf_obj.to_udf_wrapper(self._settings.batch_size),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -775,7 +783,7 @@ class DataChain:
             udf_obj.prefetch = prefetch
         return self._evolve(
             query=self._query.generate(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows),
+                udf_obj.to_udf_wrapper(self._settings.batch_size),
                 **self._settings.to_dict(),
             ),
             signal_schema=udf_obj.output,
@@ -911,7 +919,7 @@ class DataChain:
         udf_obj = self._udf_to_obj(Aggregator, func, params, output, signal_map)
         return self._evolve(
             query=self._query.generate(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows),
+                udf_obj.to_udf_wrapper(self._settings.batch_size),
                 partition_by=processed_partition_by,
                 **self._settings.to_dict(),
             ),
@@ -960,7 +968,7 @@ class DataChain:
 
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows, batch=batch),
+                udf_obj.to_udf_wrapper(self._settings.batch_size, batch=batch),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -2306,7 +2314,7 @@ class DataChain:
         table_name: str,
         connection: "ConnectionType",
         *,
-        batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
+        batch_size: int = DEFAULT_DATABASE_BATCH_SIZE,
         on_conflict: Optional[str] = None,
         conflict_columns: Optional[list[str]] = None,
         column_mapping: Optional[dict[str, Optional[str]]] = None,
@@ -2328,7 +2336,7 @@ class DataChain:
                 library. If a DBAPI2 object, only sqlite3 is supported. The user is
                 responsible for engine disposal and connection closure for the
                 SQLAlchemy connectable; str connections are closed automatically.
-            batch_rows: Number of rows to insert per batch for optimal performance.
+            batch_size: Number of rows to insert per batch for optimal performance.
                 Larger batches are faster but use more memory. Default: 10,000.
             on_conflict: Strategy for handling duplicate rows (requires table
                 constraints):
@@ -2409,7 +2417,7 @@ class DataChain:
             self,
             table_name,
             connection,
-            batch_rows=batch_rows,
+            batch_size=batch_size,
             on_conflict=on_conflict,
             conflict_columns=conflict_columns,
             column_mapping=column_mapping,

datachain/lib/dc/datasets.py

@@ -40,6 +40,7 @@ def read_dataset(
     delta_result_on: Optional[Union[str, Sequence[str]]] = None,
     delta_compare: Optional[Union[str, Sequence[str]]] = None,
     delta_retry: Optional[Union[bool, str]] = None,
+    delta_unsafe: bool = False,
     update: bool = False,
 ) -> "DataChain":
     """Get data from a saved Dataset. It returns the chain itself.
@@ -80,6 +81,8 @@ def read_dataset(
         update: If True always checks for newer versions available on Studio, even if
             some version of the dataset exists locally already. If False (default), it
             will only fetch the dataset from Studio if it is not found locally.
+        delta_unsafe: Allow restricted ops in delta: merge, agg, union, group_by,
+            distinct.
 
 
     Example:
@@ -205,6 +208,7 @@ def read_dataset(
             right_on=delta_result_on,
             compare=delta_compare,
             delta_retry=delta_retry,
+            delta_unsafe=delta_unsafe,
         )
 
     return chain

datachain/lib/dc/records.py

@@ -31,7 +31,7 @@ def read_records(
 
     Parameters:
         to_insert : records (or a single record) to insert. Each record is
-                    a dictionary of signals and theirs values.
+                    a dictionary of signals and their values.
         schema : describes chain signals and their corresponding types
 
     Example:
@@ -45,7 +45,6 @@ def read_records(
     """
     from datachain.query.dataset import adjust_outputs, get_col_types
     from datachain.sql.types import SQLType
-    from datachain.utils import batched
 
     from .datasets import read_dataset
 
@@ -96,7 +95,6 @@ def read_records(
         {c.name: c.type for c in columns if isinstance(c.type, SQLType)},
     )
     records = (adjust_outputs(warehouse, record, col_types) for record in to_insert)
-    for chunk in batched(records, READ_RECORDS_BATCH_SIZE):
-        warehouse.insert_rows(table, chunk)
+    warehouse.insert_rows(table, records, batch_size=READ_RECORDS_BATCH_SIZE)
     warehouse.insert_rows_done(table)
     return read_dataset(name=dsr.full_name, session=session, settings=settings)

datachain/lib/dc/storage.py

@@ -43,6 +43,7 @@ def read_storage(
     delta_result_on: Optional[Union[str, Sequence[str]]] = None,
     delta_compare: Optional[Union[str, Sequence[str]]] = None,
     delta_retry: Optional[Union[bool, str]] = None,
+    delta_unsafe: bool = False,
     client_config: Optional[dict] = None,
 ) -> "DataChain":
     """Get data from storage(s) as a list of file with all file attributes.
@@ -77,6 +78,9 @@ def read_storage(
               (error mode)
             - True: Reprocess records missing from the result dataset (missing mode)
             - None: No retry processing (default)
+        delta_unsafe: Allow restricted ops in delta: merge, agg, union, group_by,
+            distinct. Caller must ensure datasets are consistent and not partially
+            updated.
 
     Returns:
         DataChain: A DataChain object containing the file information.
@@ -218,6 +222,7 @@ def read_storage(
             right_on=delta_result_on,
             compare=delta_compare,
             delta_retry=delta_retry,
+            delta_unsafe=delta_unsafe,
         )
 
     return storage_chain