dont duplicate temp tables on reusing queries (#1410)

Author: shcheklein

src/datachain/query/dataset.py

@@ -231,8 +231,9 @@ def query(
 
     def apply(self, query_generator, temp_tables: list[str]) -> "StepResult":
         source_query = query_generator.exclude(("sys__id",))
+        right_before = len(self.dq.temp_table_names)
         target_query = self.dq.apply_steps().select()
-        temp_tables.extend(self.dq.temp_table_names)
+        temp_tables.extend(self.dq.temp_table_names[right_before:])
 
         # creating temp table that will hold subtract results
         temp_table_name = self.catalog.warehouse.temp_table_name()
@@ -951,10 +952,12 @@ def hash_inputs(self) -> str:
     def apply(
         self, query_generator: QueryGenerator, temp_tables: list[str]
     ) -> StepResult:
+        left_before = len(self.query1.temp_table_names)
         q1 = self.query1.apply_steps().select().subquery()
-        temp_tables.extend(self.query1.temp_table_names)
+        temp_tables.extend(self.query1.temp_table_names[left_before:])
+        right_before = len(self.query2.temp_table_names)
         q2 = self.query2.apply_steps().select().subquery()
-        temp_tables.extend(self.query2.temp_table_names)
+        temp_tables.extend(self.query2.temp_table_names[right_before:])
 
         columns1 = _drop_system_columns(q1.columns)
         columns2 = _drop_system_columns(q2.columns)
@@ -1004,8 +1007,9 @@ def hash_inputs(self) -> str:
         return hashlib.sha256(b"".join(parts)).hexdigest()
 
     def get_query(self, dq: "DatasetQuery", temp_tables: list[str]) -> sa.Subquery:
+        temp_tables_before = len(dq.temp_table_names)
         query = dq.apply_steps().select()
-        temp_tables.extend(dq.temp_table_names)
+        temp_tables.extend(dq.temp_table_names[temp_tables_before:])
 
         if not any(isinstance(step, (SQLJoin, SQLUnion)) for step in dq.steps):
             return query.subquery(dq.table.name)

tests/func/test_temp_table_tracking.py

@@ -0,0 +1,68 @@
+import datachain as dc
+from datachain.query.dataset import DatasetQuery
+
+
+def _capture_temp_tables(mocker):
+    captured: list[list[str]] = []
+    original_cleanup = DatasetQuery.cleanup
+
+    def capture(self):
+        captured.append(list(self.temp_table_names))
+        return original_cleanup(self)
+
+    mocker.patch("datachain.query.dataset.DatasetQuery.cleanup", capture)
+    return captured
+
+
+def _assert_no_duplicate_temp_tables(captured: list[list[str]]):
+    for tables in captured:
+        assert len(tables) == len(set(tables))
+
+
+def test_nested_merge_has_no_duplicate_temp_tables(test_session, mocker):
+    captured = _capture_temp_tables(mocker)
+
+    base = dc.read_values(num=[1, 2], session=test_session)
+    generated = base.map(num_plus=lambda num: str(num + 10))
+    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]
+
+    rerun = chain.select("num").to_pandas()["num"].tolist()
+    assert rerun == expected
+
+    _assert_no_duplicate_temp_tables(captured)
+
+
+def test_union_has_no_duplicate_temp_tables(test_session, mocker):
+    captured = _capture_temp_tables(mocker)
+
+    left = dc.read_values(num=[1, 2], session=test_session)
+    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]
+
+    rerun = sorted(union_chain.select("num").to_pandas()["num"].tolist())
+    assert rerun == expected
+
+    _assert_no_duplicate_temp_tables(captured)
+
+
+def test_subtract_has_no_duplicate_temp_tables(test_session, mocker):
+    captured = _capture_temp_tables(mocker)
+
+    source = dc.read_values(num=[1, 2], session=test_session)
+    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]
+
+    rerun = subtract_chain.select("num").to_pandas()["num"].tolist()
+    assert rerun == expected
+
+    _assert_no_duplicate_temp_tables(captured)