perf(file): stream File.save() to avoid OOM on large files

src/datachain/client/fsspec.py

@@ -5,10 +5,10 @@
 import os
 import posixpath
 import re
+import shutil
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Iterator, Sequence
 from datetime import datetime
-from shutil import copy2
 from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, Literal, NamedTuple
 from urllib.parse import urlparse
 
@@ -424,7 +424,7 @@ def do_instantiate_object(self, file: "File", dst: str) -> None:
             reflink(src, dst)
         except OSError:
             # Default to copy if reflinks are not supported
-            copy2(src, dst)
+            shutil.copy2(src, dst)
 
     def open_object(
         self, file: "File", use_cache: bool = True, cb: Callback = DEFAULT_CALLBACK
@@ -440,7 +440,21 @@ def open_object(
             cb,
         )  # type: ignore[return-value]
 
-    def upload(self, data: bytes, path: str) -> "File":
+    def upload(
+        self, data: bytes | bytearray | memoryview | BinaryIO, path: str
+    ) -> "File":
+        """Upload *data* to *path*.
+
+        ``data`` may be:
+          - a bytes-like object (``bytes``/``bytearray``/``memoryview``) —
+            written in a single ``pipe_file`` call.
+          - a binary readable stream — copied into the destination via
+            ``fs.open(path, 'wb')`` using ``shutil.copyfileobj``. The
+            destination handle (an ``AbstractBufferedFile`` on cloud
+            backends) flushes a multipart part every ``self.blocksize``
+            bytes (e.g. 5 MiB on s3fs), so peak RAM is bounded by the
+            backend's part size rather than the file size.
+        """
         if path.startswith(self.PREFIX):
             full_path = path
             _, rel_path = self.split_url(path)
@@ -453,7 +467,17 @@ def upload(self, data: bytes, path: str) -> "File":
         parent = posixpath.dirname(full_path)
         self.fs.makedirs(parent, exist_ok=True)
 
-        self.fs.pipe_file(full_path, data)
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            self.fs.pipe_file(full_path, data)
+        else:
+            with self.fs.open(full_path, "wb") as dst:
+                # Use a larger copy buffer than shutil's 16 KiB default to
+                # reduce per-chunk overhead on multi-GB uploads. The
+                # destination's blocksize (when exposed by fsspec backends)
+                # matches the multipart part size it will flush at, making it
+                # a natural choice; fall back to 8 MiB otherwise.
+                buf_size = getattr(dst, "blocksize", None) or 8 * 1024 * 1024
+                shutil.copyfileobj(data, dst, length=buf_size)
         file_info = self.fs.info(full_path, **self._file_info_kwargs())
         return self.info_to_file(file_info, rel_path)
 

src/datachain/client/http.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, ClassVar, cast
+from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, cast
 from urllib.parse import quote, urlparse
 
 from fsspec.implementations.http import HTTPFileSystem
@@ -145,7 +145,9 @@ def info_to_file(self, v: dict[str, Any], path: str) -> File:
             last_modified=last_modified,
         )
 
-    def upload(self, data: bytes, path: str) -> "File":
+    def upload(
+        self, data: bytes | bytearray | memoryview | BinaryIO, path: str
+    ) -> "File":
         raise NotImplementedError(
             "HTTP/HTTPS client is read-only. Upload operations are not supported."
         )

src/datachain/lib/file.py

@@ -631,7 +631,8 @@ def save(self, destination: str, client_config: dict | None = None) -> "File":
 
         destination = stringify_path(destination)
         client, rel_path = self._resolve_destination(destination, client_config)
-        result = client.upload(self.read_bytes(), rel_path)
+        with self.open(mode="rb") as src:
+            result = client.upload(src, rel_path)
         result._set_stream(self._catalog)
         return result
 

tests/unit/test_client_local.py

@@ -427,6 +427,19 @@ def test_upload_returned_file_path_matches_relative_input(tmp_path, catalog):
     assert result.source == path_to_fsspec_uri(str(tmp_path))
 
 
+def test_upload_accepts_binary_stream(tmp_path, catalog):
+    client = FileClient.from_source(str(tmp_path), catalog.cache)
+
+    src = tmp_path / "src.bin"
+    src.write_bytes(b"streamed-content")
+
+    with open(src, "rb") as fh:
+        result = client.upload(fh, "out.bin")
+
+    assert (tmp_path / "out.bin").read_bytes() == b"streamed-content"
+    assert result.path == "out.bin"
+
+
 @pytest.mark.parametrize(
     "path,match",
     [