fixing dataset dependencies when persist is used (#1112)

Author: ilongin

src/datachain/query/dataset.py

@@ -1675,13 +1675,27 @@ def generate(
         return query
 
     def _add_dependencies(self, dataset: "DatasetRecord", version: str):
-        for dependency in self.dependencies:
-            ds_dependency_name, ds_dependency_version = dependency
+        dependencies: set[DatasetDependencyType] = set()
+        for dep_name, dep_version in self.dependencies:
+            if Session.is_temp_dataset(dep_name):
+                # temp dataset are created for optimization and they will be removed
+                # afterwards. Therefore, we should not put them as dependencies, but
+                # their own direct dependencies
+                for dep in self.catalog.get_dataset_dependencies(
+                    dep_name, dep_version, indirect=False
+                ):
+                    if dep:
+                        dependencies.add((dep.name, dep.version))
+            else:
+                dependencies.add((dep_name, dep_version))
+
+        for dep_name, dep_version in dependencies:
+            # ds_dependency_name, ds_dependency_version = dependency
             self.catalog.metastore.add_dataset_dependency(
                 dataset.name,
                 version,
-                ds_dependency_name,
-                ds_dependency_version,
+                dep_name,
+                dep_version,
             )
 
     def exec(self) -> "Self":

tests/func/test_datachain.py

@@ -239,6 +239,25 @@ def test_read_storage_dependencies(cloud_test_catalog, cloud_type):
     assert dependencies[0].name == dep_name
 
 
+def test_persist_not_affects_dependencies(tmp_dir, test_session):
+    for i in range(4):
+        (tmp_dir / f"file{i}.txt").write_text(f"file{i}")
+
+    uri = tmp_dir.as_uri()
+    dep_name, _, _ = parse_listing_uri(uri, test_session.catalog.client_config)
+    chain = dc.read_storage(uri, session=test_session)  # .persist()
+    # calling multiple persists to create temp datasets
+    chain = chain.persist()
+    chain = chain.persist()
+    chain = chain.persist()
+    chain.save("test-data")
+    dependencies = test_session.catalog.get_dataset_dependencies("test-data", "1.0.0")
+
+    assert len(dependencies) == 1
+    assert dependencies[0].name == dep_name
+    assert dependencies[0].type == DatasetDependencyType.STORAGE
+
+
 @pytest.mark.parametrize("use_cache", [True, False])
 @pytest.mark.parametrize("prefetch", [0, 2])
 def test_map_file(cloud_test_catalog, use_cache, prefetch, monkeypatch):