datachain 0.36.4__tar.gz → 0.36.6__tar.gz

{datachain-0.36.4 → datachain-0.36.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.36.4
+Version: 0.36.6
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.36.4 → datachain-0.36.6}/src/datachain/delta.py

@@ -1,12 +1,16 @@
+import hashlib
 from collections.abc import Sequence
 from copy import copy
 from functools import wraps
 from typing import TYPE_CHECKING, TypeVar
 
+from attrs import frozen
+
 import datachain
 from datachain.dataset import DatasetDependency, DatasetRecord
 from datachain.error import DatasetNotFoundError
 from datachain.project import Project
+from datachain.query.dataset import Step, step_result
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -14,7 +18,9 @@ if TYPE_CHECKING:
 
     from typing_extensions import ParamSpec
 
+    from datachain.catalog import Catalog
     from datachain.lib.dc import DataChain
+    from datachain.query.dataset import QueryGenerator
 
     P = ParamSpec("P")
 
@@ -43,11 +49,38 @@ def delta_disabled(
     return _inner
 
 
+@frozen
+class _RegenerateSystemColumnsStep(Step):
+    catalog: "Catalog"
+
+    def hash_inputs(self) -> str:
+        return hashlib.sha256(b"regenerate_sys_columns").hexdigest()
+
+    def apply(self, query_generator: "QueryGenerator", temp_tables: list[str]):
+        selectable = query_generator.select()
+        regenerated = self.catalog.warehouse._regenerate_system_columns(
+            selectable,
+            keep_existing_columns=True,
+            regenerate_columns=None,
+        )
+
+        def q(*columns):
+            return regenerated.with_only_columns(*columns)
+
+        return step_result(q, regenerated.selected_columns)
+
+
 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.append(
+        _RegenerateSystemColumnsStep(
+            catalog=dc.session.catalog,
+        )
+    )
+
     dc._query.steps += other._query.steps.copy()
     dc.signals_schema = other.signals_schema
     return dc

{datachain-0.36.4 → datachain-0.36.6}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.36.4
+Version: 0.36.6
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.36.4 → datachain-0.36.6}/tests/func/test_delta.py

@@ -1,4 +1,5 @@
 import os
+import uuid
 
 import pytest
 import regex as re
@@ -224,6 +225,42 @@ def test_delta_update_unsafe(test_session):
     }
 
 
+def test_delta_replay_regenerates_system_columns(test_session):
+    source_name = f"regen_source_{uuid.uuid4().hex[:8]}"
+    result_name = f"regen_result_{uuid.uuid4().hex[:8]}"
+
+    dc.read_values(
+        measurement_id=[1, 2],
+        err=["", ""],
+        num=[1, 2],
+        session=test_session,
+    ).save(source_name)
+
+    def build_chain(delta: bool):
+        read_kwargs = {"session": test_session}
+        if delta:
+            read_kwargs.update({"delta": True, "delta_on": "measurement_id"})
+        return (
+            dc.read_dataset(source_name, **read_kwargs)
+            .filter(C.err == "")
+            .select_except("err")
+            .map(double=lambda num: num * 2, output=int)
+            .select_except("num")
+        )
+
+    build_chain(delta=False).save(result_name)
+
+    build_chain(delta=True).save(
+        result_name,
+        delta=True,
+        delta_on="measurement_id",
+    )
+
+    assert set(
+        dc.read_dataset(result_name, session=test_session).to_values("measurement_id")
+    ) == {1, 2}
+
+
 def test_delta_update_from_storage(test_session, tmp_dir, tmp_path):
     ds_name = "delta_ds"
     path = tmp_dir.as_uri()

{datachain-0.36.4 → datachain-0.36.6}/tests/func/test_temp_table_tracking.py

@@ -27,11 +27,13 @@ def test_nested_merge_has_no_duplicate_temp_tables(test_session, mocker):
     inner = generated.merge(base, on="num", inner=True)
     chain = base.merge(inner, on="num", inner=True)
 
-    expected = chain.select("num").to_pandas()["num"].tolist()
-    assert expected == [1, 2]
+    values = chain.select("num").to_pandas()["num"].tolist()
+    assert set(values) == {1, 2}
+    assert len(values) == 2
 
     rerun = chain.select("num").to_pandas()["num"].tolist()
-    assert rerun == expected
+    assert set(rerun) == {1, 2}
+    assert len(rerun) == len(values)
 
     _assert_no_duplicate_temp_tables(captured)
 
@@ -43,11 +45,13 @@ def test_union_has_no_duplicate_temp_tables(test_session, mocker):
     right = dc.read_values(num=[3], session=test_session)
     union_chain = left.union(right)
 
-    expected = sorted(union_chain.select("num").to_pandas()["num"].tolist())
-    assert expected == [1, 2, 3]
+    values = union_chain.select("num").to_pandas()["num"].tolist()
+    assert set(values) == {1, 2, 3}
+    assert len(values) == 3
 
-    rerun = sorted(union_chain.select("num").to_pandas()["num"].tolist())
-    assert rerun == expected
+    rerun = union_chain.select("num").to_pandas()["num"].tolist()
+    assert set(rerun) == {1, 2, 3}
+    assert len(rerun) == len(values)
 
     _assert_no_duplicate_temp_tables(captured)
 
@@ -59,10 +63,10 @@ def test_subtract_has_no_duplicate_temp_tables(test_session, mocker):
     target = dc.read_values(num=[2], session=test_session)
     subtract_chain = source.subtract(target, on="num")
 
-    expected = subtract_chain.select("num").to_pandas()["num"].tolist()
-    assert expected == [1]
+    values = subtract_chain.select("num").to_pandas()["num"].tolist()
+    assert values == [1]
 
     rerun = subtract_chain.select("num").to_pandas()["num"].tolist()
-    assert rerun == expected
+    assert rerun == values
 
     _assert_no_duplicate_temp_tables(captured)