ygg 0.1.44__py3-none-any.whl → 0.1.46__py3-none-any.whl

{ygg-0.1.44.dist-info → ygg-0.1.46.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.44
+Version: 0.1.46
 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.44.dist-info → ygg-0.1.46.dist-info}/RECORD

@@ -1,24 +1,25 @@
-ygg-0.1.44.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.46.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=PfH7Xwt6uue6oqe6S5V8NhDJcVQClkKrBE1KXhdelZc,117
-yggdrasil/version.py,sha256=DRHimXBBU-nUJWLj3HfOg9_glzt2k0luj068qgL5e5E,22
+yggdrasil/version.py,sha256=Jfjdb6r4PUMCYLyBcF_gpjjUzE1bkfL7XTvU1GuAtdA,22
 yggdrasil/databricks/__init__.py,sha256=skctY2c8W-hI81upx9F_PWRe5ishL3hrdiTuizgDjdw,152
 yggdrasil/databricks/compute/__init__.py,sha256=NvdzmaJSNYY1uJthv1hHdBuNu3bD_-Z65DWnaJt9yXg,289
-yggdrasil/databricks/compute/cluster.py,sha256=I2Iy0WInwNifL72jZCPg2WgpfbNQjXV4bF0emRT732g,41137
-yggdrasil/databricks/compute/execution_context.py,sha256=E86b-VF9hZ4WbqHt43ATR-YYvgPZ91_17DQ5oaQFNgg,22688
+yggdrasil/databricks/compute/cluster.py,sha256=HI9811oBCpWeo4V921FVAlRUXKXM4XO7HS9DQVOuzpM,41340
+yggdrasil/databricks/compute/execution_context.py,sha256=SMFjgWqp9CEgruJszxTAmqg3iSREb7OG0nYTUn2H3ec,21946
 yggdrasil/databricks/compute/remote.py,sha256=nEN_Fr1Ouul_iKOf4B5QjEGscYAcl7nHjGsl2toRzrU,2874
 yggdrasil/databricks/jobs/__init__.py,sha256=snxGSJb0M5I39v0y3IR-uEeSlZR248cQ_4DJ1sYs-h8,154
 yggdrasil/databricks/jobs/config.py,sha256=9LGeHD04hbfy0xt8_6oobC4moKJh4_DTjZiK4Q2Tqjk,11557
 yggdrasil/databricks/sql/__init__.py,sha256=y1n5yg-drZ8QVZbEgznsRG24kdJSnFis9l2YfYCsaCM,234
-yggdrasil/databricks/sql/engine.py,sha256=Azx3gKtWOMy3D9I2FhkLmpthZPWAJZ9iZkaDivmt_0s,41002
+yggdrasil/databricks/sql/engine.py,sha256=K5WmGKpXU78JA3UdK8dLxBD_GXKidZJFe7hytuC5UHg,41029
 yggdrasil/databricks/sql/exceptions.py,sha256=Jqd_gT_VyPL8klJEHYEzpv5eHtmdY43WiQ7HZBaEqSk,53
 yggdrasil/databricks/sql/statement_result.py,sha256=_mBolHae0AASfe1Tlld1KTXs-K4-oy9dniHDyR2ILYc,16736
 yggdrasil/databricks/sql/types.py,sha256=5G-BM9_eOsRKEMzeDTWUsWW5g4Idvs-czVCpOCrMhdA,6412
+yggdrasil/databricks/sql/warehouse.py,sha256=1J0dyQLJb-OS1_1xU1eAVZ4CoL2-FhFeowKSvU3RzFc,9773
 yggdrasil/databricks/workspaces/__init__.py,sha256=Ti1I99JTC3koYJaCy8WYvkAox4KdcuMRk8b2rHroWCY,133
 yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IAP2KemXBMPkUoY4w,9805
 yggdrasil/databricks/workspaces/io.py,sha256=Tdde4LaGNJNT50R11OkEYZyNacyIW9QrOXMAicAlIr4,32208
 yggdrasil/databricks/workspaces/path.py,sha256=-XnCD9p42who3DAwnITVE1KyrZUSoXDKHA8iZi-7wk4,47743
 yggdrasil/databricks/workspaces/path_kind.py,sha256=Xc319NysH8_6E9C0Q8nCxDHYG07_SnzyUVKHe0dNdDQ,305
-yggdrasil/databricks/workspaces/workspace.py,sha256=MW-BEyldROqbX9SBbDspvlys_zehJjK5YgM3sGLfW-g,23382
+yggdrasil/databricks/workspaces/workspace.py,sha256=c6CBBun2BskEnsP74pbLVOe_TKXZs4L4r4gPQtIzlQE,23821
 yggdrasil/dataclasses/__init__.py,sha256=6SdfIyTsoM4AuVw5TW4Q-UWXz41EyfsMcpD30cmjbSM,125
 yggdrasil/dataclasses/dataclass.py,sha256=fKokFUnqe4CmXXGMTdF4XDWbCUl_c_-se-UD48L5s1E,6594
 yggdrasil/libs/__init__.py,sha256=ulzk-ZkFUI2Pfo93YKtO8MBsEWtRZzLos7HAxN74R0w,168
@@ -30,10 +31,10 @@ yggdrasil/libs/extensions/__init__.py,sha256=mcXW5Li3Cbprbs4Ci-b5A0Ju0wmLcfvEiFu
 yggdrasil/libs/extensions/polars_extensions.py,sha256=RTkGi8llhPJjX7x9egix7-yXWo2X24zIAPSKXV37SSA,12397
 yggdrasil/libs/extensions/spark_extensions.py,sha256=E64n-3SFTDgMuXwWitX6vOYP9ln2lpGKb0htoBLEZgc,16745
 yggdrasil/pyutils/__init__.py,sha256=tl-LapAc71TV7RMgf2ftKwrzr8iiLOGHeJgA3RvO93w,293
-yggdrasil/pyutils/callable_serde.py,sha256=FtbY2PUBHwDe2IyX28gLJwQYjtrvhvqtTSIIBc5x-dk,23538
+yggdrasil/pyutils/callable_serde.py,sha256=UNxLfhkE1xunbGdos4BF7TOMZbBxDRNMlj8lkZx3xq0,23197
 yggdrasil/pyutils/equality.py,sha256=Xyf8D1dLUCm3spDEir8Zyj7O4US_fBJwEylJCfJ9slI,3080
 yggdrasil/pyutils/exceptions.py,sha256=ssKNm-rjhavHUOZmGA7_1Gq9tSHDrb2EFI-cnBuWgng,3388
-yggdrasil/pyutils/expiring_dict.py,sha256=q9gb09-2EUN-jQZumUw5BXOQGYcj1wb85qKtQlciSxg,5825
+yggdrasil/pyutils/expiring_dict.py,sha256=pr2u25LGwPVbLfsLptiHGovUtYRRo0AMjaJtCtJl7nQ,8477
 yggdrasil/pyutils/modules.py,sha256=B7IP99YqUMW6-DIESFzBx8-09V1d0a8qrIJUDFhhL2g,11424
 yggdrasil/pyutils/parallel.py,sha256=ubuq2m9dJzWYUyKCga4Y_9bpaeMYUrleYxdp49CHr44,6781
 yggdrasil/pyutils/python_env.py,sha256=tuglnjdqHQjNh18qDladVoSEOjCD0RcnMEPYJ0tArOs,50985
@@ -55,8 +56,8 @@ yggdrasil/types/cast/registry.py,sha256=_zdFGmUBB7P-e_LIcJlOxMcxAkXoA-UXB6HqLMgT
 yggdrasil/types/cast/spark_cast.py,sha256=_KAsl1DqmKMSfWxqhVE7gosjYdgiL1C5bDQv6eP3HtA,24926
 yggdrasil/types/cast/spark_pandas_cast.py,sha256=BuTiWrdCANZCdD_p2MAytqm74eq-rdRXd-LGojBRrfU,5023
 yggdrasil/types/cast/spark_polars_cast.py,sha256=btmZNHXn2NSt3fUuB4xg7coaE0RezIBdZD92H8NK0Jw,9073
-ygg-0.1.44.dist-info/METADATA,sha256=_5XOQ56SkOMJb6dcbRc7CyOlp6jewJPNZdDi05_9Ck8,19204
-ygg-0.1.44.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ygg-0.1.44.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
-ygg-0.1.44.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.44.dist-info/RECORD,,
+ygg-0.1.46.dist-info/METADATA,sha256=V8VQXfyZROSBPEU21Rv-xykPqQkg3hhvG4k2OdP8gNc,19204
+ygg-0.1.46.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ygg-0.1.46.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.46.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.46.dist-info/RECORD,,

yggdrasil/databricks/compute/cluster.py

@@ -22,7 +22,7 @@ from typing import Any, Iterator, Optional, Union, List, Callable, Dict, ClassVa
 
 from .execution_context import ExecutionContext
 from ..workspaces.workspace import WorkspaceService, Workspace
-from ... import retry, CallableSerde
+from ... import CallableSerde
 from ...libs.databrickslib import databricks_sdk
 from ...pyutils.equality import dicts_equal, dict_diff
 from ...pyutils.expiring_dict import ExpiringDict
@@ -47,6 +47,7 @@ else:  # pragma: no cover - runtime fallback when SDK is missing
 __all__ = ["Cluster"]
 
 
+LOGGER = logging.getLogger(__name__)
 NAME_ID_CACHE: dict[str, ExpiringDict] = {}
 
 
@@ -72,9 +73,6 @@ def get_cached_cluster_id(
     return existing.get(cluster_name) if existing else None
 
 
-logger = logging.getLogger(__name__)
-
-
 # module-level mapping Databricks Runtime -> (major, minor) Python version
 _PYTHON_BY_DBR: dict[str, tuple[int, int]] = {
     "10.4": (3, 8),
@@ -363,7 +361,8 @@ class Cluster(WorkspaceService):
         tick: float = 0.5,
         timeout: Union[float, dt.timedelta] = 600,
         backoff: int = 2,
-        max_sleep_time: float = 15
+        max_sleep_time: float = 15,
+        wait_libraries: bool = True
     ):
         """Wait for the cluster to exit pending states.
 
@@ -372,6 +371,7 @@ class Cluster(WorkspaceService):
             timeout: Max seconds to wait before timing out.
             backoff: Backoff multiplier for the sleep interval.
             max_sleep_time: Maximum sleep interval in seconds.
+            wait_libraries: Wait libraries to install fully
 
         Returns:
             The current Cluster instance.
@@ -390,7 +390,8 @@ class Cluster(WorkspaceService):
 
             sleep_time = min(max_sleep_time, sleep_time * backoff)
 
-        self.wait_installed_libraries()
+        if wait_libraries:
+            self.wait_installed_libraries()
 
         self.raise_for_status()
 
@@ -638,7 +639,7 @@ class Cluster(WorkspaceService):
             if k in _CREATE_ARG_NAMES
         }
 
-        logger.debug(
+        LOGGER.debug(
             "Creating Databricks cluster %s with %s",
             update_details["cluster_name"],
             update_details,
@@ -646,7 +647,7 @@ class Cluster(WorkspaceService):
 
         self.details = self.clusters_client().create_and_wait(**update_details)
 
-        logger.info(
+        LOGGER.info(
             "Created %s",
             self
         )
@@ -699,7 +700,7 @@ class Cluster(WorkspaceService):
                 for k, v in dict_diff(existing_details, update_details, keys=_EDIT_ARG_NAMES).items()
             }
 
-            logger.debug(
+            LOGGER.debug(
                 "Updating %s with %s",
                 self, diff
             )
@@ -707,7 +708,7 @@ class Cluster(WorkspaceService):
             self.wait_for_status()
             self.clusters_client().edit(**update_details)
 
-            logger.info(
+            LOGGER.info(
                 "Updated %s",
                 self
             )
@@ -811,7 +812,7 @@ class Cluster(WorkspaceService):
         self.wait_for_status()
 
         if not self.is_running:
-            logger.info("Starting %s", self)
+            LOGGER.debug("Starting %s", self)
 
             if wait_timeout:
                 self.clusters_client().start(cluster_id=self.cluster_id)
@@ -819,6 +820,8 @@ class Cluster(WorkspaceService):
             else:
                 self.clusters_client().start(cluster_id=self.cluster_id)
 
+            LOGGER.info("Started %s", self)
+
         return self
 
     def restart(
@@ -832,7 +835,6 @@ class Cluster(WorkspaceService):
         self.wait_for_status()
 
         if self.is_running:
-            logger.info("Restarting %s", self)
             self.details = self.clusters_client().restart_and_wait(cluster_id=self.cluster_id)
             return self.wait_for_status()
 
@@ -846,8 +848,10 @@ class Cluster(WorkspaceService):
         Returns:
             The SDK delete response.
         """
-        logger.info("Deleting %s", self)
-        return self.clusters_client().delete(cluster_id=self.cluster_id)
+        if self.cluster_id:
+            LOGGER.debug("Deleting %s", self)
+            self.clusters_client().delete(cluster_id=self.cluster_id)
+            LOGGER.info("Deleted %s", self)
 
     def context(
         self,
@@ -1150,7 +1154,7 @@ class Cluster(WorkspaceService):
                 if raise_error:
                     raise DatabricksError("Libraries %s in %s failed to install" % (failed, self))
 
-                logger.warning(
+                LOGGER.exception(
                     "Libraries %s in %s failed to install",
                     failed, self
                 )
@@ -1187,7 +1191,7 @@ class Cluster(WorkspaceService):
         Returns:
             The uploaded library argument(s).
         """
-        return self.context().install_temporary_libraries(libraries=libraries)
+        return self.system_context.install_temporary_libraries(libraries=libraries)
 
     def _check_library(
         self,

yggdrasil/databricks/compute/execution_context.py

@@ -12,11 +12,13 @@ import re
 import sys
 import threading
 import zipfile
+from threading import Thread
 from types import ModuleType
 from typing import TYPE_CHECKING, Optional, Any, Callable, List, Dict, Union, Iterable, Tuple
 
 from ...libs.databrickslib import databricks_sdk
 from ...pyutils.exceptions import raise_parsed_traceback
+from ...pyutils.expiring_dict import ExpiringDict
 from ...pyutils.modules import resolve_local_lib_path
 from ...pyutils.callable_serde import CallableSerde
 
@@ -30,7 +32,7 @@ __all__ = [
     "ExecutionContext"
 ]
 
-logger = logging.getLogger(__name__)
+LOGGER = logging.getLogger(__name__)
 
 
 @dc.dataclass
@@ -38,7 +40,6 @@ class RemoteMetadata:
     """Metadata describing the remote cluster execution environment."""
     site_packages_path: Optional[str] = dc.field(default=None)
     os_env: Dict[str, str] = dc.field(default_factory=dict)
-    requirements: Optional[str] = dc.field(default=None)
     version_info: Tuple[int, int, int] = dc.field(default=(0, 0, 0))
 
     def os_env_diff(
@@ -80,6 +81,7 @@ class ExecutionContext:
 
     _was_connected: Optional[bool] = dc.field(default=None, repr=False)
     _remote_metadata: Optional[RemoteMetadata] = dc.field(default=None, repr=False)
+    _uploaded_package_roots: Optional[ExpiringDict] = dc.field(default_factory=ExpiringDict, repr=False)
 
     _lock: threading.RLock = dc.field(default_factory=threading.RLock, init=False, repr=False)
 
@@ -113,7 +115,11 @@ class ExecutionContext:
 
     def __del__(self):
         """Best-effort cleanup for the remote execution context."""
-        self.close()
+        if self.context_id:
+            try:
+                Thread(target=self.close).start()
+            except BaseException:
+                pass
 
     @property
     def remote_metadata(self) -> RemoteMetadata:
@@ -127,9 +133,7 @@ class ExecutionContext:
         with self._lock:
             # double-check after acquiring lock
             if self._remote_metadata is None:
-                cmd = r"""import glob
-import json
-import os
+                cmd = r"""import glob, json, os
 from yggdrasil.pyutils.python_env import PythonEnv
 
 current_env = PythonEnv.get_current()
@@ -144,7 +148,6 @@ os_env = meta["os_env"] = {}
 for k, v in os.environ.items():
     os_env[k] = v
     
-meta["requirements"] = current_env.requirements()
 meta["version_info"] = current_env.version_info
 
 print(json.dumps(meta))"""
@@ -191,7 +194,7 @@ print(json.dumps(meta))"""
         """
         self.cluster.ensure_running()
 
-        logger.debug(
+        LOGGER.debug(
             "Creating Databricks command execution context for %s",
             self.cluster
         )
@@ -217,7 +220,7 @@ print(json.dumps(meta))"""
             The connected ExecutionContext instance.
         """
         if self.context_id is not None:
-            logger.debug(
+            LOGGER.debug(
                 "Execution context already open for %s",
                 self
             )
@@ -235,7 +238,7 @@ print(json.dumps(meta))"""
             raise RuntimeError("Failed to create command execution context")
 
         self.context_id = context_id
-        logger.info(
+        LOGGER.info(
             "Opened execution context for %s",
             self
         )
@@ -247,13 +250,9 @@ print(json.dumps(meta))"""
         Returns:
             None.
         """
-        if self.context_id is None:
+        if not self.context_id:
             return
 
-        logger.debug(
-            "Closing execution context for %s",
-            self
-        )
         try:
             self._workspace_client().command_execution.destroy(
                 cluster_id=self.cluster.cluster_id,
@@ -349,7 +348,7 @@ print(json.dumps(meta))"""
 
         self.connect(language=Language.PYTHON)
 
-        logger.debug(
+        LOGGER.debug(
             "Executing callable %s with %s",
             getattr(func, "__name__", type(func)),
             self,
@@ -386,7 +385,11 @@ print(json.dumps(meta))"""
         )
 
         try:
-            result = serialized.parse_command_result(raw_result, result_tag=result_tag)
+            result = serialized.parse_command_result(
+                raw_result,
+                result_tag=result_tag,
+                workspace=self.cluster.workspace
+            )
         except ModuleNotFoundError as remote_module_error:
             _MOD_NOT_FOUND_RE = re.compile(r"No module named ['\"]([^'\"]+)['\"]")
             module_name = _MOD_NOT_FOUND_RE.search(str(remote_module_error))
@@ -394,12 +397,18 @@ print(json.dumps(meta))"""
             module_name = module_name.split(".")[0]
 
             if module_name and "yggdrasil" not in module_name:
-                self.close()
+                LOGGER.debug(
+                    "Installing missing module %s from local environment",
+                    module_name,
+                )
 
-                self.cluster.install_libraries(
+                self.install_temporary_libraries(
                     libraries=[module_name],
-                    raise_error=True,
-                    restart=True
+                )
+
+                LOGGER.warning(
+                    "Installed missing module %s from local environment",
+                    module_name,
                 )
 
                 return self.execute_callable(
@@ -412,6 +421,7 @@ print(json.dumps(meta))"""
                     timeout=timeout,
                     command=command,
                 )
+
             raise remote_module_error
 
         return result
@@ -446,29 +456,7 @@ print(json.dumps(meta))"""
             timeout=timeout or dt.timedelta(minutes=20)
         )
 
-        try:
-            return self._decode_result(result, result_tag=result_tag, print_stdout=print_stdout)
-        except ModuleNotFoundError as remote_module_error:
-            _MOD_NOT_FOUND_RE = re.compile(r"No module named ['\"]([^'\"]+)['\"]")
-            module_name = _MOD_NOT_FOUND_RE.search(str(remote_module_error))
-            module_name = module_name.group(1) if module_name else None
-            module_name = module_name.split(".")[0]
-
-            if module_name and "yggdrasil" not in module_name:
-                self.close()
-                self.cluster.install_libraries(
-                    libraries=[module_name],
-                    raise_error=True,
-                    restart=True
-                )
-
-                return self.execute_command(
-                    command=command,
-                    timeout=timeout,
-                    result_tag=result_tag,
-                    print_stdout=print_stdout
-                )
-            raise remote_module_error
+        return self._decode_result(result, result_tag=result_tag, print_stdout=print_stdout)
 
     # ------------------------------------------------------------------
     # generic local → remote uploader, via remote python
@@ -589,16 +577,22 @@ with zipfile.ZipFile(buf, "r") as zf:
             ]
 
         resolved = resolve_local_lib_path(libraries)
+        str_resolved = str(resolved)
+        existing = self._uploaded_package_roots.get(str_resolved)
 
-        remote_site_packages_path = self.remote_metadata.site_packages_path
-        if resolved.is_dir():
-            # site-packages/<package_name>/
-            remote_target = posixpath.join(remote_site_packages_path, resolved.name)
-        else:
-            # site-packages/<module_file>
-            remote_target = posixpath.join(remote_site_packages_path, resolved.name)
+        if not existing:
+            remote_site_packages_path = self.remote_metadata.site_packages_path
 
-        self.upload_local_path(resolved, remote_target)
+            if resolved.is_dir():
+                # site-packages/<package_name>/
+                remote_target = posixpath.join(remote_site_packages_path, resolved.name)
+            else:
+                # site-packages/<module_file>
+                remote_target = posixpath.join(remote_site_packages_path, resolved.name)
+
+            self.upload_local_path(resolved, remote_target)
+
+            self._uploaded_package_roots[str_resolved] = remote_target
 
         return libraries
 
@@ -649,16 +643,4 @@ with zipfile.ZipFile(buf, "r") as zf:
         else:
             output = ""
 
-        # result_tag slicing
-        if result_tag:
-            start = output.find(result_tag)
-            if start != -1:
-                content_start = start + len(result_tag)
-                end = output.find(result_tag, content_start)
-                if end != -1:
-                    before = output[:start].strip()
-                    if before and print_stdout:
-                        print(before)
-                    return output[content_start:end]
-
         return output

yggdrasil/databricks/sql/engine.py

@@ -130,7 +130,10 @@ class SQLEngine(WorkspaceService):
         schema_name = schema_name or self.schema_name
         return catalog_name, schema_name, table_name
 
-    def _default_warehouse(self, cluster_size: str = "Small"):
+    def _default_warehouse(
+        self,
+        cluster_size: str = "Small"
+    ):
         """Pick a default SQL warehouse (best-effort) matching the desired size.
 
         Args:
@@ -951,7 +954,7 @@ FROM parquet.`{temp_volume_path}`"""
         """Convert an Arrow Field to a Databricks SQL column DDL fragment."""
         name = field.name
         nullable_str = " NOT NULL" if put_not_null and not field.nullable else ""
-        name_str = f"{name} " if put_name else ""
+        name_str = f"`{name}` " if put_name else ""
 
         comment_str = ""
         if put_comment and field.metadata and b"comment" in field.metadata: