ygg 0.1.56__py3-none-any.whl → 0.1.60__py3-none-any.whl

yggdrasil/databricks/workspaces/path.py

@@ -5,12 +5,15 @@ from __future__ import annotations
 
 import dataclasses
 import datetime as dt
+import io
 import random
 import string
 import time
 from pathlib import PurePosixPath
-from typing import Optional, Tuple, Union, TYPE_CHECKING, List
+from threading import Thread
+from typing import Optional, Tuple, Union, TYPE_CHECKING, List, Any, IO
 
+import dill
 import pyarrow as pa
 import pyarrow.dataset as ds
 from pyarrow import ArrowInvalid
@@ -23,14 +26,15 @@ from .volumes_path import get_volume_status, get_volume_metadata
 from ...libs.databrickslib import databricks
 from ...libs.pandaslib import PandasDataFrame
 from ...libs.polarslib import polars, PolarsDataFrame
-from ...types.cast.arrow_cast import cast_arrow_tabular
 from ...types.cast.cast_options import CastOptions
+from ...types.cast.pandas_cast import pandas_converter, cast_pandas_dataframe
 from ...types.cast.polars_cast import polars_converter, cast_polars_dataframe
 from ...types.cast.registry import convert, register_converter
 from ...types.file_format import ExcelFileFormat
 
 if databricks is not None:
-    from databricks.sdk.service.catalog import VolumeType, VolumeInfo
+    from databricks.sdk.errors import InternalError
+    from databricks.sdk.service.catalog import VolumeType, VolumeInfo, PathOperation
     from databricks.sdk.service.workspace import ObjectType
     from databricks.sdk.errors.platform import (
         NotFound,
@@ -176,6 +180,8 @@ class DatabricksPath:
         if not obj:
             return cls.empty_instance(workspace=workspace)
 
+        if isinstance(obj, str):
+            obj = [obj]
         if not isinstance(obj, (str, list)):
             if isinstance(obj, DatabricksPath):
                 if workspace is not None and obj._workspace is None:
@@ -191,6 +197,7 @@ class DatabricksPath:
                 obj = str(obj)
 
 
+
         obj = _flatten_parts(obj)
 
         if obj and not obj[0]:
@@ -246,16 +253,23 @@ class DatabricksPath:
         if self._workspace is not None:
             self._workspace.__exit__(exc_type, exc_val, exc_tb)
 
+        self.close(wait=False)
+
     def __str__(self):
         return self.full_path()
 
     def __repr__(self):
         return self.url()
 
+    def __del__(self):
+        self.close(wait=False)
+
     def __fspath__(self):
         return self.full_path()
 
     def url(self):
+        if self._workspace is not None:
+            return self._workspace.safe_host + self.full_path()
         return "dbfs://%s" % self.full_path()
 
     def full_path(self) -> str:
@@ -282,7 +296,7 @@ class DatabricksPath:
         Returns:
             A PyArrow FileSystem instance.
         """
-        return self.workspace.filesytem(workspace=workspace)
+        return self.workspace.filesystem(workspace=workspace)
 
     @property
     def parent(self):
@@ -496,9 +510,15 @@ class DatabricksPath:
 
         return self
 
-    def close(self):
+    def close(self, wait: bool = True):
         if self.temporary:
-            self.remove(recursive=True)
+            if wait:
+                self.remove(recursive=True)
+            else:
+                Thread(
+                    target=self.remove,
+                    kwargs={"recursive": True}
+                ).start()
 
     def storage_location(self) -> str:
         info = self.volume_info()
@@ -586,7 +606,7 @@ class DatabricksPath:
             mtime = float(info.modified_at) / 1000.0 if info.modified_at is not None else None
 
             return self.reset_metadata(is_file=is_file, is_dir=is_dir, size=size, mtime=mtime)
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             pass
 
         found = next(self.ls(fetch_size=1, recursive=False, allow_not_found=True), None)
@@ -730,7 +750,7 @@ class DatabricksPath:
                     properties=default_tags,
                     comment="Catalog auto generated by yggdrasil"
                 )
-            except (AlreadyExists, ResourceAlreadyExists, PermissionDenied, BadRequest):
+            except (AlreadyExists, ResourceAlreadyExists, PermissionDenied, BadRequest, InternalError):
                 if not exist_ok:
                     raise
 
@@ -742,7 +762,7 @@ class DatabricksPath:
                     properties=default_tags,
                     comment="Schema auto generated by yggdrasil"
                 )
-            except (AlreadyExists, ResourceAlreadyExists, PermissionDenied, BadRequest):
+            except (AlreadyExists, ResourceAlreadyExists, PermissionDenied, BadRequest, InternalError):
                 if not exist_ok:
                     raise
 
@@ -808,37 +828,54 @@ class DatabricksPath:
 
     def remove(
         self,
-        recursive: bool = True
+        recursive: bool = True,
+        allow_not_found: bool = True
     ):
         """Remove the path as a file or directory.
 
         Args:
             recursive: Whether to delete directories recursively.
+            allow_not_found: Allow not found path
 
         Returns:
             The DatabricksPath instance.
         """
         if self.kind == DatabricksPathKind.VOLUME:
-            return self._remove_volume_obj(recursive=recursive)
+            return self._remove_volume_obj(recursive=recursive, allow_not_found=allow_not_found)
         elif self.kind == DatabricksPathKind.WORKSPACE:
-            return self._remove_workspace_obj(recursive=recursive)
+            return self._remove_workspace_obj(recursive=recursive, allow_not_found=allow_not_found)
         elif self.kind == DatabricksPathKind.DBFS:
-            return self._remove_dbfs_obj(recursive=recursive)
+            return self._remove_dbfs_obj(recursive=recursive, allow_not_found=allow_not_found)
 
-    def _remove_volume_obj(self, recursive: bool = True):
+    def _remove_volume_obj(
+        self,
+        recursive: bool = True,
+        allow_not_found: bool = True
+    ):
         if self.is_file():
-            return self._remove_volume_file()
-        return self._remove_volume_dir(recursive=recursive)
+            return self._remove_volume_file(allow_not_found=allow_not_found)
+        elif self.is_dir():
+            return self._remove_volume_dir(recursive=recursive, allow_not_found=allow_not_found)
 
-    def _remove_workspace_obj(self, recursive: bool = True):
+    def _remove_workspace_obj(
+        self,
+        recursive: bool = True,
+        allow_not_found: bool = True
+    ):
         if self.is_file():
-            return self._remove_workspace_file()
-        return self._remove_workspace_dir(recursive=recursive)
+            return self._remove_workspace_file(allow_not_found=allow_not_found)
+        elif self.is_dir():
+            return self._remove_workspace_dir(recursive=recursive, allow_not_found=allow_not_found)
 
-    def _remove_dbfs_obj(self, recursive: bool = True):
+    def _remove_dbfs_obj(
+        self,
+        recursive: bool = True,
+        allow_not_found: bool = True
+    ):
         if self.is_file():
-            return self._remove_dbfs_file()
-        return self._remove_dbfs_dir(recursive=recursive)
+            return self._remove_dbfs_file(allow_not_found=allow_not_found)
+        elif self.is_dir():
+            return self._remove_dbfs_dir(recursive=recursive, allow_not_found=allow_not_found)
 
     def rmfile(self, allow_not_found: bool = True):
         """Remove the path as a file.
@@ -859,7 +896,7 @@ class DatabricksPath:
         sdk = self.workspace.sdk()
         try:
             sdk.files.delete(self.files_full_path())
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             if not allow_not_found:
                 raise
         finally:
@@ -871,7 +908,7 @@ class DatabricksPath:
         sdk = self.workspace.sdk()
         try:
             sdk.workspace.delete(self.workspace_full_path(), recursive=True)
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             if not allow_not_found:
                 raise
         finally:
@@ -883,7 +920,7 @@ class DatabricksPath:
         sdk = self.workspace.sdk()
         try:
             sdk.dbfs.delete(self.dbfs_full_path(), recursive=True)
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             if not allow_not_found:
                 raise
         finally:
@@ -940,7 +977,7 @@ class DatabricksPath:
 
             if not with_root:
                 sdk.workspace.mkdirs(full_path)
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             if not allow_not_found:
                 raise
         finally:
@@ -962,7 +999,7 @@ class DatabricksPath:
 
             if not with_root:
                 sdk.dbfs.mkdirs(full_path)
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             if not allow_not_found:
                 raise
         finally:
@@ -983,7 +1020,7 @@ class DatabricksPath:
         if rel:
             try:
                 sdk.files.delete_directory(full_path)
-            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied) as e:
+            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError) as e:
                 message = str(e)
 
                 if recursive and "directory is not empty" in message:
@@ -998,13 +1035,13 @@ class DatabricksPath:
         elif volume_name:
             try:
                 sdk.volumes.delete(f"{catalog_name}.{schema_name}.{volume_name}")
-            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
                 if not allow_not_found:
                     raise
         elif schema_name:
             try:
                 sdk.schemas.delete(f"{catalog_name}.{schema_name}", force=True)
-            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
                 if not allow_not_found:
                     raise
 
@@ -1064,7 +1101,7 @@ class DatabricksPath:
                             yield from base._ls_volume(recursive=recursive)
                         else:
                             yield base
-                except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+                except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
                     if not allow_not_found:
                         raise
             elif schema_name is None:
@@ -1082,7 +1119,7 @@ class DatabricksPath:
                             yield from base._ls_volume(recursive=recursive)
                         else:
                             yield base
-                except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+                except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
                     if not allow_not_found:
                         raise
             else:
@@ -1100,7 +1137,7 @@ class DatabricksPath:
                             yield from base._ls_volume(recursive=recursive)
                         else:
                             yield base
-                except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+                except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
                     if not allow_not_found:
                         raise
         else:
@@ -1121,7 +1158,7 @@ class DatabricksPath:
                         yield from base._ls_volume(recursive=recursive)
                     else:
                         yield base
-            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+            except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
                 if not allow_not_found:
                     raise
 
@@ -1140,7 +1177,7 @@ class DatabricksPath:
                     _is_dir=is_dir,
                     _size=info.size,
                 )
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             if not allow_not_found:
                 raise
 
@@ -1158,7 +1195,7 @@ class DatabricksPath:
                     _is_dir=info.is_dir,
                     _size=info.file_size,
                 )
-        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied):
+        except (NotFound, ResourceDoesNotExist, BadRequest, PermissionDenied, InternalError):
             if not allow_not_found:
                 raise
 
@@ -1225,19 +1262,10 @@ class DatabricksPath:
         else:
             raise FileNotFoundError(f"Path {self} does not exist, or dest is not same file or folder type")
 
-    def write_bytes(self, data: bytes):
-        if hasattr(data, "read"):
-            data = data.read()
-
-        with self.open("wb") as f:
-            f.write_all_bytes(data=data)
-
     def temporary_credentials(
         self,
         operation: Optional["PathOperation"] = None
     ):
-        from databricks.sdk.service.catalog import PathOperation
-
         if self.kind != DatabricksPathKind.VOLUME:
             raise ValueError(f"Cannot generate temporary credentials for {repr(self)}")
 
@@ -1250,6 +1278,14 @@ class DatabricksPath:
             operation=operation or PathOperation.PATH_READ,
         )
 
+    def read_bytes(self, use_cache: bool = False):
+        with self.open("rb") as f:
+            return f.read_all_bytes(use_cache=use_cache)
+
+    def write_bytes(self, data: Union[bytes, IO[bytes]]):
+        with self.open("wb") as f:
+            f.write_all_bytes(data=data)
+
     # -------------------------
     # Data ops (Arrow / Pandas / Polars)
     # -------------------------
@@ -1358,7 +1394,6 @@ class DatabricksPath:
         table: pa.Table,
         file_format: Optional[FileFormat] = None,
         batch_size: Optional[int] = None,
-        **kwargs
     ):
         """Write an Arrow table to the path, sharding if needed.
 
@@ -1366,7 +1401,6 @@ class DatabricksPath:
             table: Arrow table to write.
             file_format: Optional file format override.
             batch_size: Optional batch size for writes.
-            **kwargs: Format-specific options.
 
         Returns:
             The DatabricksPath instance.
@@ -1379,7 +1413,11 @@ class DatabricksPath:
                     part_path = connected / f"{seed}-{i:05d}-{_rand_str(4)}.parquet"
 
                     with part_path.open(mode="wb") as f:
-                        f.write_arrow_batch(batch, file_format=file_format)
+                        f.write_arrow_batch(
+                            batch,
+                            file_format=file_format,
+                            batch_size=batch_size,
+                        )
 
                 return connected
 
@@ -1389,7 +1427,6 @@ class DatabricksPath:
                         table,
                         file_format=file_format,
                         batch_size=batch_size,
-                        **kwargs
                     )
 
         return self
@@ -1399,7 +1436,6 @@ class DatabricksPath:
         file_format: Optional[FileFormat] = None,
         batch_size: Optional[int] = None,
         concat: bool = True,
-        **kwargs
     ):
         """Read the path into a pandas DataFrame.
 
@@ -1407,7 +1443,6 @@ class DatabricksPath:
             file_format: Optional file format override.
             batch_size: Optional batch size for reads.
             concat: Whether to concatenate results for directories.
-            **kwargs: Format-specific options.
 
         Returns:
             A pandas DataFrame or list of DataFrames if concat=False.
@@ -1417,14 +1452,12 @@ class DatabricksPath:
                 file_format=file_format,
                 batch_size=batch_size,
                 concat=True,
-                **kwargs
             ).to_pandas()
 
         tables = self.read_arrow_table(
             batch_size=batch_size,
             file_format=file_format,
             concat=False,
-            **kwargs
         )
 
         return [t.to_pandas() for t in tables]  # type: ignore[arg-type]
@@ -1434,7 +1467,6 @@ class DatabricksPath:
         df: PandasDataFrame,
         file_format: Optional[FileFormat] = None,
         batch_size: Optional[int] = None,
-        **kwargs
     ):
         """Write a pandas DataFrame to the path.
 
@@ -1442,7 +1474,6 @@ class DatabricksPath:
             df: pandas DataFrame to write.
             file_format: Optional file format override.
             batch_size: Optional batch size for writes.
-            **kwargs: Format-specific options.
 
         Returns:
             The DatabricksPath instance.
@@ -1463,7 +1494,6 @@ class DatabricksPath:
                             batch,
                             file_format=file_format,
                             batch_size=batch_size,
-                            **kwargs
                         )
             else:
                 with connected.open(mode="wb", clone=False) as f:
@@ -1471,7 +1501,6 @@ class DatabricksPath:
                         df,
                         file_format=file_format,
                         batch_size=batch_size,
-                        **kwargs
                     )
 
         return self
@@ -1523,7 +1552,6 @@ class DatabricksPath:
         df,
         file_format: Optional[FileFormat] = None,
         batch_size: Optional[int] = None,
-        **kwargs
     ):
         """
         Write Polars to a DatabricksPath.
@@ -1538,7 +1566,6 @@ class DatabricksPath:
             df: polars DataFrame or LazyFrame to write.
             file_format: Optional file format override.
             batch_size: Optional rows per part for directory sinks.
-            **kwargs: Format-specific options.
 
         Returns:
             The DatabricksPath instance.
@@ -1552,7 +1579,7 @@ class DatabricksPath:
         with self.connect() as connected:
             if connected.is_dir_sink():
                 seed = int(time.time() * 1000)
-                rows_per_part = batch_size or 1_000_000
+                rows_per_part = batch_size or 1024 * 1024
 
                 # Always parquet for directory sinks (lake layout standard)
                 for i, chunk in enumerate(df.iter_slices(n_rows=rows_per_part)):
@@ -1563,7 +1590,6 @@ class DatabricksPath:
                             df,
                             file_format=file_format,
                             batch_size=batch_size,
-                            **kwargs
                         )
             else:
                 with connected.open(mode="wb", clone=False) as f:
@@ -1571,11 +1597,33 @@ class DatabricksPath:
                         df,
                         file_format=file_format,
                         batch_size=batch_size,
-                        **kwargs
                     )
 
         return self
 
+    def read_pickle(
+        self,
+    ) -> Any:
+        content = self.read_bytes()
+        obj = dill.loads(content)
+
+        return obj
+
+    def write_pickle(
+        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_bytes(data=buffer.getvalue())
+
+
     def sql(
         self,
         query: str,
@@ -1637,19 +1685,7 @@ class DatabricksPath:
                 "Invalid engine %s, must be in duckdb, polars" % engine
             )
 
-
 if databricks is not None:
-    @register_converter(DatabricksPath, pa.Table)
-    def databricks_path_to_arrow_table(
-        data: DatabricksPath,
-        options: Optional[CastOptions] = None,
-    ) -> pa.Table:
-        return cast_arrow_tabular(
-            data.read_arrow_table(),
-            options
-        )
-
-
     @register_converter(DatabricksPath, ds.Dataset)
     def databricks_path_to_arrow_table(
         data: DatabricksPath,
@@ -1658,6 +1694,16 @@ if databricks is not None:
         return data.arrow_dataset()
 
 
+    @pandas_converter(DatabricksPath, PandasDataFrame)
+    def databricks_path_to_pandas(
+        data: DatabricksPath,
+        options: Optional[CastOptions] = None,
+    ) -> PolarsDataFrame:
+        return cast_pandas_dataframe(
+            data.read_pandas(),
+            options
+        )
+
     @polars_converter(DatabricksPath, PolarsDataFrame)
     def databricks_path_to_polars(
         data: DatabricksPath,