ygg 0.1.57__py3-none-any.whl → 0.1.64__py3-none-any.whl

yggdrasil/databricks/workspaces/io.py

@@ -2,27 +2,26 @@
 
 import base64
 import io
+import logging
+import os
 import time
 from abc import ABC, abstractmethod
+from tempfile import SpooledTemporaryFile
 from threading import Thread
-from typing import TYPE_CHECKING, Optional, IO, AnyStr, Union
+from typing import TYPE_CHECKING, Optional, IO, AnyStr, Union, Any, BinaryIO
 
+import dill
 import pyarrow as pa
 import pyarrow.csv as pcsv
 import pyarrow.parquet as pq
-from pyarrow.dataset import (
-    FileFormat,
-    ParquetFileFormat,
-    CsvFileFormat,
-)
 
 from .path_kind import DatabricksPathKind
 from ...libs.databrickslib import databricks
 from ...libs.pandaslib import PandasDataFrame
 from ...libs.polarslib import polars, PolarsDataFrame
-from ...pyutils import retry
+from ...pyutils.retry import retry
 from ...types.cast.registry import convert
-from ...types.file_format import ExcelFileFormat
+from ...types.file_format import FileFormat, ParquetFileFormat, CsvFileFormat, ExcelFileFormat
 
 if databricks is not None:
     from databricks.sdk.service.workspace import ImportFormat, ExportFormat
@@ -42,6 +41,66 @@ __all__ = [
 ]
 
 
+LOGGER = logging.getLogger(__name__)
+_SPOOL_MAX = 64 * 1024 * 1024   # 64MB in RAM then spill to disk
+_COPY_CHUNK = 8 * 1024 * 1024   # 8MB chunks
+
+def _prepare_binaryio_and_size(
+    data: Union[bytes, bytearray, memoryview, BinaryIO]
+) -> tuple[int, BinaryIO, bool]:
+    """
+    Returns (size, bio, should_close).
+
+    - bytes-like -> wrap in BytesIO (closeable by us).
+    - seekable file -> compute size via fstat or seek/tell.
+    - non-seekable stream -> spool into SpooledTemporaryFile, count bytes.
+    """
+    # bytes-like
+    if isinstance(data, (bytes, bytearray, memoryview)):
+        b = bytes(data)
+        return len(b), io.BytesIO(b), True
+
+    f: BinaryIO = data
+
+    # 1) try OS-level size for real files
+    try:
+        fileno = f.fileno()  # type: ignore[attr-defined]
+    except Exception:
+        fileno = None
+
+    if fileno is not None:
+        try:
+            st = os.fstat(fileno)
+            # rewind if possible
+            try:
+                f.seek(0)
+            except Exception:
+                pass
+            return int(st.st_size), f, False
+        except Exception:
+            pass
+
+    # 2) try seek/tell (seekable streams)
+    try:
+        f.seek(0, io.SEEK_END)
+        end = f.tell()
+        f.seek(0)
+        return int(end), f, False
+    except Exception:
+        pass
+
+    # 3) non-seekable stream: spool + count
+    spooled = SpooledTemporaryFile(max_size=_SPOOL_MAX, mode="w+b")
+    size = 0
+    while True:
+        chunk = f.read(_COPY_CHUNK)
+        if not chunk:
+            break
+        spooled.write(chunk)
+        size += len(chunk)
+    spooled.seek(0)
+    return size, spooled, True
+
 class DatabricksIO(ABC, IO):
     """File-like interface for Databricks workspace, volume, or DBFS paths."""
 
@@ -96,6 +155,18 @@ class DatabricksIO(ABC, IO):
     def __hash__(self):
         return self.path.__hash__()
 
+    def __str__(self):
+        return "%s(path=%s)" % (
+            self.__class__.__name__,
+            self.path.__repr__()
+        )
+
+    def __repr__(self):
+        return "%s(path=%s)" % (
+            self.__class__.__name__,
+            self.path.__repr__()
+        )
+
     @classmethod
     def create_instance(
         cls,
@@ -504,7 +575,7 @@ class DatabricksIO(ABC, IO):
         return True
 
     @abstractmethod
-    def write_all_bytes(self, data: bytes):
+    def write_all_bytes(self, data: Union[bytes, IO[bytes]]):
         """Write raw bytes to the remote path.
 
         Args:
@@ -921,10 +992,32 @@ class DatabricksIO(ABC, IO):
 
         self.write_all_bytes(data=buffer.getvalue())
 
+    def read_object(
+        self,
+    ):
+        content = self.read_all_bytes()
+        obj = dill.loads(content)
+
+        return obj
+
+    def write_object(
+        self,
+        obj: Any,
+        file_format: Optional[FileFormat] = None,
+    ):
+        buffer = io.BytesIO()
+
+        if isinstance(obj, PandasDataFrame):
+            obj.to_pickle(buffer)
+        else:
+            buffer.write(dill.dumps(obj))
+
+        self.write_all_bytes(data=buffer.getvalue())
 
 class DatabricksWorkspaceIO(DatabricksIO):
     """IO adapter for Workspace files."""
 
+    @retry(exceptions=(InternalError,))
     def read_byte_range(self, start: int, length: int, allow_not_found: bool = False) -> bytes:
         """Read bytes from a Workspace file.
 
@@ -956,11 +1049,12 @@ class DatabricksWorkspaceIO(DatabricksIO):
         end = start + length
         return data[start:end]
 
-    def write_all_bytes(self, data: bytes):
+    @retry(exceptions=(InternalError,))
+    def write_all_bytes(self, data: Union[bytes, IO[bytes]]):
         """Write bytes to a Workspace file.
 
         Args:
-            data: Bytes to write.
+            data: Union[bytes, IO[bytes]] to write.
 
         Returns:
             The DatabricksWorkspaceIO instance.
@@ -969,6 +1063,20 @@ class DatabricksWorkspaceIO(DatabricksIO):
         workspace_client = sdk.workspace
         full_path = self.path.workspace_full_path()
 
+        if isinstance(data, bytes):
+            bsize = len(data)
+        elif isinstance(data, io.BytesIO):
+            bsize = len(data.getvalue())
+        else:
+            bsize = None
+
+        LOGGER.debug(
+            "Writing %s(size=%s) in %s",
+            type(data),
+            bsize,
+            self
+        )
+
         try:
             workspace_client.upload(
                 full_path,
@@ -989,16 +1097,23 @@ class DatabricksWorkspaceIO(DatabricksIO):
         self.path.reset_metadata(
             is_file=True,
             is_dir=False,
-            size=len(data),
+            size=bsize,
             mtime=time.time()
         )
 
+        LOGGER.info(
+            "Written %s bytes in %s",
+            bsize,
+            self
+        )
+
         return self
 
 
 class DatabricksVolumeIO(DatabricksIO):
     """IO adapter for Unity Catalog volume files."""
 
+    @retry(exceptions=(InternalError,))
     def read_byte_range(self, start: int, length: int, allow_not_found: bool = False) -> bytes:
         """Read bytes from a volume file.
 
@@ -1023,9 +1138,9 @@ class DatabricksVolumeIO(DatabricksIO):
 
         try:
             resp = client.download(full_path)
-        except Exception as e:
+        except (NotFound, ResourceDoesNotExist, BadRequest, InternalError) as e:
             # Databricks SDK exceptions vary a bit by version; keep it pragmatic.
-            if allow_not_found and any(s in str(e).lower() for s in ("not found", "not exist", "404")):
+            if allow_not_found:
                 return b""
             raise
 
@@ -1038,48 +1153,67 @@ class DatabricksVolumeIO(DatabricksIO):
         end = start + length
         return data[start:end]
 
-    @retry(exceptions=(InternalError,))
-    def write_all_bytes(self, data: bytes):
-        """Write bytes to a volume file.
-
-        Args:
-            data: Bytes to write.
-
-        Returns:
-            The DatabricksVolumeIO instance.
-        """
+    def write_all_bytes(
+        self,
+        data: Union[bytes, bytearray, memoryview, BinaryIO],
+        *,
+        overwrite: bool = True,
+        part_size: Optional[int] = None,
+        use_parallel: bool = True,
+        parallelism: Optional[int] = None,
+    ):
+        """Write bytes/stream to a volume file safely (BinaryIO upload)."""
         sdk = self.workspace.sdk()
         client = sdk.files
         full_path = self.path.files_full_path()
 
-        try:
-            client.upload(
-                full_path,
-                io.BytesIO(data),
-                overwrite=True
-            )
-        except (NotFound, ResourceDoesNotExist, BadRequest):
-            self.path.parent.mkdir(parents=True, exist_ok=True)
+        LOGGER.debug("Writing all bytes in %s", self)
 
-            client.upload(
+        size, bio, should_close = _prepare_binaryio_and_size(data)
+
+        def _upload():
+            return client.upload(
                 full_path,
-                io.BytesIO(data),
-                overwrite=True
+                bio,
+                overwrite=overwrite,
+                part_size=part_size,
+                use_parallel=use_parallel,
+                parallelism=parallelism,
             )
 
+        try:
+            _ = _upload()
+        except (NotFound, ResourceDoesNotExist, BadRequest, InternalError):
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+            # Important: rewind if possible before retry
+            try:
+                bio.seek(0)
+            except Exception:
+                pass
+            _ = _upload()
+        finally:
+            if should_close:
+                try:
+                    bio.close()
+                except Exception:
+                    pass
+
         self.path.reset_metadata(
             is_file=True,
             is_dir=False,
-            size=len(data),
-            mtime=time.time()
+            size=size,
+            mtime=time.time(),
         )
 
-        return self
+        LOGGER.info("Written %s bytes in %s", size or "all", self.path)
+
+        return self  # or return result if your API prefers that
 
 
 class DatabricksDBFSIO(DatabricksIO):
     """IO adapter for DBFS files."""
 
+    @retry(exceptions=(InternalError,))
     def read_byte_range(self, start: int, length: int, allow_not_found: bool = False) -> bytes:
         """Read bytes from a DBFS file.
 
@@ -1124,11 +1258,12 @@ class DatabricksDBFSIO(DatabricksIO):
 
         return bytes(read_bytes)
 
-    def write_all_bytes(self, data: bytes):
+    @retry(exceptions=(InternalError,))
+    def write_all_bytes(self, data: Union[bytes, IO[bytes]]):
         """Write bytes to a DBFS file.
 
         Args:
-            data: Bytes to write.
+            data: Union[bytes, IO[bytes]] to write.
 
         Returns:
             The DatabricksDBFSIO instance.
@@ -1137,6 +1272,11 @@ class DatabricksDBFSIO(DatabricksIO):
         client = sdk.dbfs
         full_path = self.path.dbfs_full_path()
 
+        LOGGER.debug(
+            "Writing all bytes in %s",
+            self
+        )
+
         try:
             with client.open(
                 path=full_path,
@@ -1156,6 +1296,11 @@ class DatabricksDBFSIO(DatabricksIO):
             ) as f:
                 f.write(data)
 
+        LOGGER.info(
+            "Written all bytes in %s",
+            self
+        )
+
         self.path.reset_metadata(
             is_file=True,
             is_dir=False,