ygg 0.1.29__tar.gz → 0.1.31__tar.gz

{ygg-0.1.29 → ygg-0.1.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.29
+Version: 0.1.31
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.29 → ygg-0.1.31}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ygg"
-version = "0.1.29"
+version = "0.1.31"
 description = "Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks"
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }

{ygg-0.1.29 → ygg-0.1.31}/src/ygg.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.29
+Version: 0.1.31
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.29 → ygg-0.1.31}/src/ygg.egg-info/SOURCES.txt

@@ -21,7 +21,10 @@ src/yggdrasil/databricks/sql/exceptions.py
 src/yggdrasil/databricks/sql/statement_result.py
 src/yggdrasil/databricks/sql/types.py
 src/yggdrasil/databricks/workspaces/__init__.py
-src/yggdrasil/databricks/workspaces/databricks_path.py
+src/yggdrasil/databricks/workspaces/filesytem.py
+src/yggdrasil/databricks/workspaces/io.py
+src/yggdrasil/databricks/workspaces/path.py
+src/yggdrasil/databricks/workspaces/path_kind.py
 src/yggdrasil/databricks/workspaces/workspace.py
 src/yggdrasil/dataclasses/__init__.py
 src/yggdrasil/dataclasses/dataclass.py

{ygg-0.1.29 → ygg-0.1.31}/src/yggdrasil/databricks/compute/cluster.py

@@ -141,7 +141,7 @@ class Cluster(WorkspaceService):
         source: Optional[PythonEnv] = None,
         cluster_id: Optional[str] = None,
         cluster_name: Optional[str] = None,
-        single_user_name: Optional[str] = None,
+        single_user_name: Optional[str] = "current",
         runtime_engine: Optional["RuntimeEngine"] = None,
         libraries: Optional[list[str]] = None,
         **kwargs
@@ -153,7 +153,6 @@ class Cluster(WorkspaceService):
         libraries.extend([
             _ for _ in [
                 "ygg",
-                "dill",
                 "uv",
             ] if _ not in libraries
         ])
@@ -165,11 +164,22 @@ class Cluster(WorkspaceService):
         elif python_version[1] < 11:
             python_version = None
 
+        current_user_name = self.workspace.current_user.user_name
+
+        if single_user_name == "current":
+            single_user_name = current_user_name
+
+        cluster_id = cluster_id or self.cluster_id
+        cluster_name = cluster_name or self.cluster_name
+
+        if not cluster_id and not cluster_name:
+            cluster_name = current_user_name
+
         inst = self.create_or_update(
             cluster_id=cluster_id,
-            cluster_name=cluster_name or self.cluster_name or self.workspace.current_user.user_name,
+            cluster_name=cluster_name,
             python_version=python_version,
-            single_user_name=single_user_name or self.workspace.current_user.user_name,
+            single_user_name=single_user_name,
             runtime_engine=runtime_engine or RuntimeEngine.PHOTON,
             libraries=libraries,
             **kwargs
@@ -180,11 +190,10 @@ class Cluster(WorkspaceService):
     def pull_python_environment(
         self,
         name: Optional[str] = None,
-        target: Optional[PythonEnv] = None,
+        target: PythonEnv | str | None = None,
     ):
         with self.context() as c:
             m = c.remote_metadata
-            requirements = m.requirements
             version_info = m.version_info
 
         python_version = ".".join(str(_) for _ in version_info)
@@ -192,14 +201,20 @@ class Cluster(WorkspaceService):
         if target is None:
             target = PythonEnv.create(
                 name=name or self.name,
-                requirements=requirements,
                 python=python_version
             )
-        else:
-            target.update(
-                requirements=requirements,
-                python=python_version,
-            )
+        elif isinstance(target, str):
+            if target.casefold() == "current":
+                target = PythonEnv.get_current()
+            else:
+                target = PythonEnv.create(
+                    name=target,
+                    python=python_version
+                )
+
+        target.update(
+            python=python_version,
+        )
 
         return target
     
@@ -670,6 +685,9 @@ class Cluster(WorkspaceService):
             def h(z): ...
         """
         def decorator(func: Callable):
+            if os.getenv("DATABRICKS_RUNTIME_VERSION") is not None:
+                return func
+
             context = self.context(language=language or Language.PYTHON)
             serialized = CallableSerde.from_callable(func)
 
@@ -850,8 +868,11 @@ class Cluster(WorkspaceService):
                 target_path = self.workspace.shared_cache_path(
                     suffix=f"/clusters/{self.cluster_id}/{os.path.basename(value)}"
                 )
-                self.workspace.upload_local_path(local_path=value, target_path=target_path)
-                value = target_path
+
+                with open(value, mode="rb") as f:
+                    target_path.write_bytes(f.read())
+
+                value = str(target_path)
             elif "." in value and not "/" in value:
                 value = value.split(".")[0]
 
@@ -865,13 +886,12 @@ class Cluster(WorkspaceService):
 
             repo = None
 
-            if pip_settings.extra_index_url:
-                if (
-                    value.startswith("datamanagement")
-                    or value.startswith("TSSecrets")
-                    or value.startswith("tgp_")
-                ):
-                    repo = pip_settings.extra_index_url
+            if pip_settings.extra_index_url and (
+                value.startswith("datamanagement")
+                or value.startswith("TSSecrets")
+                or value.startswith("tgp_")
+            ):
+                repo = pip_settings.extra_index_url
 
             return Library(
                 pypi=PythonPyPiLibrary(

{ygg-0.1.29 → ygg-0.1.31}/src/yggdrasil/databricks/compute/execution_context.py

@@ -273,7 +273,6 @@ print(json.dumps(meta))"""
         print_stdout: Optional[bool] = True,
         timeout: Optional[dt.timedelta] = None,
         command: Optional[str] = None,
-        use_dill: Optional[bool] = None
     ) -> Any:
         if self.is_in_databricks_environment():
             args = args or []
@@ -291,16 +290,18 @@ print(json.dumps(meta))"""
 
         serialized = CallableSerde.from_callable(func)
 
-        self.install_temporary_libraries(libraries=serialized.pkg_root)
+        if serialized.pkg_root:
+            self.install_temporary_libraries(libraries=serialized.pkg_root)
 
-        # Use dill of same version
         current_version = (sys.version_info.major, sys.version_info.minor)
 
-        if use_dill is None:
-            if current_version == self.cluster.python_version:
-                use_dill = True
-            else:
-                use_dill = False
+        if current_version != self.cluster.python_version[:2]:
+            raise RuntimeError(
+                f"Cannot execute callable: local Python version "
+                f"{current_version[0]}.{current_version[1]} does not match "
+                f"remote cluster Python version "
+                f"{self.cluster.python_version[0]}.{self.cluster.python_version[1]}"
+            )
 
         result_tag = "<<<RESULT>>>"
 
@@ -340,7 +341,6 @@ print(json.dumps(meta))"""
                     print_stdout=print_stdout,
                     timeout=timeout,
                     command=command,
-                    use_dill=use_dill
                 )
             raise remote_module_error
 
@@ -497,7 +497,6 @@ with zipfile.ZipFile(buf, "r") as zf:
             ]
 
         resolved = resolve_local_lib_path(libraries)
-        resolved_str = str(resolved)
 
         remote_site_packages_path = self.remote_metadata.site_packages_path
         if resolved.is_dir():

{ygg-0.1.29 → ygg-0.1.31}/src/yggdrasil/databricks/compute/remote.py

@@ -4,9 +4,12 @@ from typing import (
     Callable,
     Optional,
     TypeVar,
-    List,
+    List, TYPE_CHECKING,
 )
 
+if TYPE_CHECKING:
+    from .cluster import Cluster
+
 from ..workspaces.workspace import Workspace
 
 ReturnType = TypeVar("ReturnType")
@@ -23,16 +26,17 @@ def databricks_remote_compute(
     env_keys: Optional[List[str]] = None,
     **options
 ) -> Callable[[Callable[..., ReturnType]], Callable[..., ReturnType]]:
-    from .. import Cluster
-
     if isinstance(workspace, str):
         workspace = Workspace(host=workspace)
 
     if cluster is None:
-        if cluster_id:
-            cluster = Cluster(workspace=workspace, cluster_id=cluster_id)
+        if cluster_id or cluster_name:
+            cluster = workspace.clusters(
+                cluster_id=cluster_id,
+                cluster_name=cluster_name
+            )
         else:
-            cluster = Cluster.replicated_current_environment(
+            cluster = workspace.clusters().replicated_current_environment(
                 workspace=workspace,
                 cluster_name=cluster_name
             )

{ygg-0.1.29 → ygg-0.1.31}/src/yggdrasil/databricks/jobs/config.py

@@ -2,11 +2,11 @@ import builtins
 import dataclasses
 import datetime as dt
 import inspect
+import logging
 from dataclasses import dataclass, fields
 from enum import Enum
 from inspect import isclass
-from typing import Any, Dict, List, get_type_hints, Optional, get_origin
-import logging
+from typing import Any, Dict, List, get_type_hints, get_origin
 
 from ...libs.sparklib import SparkSession
 from ...types.cast.registry import convert
@@ -308,31 +308,3 @@ class NotebookConfig:
                 spark_session.conf.set("spark.sql.session.timeZone", "UTC")
 
         return cls.from_environment()
-
-
-class ExampleEnum(Enum):
-    """Example enum for widget demonstration"""
-    OPTION1 = "option1"
-    OPTION2 = "option2"
-    OPTION3 = "option3"
-
-
-@dataclass
-class CompleteNotebookConfig(NotebookConfig):
-    """Example JobConfig with various field types to demonstrate widget handling"""
-    # Basic types
-    text_field: str
-    integer_field: int = 42
-    float_field: float = 3.14
-    boolean_field: bool = True
-
-    # Special types
-    date_field: dt.date = dt.date(2023, 1, 1)
-    datetime_field: dt.datetime = dt.datetime(2023, 1, 1, 12, 0, 0)
-    enum_field: ExampleEnum = ExampleEnum.OPTION1
-
-    # Collection types
-    list_of_strings: List[str] = None  # Will be displayed as multiselect
-
-    # Optional fields
-    optional_text: Optional[str] = None

{ygg-0.1.29 → ygg-0.1.31}/src/yggdrasil/databricks/sql/engine.py

@@ -11,6 +11,7 @@ import pyarrow.parquet as pq
 
 from .statement_result import StatementResult
 from .types import column_info_to_arrow_field
+from .. import DatabricksPathKind
 from ..workspaces import WorkspaceService
 from ...libs.databrickslib import databricks_sdk
 from ...libs.sparklib import SparkSession, SparkDataFrame, pyspark
@@ -374,8 +375,9 @@ class SQLEngine(WorkspaceService):
             data = convert(data, pa.Table, options=cast_options, target_field=existing_schema)
 
             # Write in temp volume
-            databricks_tmp_path = connected.path(
-                "/Volumes", catalog_name, schema_name, "tmp", transaction_id, "data.parquet",
+            databricks_tmp_path = connected.dbfs_path(
+                kind=DatabricksPathKind.VOLUME,
+                parts=[catalog_name, schema_name, "tmp", transaction_id, "data.parquet"]
             )
             databricks_tmp_folder = databricks_tmp_path.parent
 

{ygg-0.1.29 → ygg-0.1.31}/src/yggdrasil/databricks/sql/statement_result.py

@@ -96,7 +96,7 @@ class StatementResult:
         if self.is_spark_sql:
             return self._response
 
-        if not self.done and self.statement_id and time.time() - self._response_refresh_time > delay:
+        if self.statement_id and not self.done and time.time() - self._response_refresh_time > delay:
             self.response = self.workspace.sdk().statement_execution.get_statement(self.statement_id)
 
         return self._response
@@ -143,11 +143,25 @@ class StatementResult:
 
     @property
     def done(self):
-        return self.persisted or self.state in [StatementState.CANCELED, StatementState.CLOSED, StatementState.FAILED, StatementState.SUCCEEDED]
+        if self.persisted:
+            return True
+
+        if self._response is None:
+            return False
+
+        return self._response.status.state in [
+            StatementState.CANCELED, StatementState.CLOSED, StatementState.FAILED, StatementState.SUCCEEDED
+        ]
 
     @property
     def failed(self):
-        return self.state in [StatementState.CANCELED, StatementState.FAILED]
+        if self.persisted:
+            return True
+
+        if self._response is None:
+            return False
+
+        return self._response.status.state in [StatementState.CANCELED, StatementState.FAILED]
 
     @property
     def persisted(self):
@@ -163,6 +177,7 @@ class StatementResult:
             self, self.disposition, Disposition.EXTERNAL_LINKS
         )
 
+        self.wait()
         result_data = self.result
         wsdk = self.workspace.sdk()
 

{ygg-0.1.29 → ygg-0.1.31}/src/yggdrasil/databricks/sql/types.py

@@ -102,6 +102,18 @@ def _split_top_level_commas(s: str):
     return parts
 
 
+def _safe_bytes(obj):
+    if not isinstance(obj, bytes):
+        if not obj:
+            return b""
+
+        if not isinstance(obj, str):
+            obj = str(obj)
+
+        return obj.encode("utf-8")
+    return obj
+
+
 def parse_sql_type_to_pa(type_str: str) -> pa.DataType:
     """
     Adapted parser that:
@@ -170,6 +182,10 @@ def column_info_to_arrow_field(col: Union[SQLColumnInfo, CatalogColumnInfo]):
     if isinstance(col, CatalogColumnInfo):
         parsed = json.loads(col.type_json)
         md = parsed.get("metadata", {}) or {}
+        md = {
+            _safe_bytes(k): _safe_bytes(v)
+            for k, v in md.items()
+        }
         nullable = col.nullable
     elif isinstance(col, SQLColumnInfo):
         md = {}

ygg-0.1.31/src/yggdrasil/databricks/workspaces/__init__.py

@@ -0,0 +1,4 @@
+from .workspace import *
+from .path import *
+from .io import *
+

ygg-0.1.31/src/yggdrasil/databricks/workspaces/filesytem.py

@@ -0,0 +1,161 @@
+__all__ = [
+    "DatabricksFileSystem",
+    "DatabricksFileSystemHandler"
+]
+
+from typing import TYPE_CHECKING, Any, Union, List, Optional
+
+from pyarrow import PythonFile
+from pyarrow.fs import FileSystem, FileInfo, FileSelector, PyFileSystem, FileSystemHandler
+
+if TYPE_CHECKING:
+    from ..workspaces.workspace import Workspace
+    from .path import DatabricksPath
+
+
+class DatabricksFileSystemHandler(FileSystemHandler):
+
+    def __init__(
+        self,
+        workspace: "Workspace",
+    ):
+        super().__init__()
+        self.workspace = workspace
+
+    def __enter__(self):
+        return self.connect(clone=True)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.workspace.__exit__(exc_type, exc_val, exc_tb)
+
+    def _parse_path(self, obj: Any) -> "DatabricksPath":
+        from .path import DatabricksPath
+
+        return DatabricksPath.parse(obj, workspace=self.workspace)
+
+    def connect(self, clone: bool = True):
+        workspace = self.connect(clone=clone)
+
+        if clone:
+            return DatabricksFileSystemHandler(
+                workspace=workspace
+            )
+
+        self.workspace = workspace
+        return self
+
+    def close(self):
+        self.workspace.close()
+
+    def copy_file(self, src, dest, *, chunk_size: int = 4 * 1024 * 1024):
+        src = self._parse_path(src)
+        dest = self._parse_path(dest)
+
+        with src.open("rb") as r, dest.open("wb") as w:
+            while True:
+                chunk = r.read(chunk_size)
+                if not chunk:
+                    break
+                w.write(chunk)
+
+    def create_dir(self, path, *args, recursive: bool = True, **kwargs):
+        return self._parse_path(path).mkdir(parents=recursive)
+
+    def delete_dir(self, path):
+        return self._parse_path(path).rmdir(recursive=True)
+
+    def delete_dir_contents(self, path, *args, accept_root_dir: bool = False, **kwargs):
+        return self._parse_path(path).rmdir(recursive=True)
+
+    def delete_root_dir_contents(self):
+        return self.delete_dir_contents("/", accept_root_dir=True)
+
+    def delete_file(self, path):
+        return self._parse_path(path).rmfile()
+
+    def equals(self, other: FileSystem):
+        return self == other
+
+    def from_uri(self, uri):
+        uri = self._parse_path(uri)
+
+        return self.__class__(
+            workspace=uri.workspace
+        )
+
+    def get_file_info(
+        self,
+        paths_or_selector: Union[FileSelector, str, "DatabricksPath", List[Union[str, "DatabricksPath"]]]
+    ) -> Union[FileInfo, List[FileInfo]]:
+        from .path import DatabricksPath
+
+        if isinstance(paths_or_selector, (str, DatabricksPath)):
+            result = self._parse_path(paths_or_selector).file_info
+
+            return result
+
+        if isinstance(paths_or_selector, FileSelector):
+            return self.get_file_info_selector(paths_or_selector)
+
+        return [
+            self.get_file_info(obj)
+            for obj in paths_or_selector
+        ]
+
+    def get_file_info_selector(
+        self,
+        selector: FileSelector
+    ):
+        base_dir = self._parse_path(selector.base_dir)
+
+        return [
+            p.file_info
+            for p in base_dir.ls(
+                recursive=selector.recursive,
+                allow_not_found=selector.allow_not_found
+            )
+        ]
+
+    def get_type_name(self):
+        return "dbfs"
+
+    def move(self, src, dest):
+        src = self._parse_path(src)
+
+        src.copy_to(dest)
+
+        src.remove(recursive=True)
+
+    def normalize_path(self, path):
+        return self._parse_path(path).full_path()
+
+    def open(
+        self,
+        path,
+        mode: str = "r+",
+        encoding: Optional[str] = None,
+    ):
+        return self._parse_path(path).open(mode=mode, encoding=encoding, clone=False)
+
+    def open_append_stream(self, path, compression='detect', buffer_size=None, metadata=None):
+        return self._parse_path(path).open(mode="ab")
+
+    def open_input_file(self, path, mode: str = "rb", **kwargs):
+        buf = self._parse_path(path).open(mode=mode).connect(clone=True)
+
+        return PythonFile(
+            buf,
+            mode=mode
+        )
+
+    def open_input_stream(self, path, compression='detect', buffer_size=None):
+        return self._parse_path(path).open(mode="rb")
+
+    def open_output_stream(self, path, compression='detect', buffer_size=None, metadata=None):
+        return self._parse_path(path).open(mode="wb")
+
+
+class DatabricksFileSystem(PyFileSystem):
+
+    def __init__(self, handler): # real signature unknown; restored from __doc__
+        super().__init__(handler)