Incremental (delta) update (#928)

Incremental (delta) update

Author: ilongin

src/datachain/dataset.py

@@ -107,24 +107,21 @@ def parse(
         dataset_version: Optional[str],
         dataset_version_created_at: Optional[datetime],
     ) -> Optional["DatasetDependency"]:
-        from datachain.client import Client
-        from datachain.lib.listing import is_listing_dataset, listing_uri_from_name
+        from datachain.lib.listing import is_listing_dataset
 
         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, _ = Client.parse_url(listing_uri_from_name(dataset_name))
 
         return cls(
             id,
-            dependency_type,
-            dependency_name,
+            (
+                DatasetDependencyType.STORAGE
+                if is_listing_dataset(dataset_name)
+                else DatasetDependencyType.DATASET
+            ),
+            dataset_name,
             (
                 dataset_version  # type: ignore[arg-type]
                 if dataset_version

src/datachain/delta.py

@@ -0,0 +1,119 @@
+from collections.abc import Sequence
+from copy import copy
+from functools import wraps
+from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union
+
+import datachain
+from datachain.dataset import DatasetDependency
+from datachain.error import DatasetNotFoundError
+
+if TYPE_CHECKING:
+    from typing_extensions import Concatenate, ParamSpec
+
+    from datachain.lib.dc import DataChain
+
+    P = ParamSpec("P")
+
+
+T = TypeVar("T", bound="DataChain")
+
+
+def delta_disabled(
+    method: "Callable[Concatenate[T, P], T]",
+) -> "Callable[Concatenate[T, P], T]":
+    """
+    Decorator for disabling DataChain methods (e.g `.agg()` or `.union()`) to
+    work with delta updates. It throws `NotImplementedError` if chain on which
+    method is called is marked as delta.
+    """
+
+    @wraps(method)
+    def _inner(self: T, *args: "P.args", **kwargs: "P.kwargs") -> T:
+        if self.delta:
+            raise NotImplementedError(
+                f"Delta update cannot be used with {method.__name__}"
+            )
+        return method(self, *args, **kwargs)
+
+    return _inner
+
+
+def _append_steps(dc: "DataChain", other: "DataChain"):
+    """Returns cloned chain with appended steps from other chain.
+    Steps are all those modification methods applied like filters, mappers etc.
+    """
+    dc = dc.clone()
+    dc._query.steps += other._query.steps.copy()
+    dc.signals_schema = other.signals_schema
+    return dc
+
+
+def delta_update(
+    dc: "DataChain",
+    name: str,
+    on: Union[str, Sequence[str]],
+    right_on: Optional[Union[str, Sequence[str]]] = None,
+    compare: Optional[Union[str, Sequence[str]]] = None,
+) -> tuple[Optional["DataChain"], Optional[list[DatasetDependency]], bool]:
+    """
+    Creates new chain that consists of the last version of current delta dataset
+    plus diff from the source with all needed modifications.
+    This way we don't need to re-calculate the whole chain from the source again(
+    apply all the DataChain methods like filters, mappers, generators etc.)
+    but just the diff part which is very important for performance.
+
+    Note that currently delta update works only if there is only one direct dependency.
+    """
+    catalog = dc.session.catalog
+    dc._query.apply_listing_pre_step()
+
+    try:
+        latest_version = catalog.get_dataset(name).latest_version
+    except DatasetNotFoundError:
+        # first creation of delta update dataset
+        return None, None, True
+
+    dependencies = catalog.get_dataset_dependencies(
+        name, latest_version, indirect=False
+    )
+
+    dep = dependencies[0]
+    if not dep:
+        # starting dataset (e.g listing) was removed so we are backing off to normal
+        # dataset creation, as it was created first time
+        return None, None, True
+
+    source_ds_name = dep.name
+    source_ds_version = dep.version
+    source_ds_latest_version = catalog.get_dataset(source_ds_name).latest_version
+    dependencies = copy(dependencies)
+    dependencies = [d for d in dependencies if d is not None]  # filter out removed dep
+    dependencies[0].version = source_ds_latest_version  # type: ignore[union-attr]
+
+    source_dc = datachain.read_dataset(source_ds_name, source_ds_version)
+    source_dc_latest = datachain.read_dataset(source_ds_name, source_ds_latest_version)
+
+    diff = source_dc_latest.compare(source_dc, on=on, compare=compare, deleted=False)
+    # We append all the steps from the original chain to diff, e.g filters, mappers.
+    diff = _append_steps(diff, dc)
+
+    # to avoid re-calculating diff multiple times
+    diff = diff.persist()
+
+    if diff.empty:
+        return None, None, False
+
+    # merging diff and the latest version of dataset
+    delta_chain = (
+        datachain.read_dataset(name, latest_version)
+        .compare(
+            diff,
+            on=right_on or on,
+            added=True,
+            modified=False,
+            deleted=False,
+        )
+        .union(diff)
+    )
+
+    return delta_chain, dependencies, True  # type: ignore[return-value]

src/datachain/diff/__init__.py

@@ -30,7 +30,7 @@ class CompareStatus(str, Enum):
     SAME = "S"
 
 
-def _compare(  # noqa: C901
+def _compare(  # noqa: C901, PLR0912
     left: "DataChain",
     right: "DataChain",
     on: Union[str, Sequence[str]],
@@ -77,14 +77,16 @@ def _to_list(obj: Optional[Union[str, Sequence[str]]]) -> Optional[list[str]]:
     cols_select = list(left.signals_schema.clone_without_sys_signals().values.keys())
 
     # getting correct on and right_on column names
+    on_ = on
     on = left.signals_schema.resolve(*on).db_signals()  # type: ignore[assignment]
-    right_on = right.signals_schema.resolve(*(right_on or on)).db_signals()  # type: ignore[assignment]
+    right_on = right.signals_schema.resolve(*(right_on or on_)).db_signals()  # type: ignore[assignment]
 
     # getting correct compare and right_compare column names if they are defined
     if compare:
+        compare_ = compare
         compare = left.signals_schema.resolve(*compare).db_signals()  # type: ignore[assignment]
         right_compare = right.signals_schema.resolve(
-            *(right_compare or compare)
+            *(right_compare or compare_)
         ).db_signals()  # type: ignore[assignment]
     elif not compare and len(cols) != len(right_cols):
         # here we will mark all rows that are not added or deleted as modified since
@@ -155,7 +157,11 @@ def _to_list(obj: Optional[Union[str, Sequence[str]]]) -> Optional[list[str]]:
     if status_col:
         cols_select.append(diff_col)
 
-    dc_diff = dc_diff.select(*cols_select)
+    if not dc_diff._sys:
+        # TODO workaround when sys signal is not available in diff
+        dc_diff = dc_diff.settings(sys=True).select(*cols_select).settings(sys=False)
+    else:
+        dc_diff = dc_diff.select(*cols_select)
 
     # final schema is schema from the left chain with status column added if needed
     dc_diff.signals_schema = (

src/datachain/lib/dc/datachain.py

@@ -25,6 +25,7 @@
 
 from datachain import semver
 from datachain.dataset import DatasetRecord
+from datachain.delta import delta_disabled, delta_update
 from datachain.func import literal
 from datachain.func.base import Function
 from datachain.func.func import Func
@@ -72,6 +73,9 @@
     P = ParamSpec("P")
 
 
+T = TypeVar("T", bound="DataChain")
+
+
 class DataChain:
     """DataChain - a data structure for batch data processing and evaluation.
 
@@ -164,6 +168,7 @@ def __init__(
         self.signals_schema = signal_schema
         self._setup: dict = setup or {}
         self._sys = _sys
+        self._delta = False
 
     def __repr__(self) -> str:
         """Return a string representation of the chain."""
@@ -177,6 +182,32 @@ def __repr__(self) -> str:
         self.print_schema(file=file)
         return file.getvalue()
 
+    def _as_delta(
+        self,
+        on: Optional[Union[str, Sequence[str]]] = None,
+        right_on: Optional[Union[str, Sequence[str]]] = None,
+        compare: Optional[Union[str, Sequence[str]]] = None,
+    ) -> "Self":
+        """Marks this chain as delta, which means special delta process will be
+        called on saving dataset for optimization"""
+        if on is None:
+            raise ValueError("'delta on' fields must be defined")
+        self._delta = True
+        self._delta_on = on
+        self._delta_result_on = right_on
+        self._delta_compare = compare
+        return self
+
+    @property
+    def empty(self) -> bool:
+        """Returns True if chain has zero number of rows"""
+        return not bool(self.count())
+
+    @property
+    def delta(self) -> bool:
+        """Returns True if this chain is ran in "delta" update mode"""
+        return self._delta
+
     @property
     def schema(self) -> dict[str, DataType]:
         """Get schema of the chain."""
@@ -254,9 +285,17 @@ def _evolve(
             signal_schema = copy.deepcopy(self.signals_schema)
         if _sys is None:
             _sys = self._sys
-        return type(self)(
+        chain = type(self)(
             query, settings, signal_schema=signal_schema, setup=self._setup, _sys=_sys
         )
+        if self.delta:
+            chain = chain._as_delta(
+                on=self._delta_on,
+                right_on=self._delta_result_on,
+                compare=self._delta_compare,
+            )
+
+        return chain
 
     def settings(
         self,
@@ -463,7 +502,7 @@ def save(  # type: ignore[override]
         attrs: Optional[list[str]] = None,
         update_version: Optional[str] = "patch",
         **kwargs,
-    ) -> "Self":
+    ) -> "DataChain":
         """Save to a Dataset. It returns the chain itself.
 
         Parameters:
@@ -490,6 +529,35 @@ def save(  # type: ignore[override]
             )
 
         schema = self.signals_schema.clone_without_sys_signals().serialize()
+        if self.delta and name:
+            delta_ds, dependencies, has_changes = delta_update(
+                self,
+                name,
+                on=self._delta_on,
+                right_on=self._delta_result_on,
+                compare=self._delta_compare,
+            )
+
+            if delta_ds:
+                return self._evolve(
+                    query=delta_ds._query.save(
+                        name=name,
+                        version=version,
+                        feature_schema=schema,
+                        dependencies=dependencies,
+                        **kwargs,
+                    )
+                )
+
+            if not has_changes:
+                # sources have not been changed so new version of resulting dataset
+                # would be the same as previous one. To avoid duplicating exact
+                # datasets, we won't create new version of it and we will return
+                # current latest version instead.
+                from .datasets import read_dataset
+
+                return read_dataset(name, **kwargs)
+
         return self._evolve(
             query=self._query.save(
                 name=name,
@@ -615,6 +683,7 @@ def gen(
             signal_schema=udf_obj.output,
         )
 
+    @delta_disabled
     def agg(
         self,
         func: Optional[Callable] = None,
@@ -768,6 +837,7 @@ def order_by(self, *args, descending: bool = False) -> "Self":
 
         return self._evolve(query=self._query.order_by(*args))
 
+    @delta_disabled
     def distinct(self, arg: str, *args: str) -> "Self":  # type: ignore[override]
         """Removes duplicate rows based on uniqueness of some input column(s)
         i.e if rows are found with the same value of input column(s), only one
@@ -802,6 +872,7 @@ def select_except(self, *args: str) -> "Self":
             query=self._query.select(*columns), signal_schema=new_schema
         )
 
+    @delta_disabled  # type: ignore[arg-type]
     def group_by(
         self,
         *,
@@ -1160,6 +1231,7 @@ def remove_file_signals(self) -> "Self":
         schema = self.signals_schema.clone_without_file_signals()
         return self.select(*schema.values.keys())
 
+    @delta_disabled
     def merge(
         self,
         right_ds: "DataChain",
@@ -1268,6 +1340,7 @@ def _resolve(
 
         return ds
 
+    @delta_disabled
     def union(self, other: "Self") -> "Self":
         """Return the set union of the two datasets.