Fix double mapper call when using to_storage method (#1258)

Author: dreadatour

src/datachain/lib/dc/datachain.py

@@ -2419,9 +2419,11 @@ def to_storage(
             ds.to_storage("gs://mybucket", placement="filename")
             ```
         """
+        chain = self.persist()
+        count = chain.count()
+
         if placement == "filename" and (
-            self._query.distinct(pathfunc.name(C(f"{signal}__path"))).count()
-            != self._query.count()
+            chain._query.distinct(pathfunc.name(C(f"{signal}__path"))).count() != count
         ):
             raise ValueError("Files with the same name found")
 
@@ -2433,7 +2435,7 @@ def to_storage(
             unit=" files",
             unit_scale=True,
             unit_divisor=10,
-            total=self.count(),
+            total=count,
             leave=False,
         )
         file_exporter = FileExporter(
@@ -2444,7 +2446,10 @@ def to_storage(
             max_threads=num_threads or 1,
             client_config=client_config,
         )
-        file_exporter.run(self.to_values(signal), progress_bar)
+        file_exporter.run(
+            (rows[0] for rows in chain.to_iter(signal)),
+            progress_bar,
+        )
 
     def shuffle(self) -> "Self":
         """Shuffle the rows of the chain deterministically."""

tests/func/test_datachain.py

@@ -9,7 +9,7 @@
 from collections.abc import Iterator
 from datetime import datetime, timedelta, timezone
 from pathlib import Path, PurePosixPath
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import numpy as np
 import pandas as pd
@@ -358,15 +358,24 @@ def test_to_storage(
     file_type,
     num_threads,
 ):
+    mapper = Mock(side_effect=lambda file_path: len(file_path))
+
     ctc = cloud_test_catalog
     df = dc.read_storage(ctc.src_uri, type=file_type, session=test_session)
     if use_map:
-        df.settings(cache=use_cache).map(
-            res=lambda file: file.export(tmp_dir / "output", placement=placement)
-        ).exec()
+        (
+            df.settings(cache=use_cache)
+            .map(mapper, params=["file.path"], output={"path_len": int})
+            .map(res=lambda file: file.export(tmp_dir / "output", placement=placement))
+            .exec()
+        )
     else:
-        df.settings(cache=use_cache).to_storage(
-            tmp_dir / "output", placement=placement, num_threads=num_threads
+        (
+            df.settings(cache=use_cache)
+            .map(mapper, params=["file.path"], output={"path_len": int})
+            .to_storage(
+                tmp_dir / "output", placement=placement, num_threads=num_threads
+            )
         )
 
     expected = {
@@ -387,6 +396,8 @@ def test_to_storage(
         with open(tmp_dir / "output" / file_path) as f:
             assert f.read() == expected[file.name]
 
+    assert mapper.call_count == len(expected)
+
 
 @pytest.mark.parametrize("use_cache", [True, False])
 def test_export_images_files(test_session, tmp_dir, tmp_path, use_cache):