ygg 0.1.28__tar.gz → 0.1.30__tar.gz

{ygg-0.1.28 → ygg-0.1.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.28
+Version: 0.1.30
 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.28 → ygg-0.1.30}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ygg"
-version = "0.1.28"
+version = "0.1.30"
 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.28 → ygg-0.1.30}/src/ygg.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.28
+Version: 0.1.30
 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.28 → ygg-0.1.30}/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
     
@@ -646,7 +661,6 @@ class Cluster(WorkspaceService):
         self,
         _func: Optional[Callable] = None,
         *,
-        before: Optional[Callable] = None,
         language: Optional["Language"] = None,
         env_keys: Optional[List[str]] = None,
         env_variables: Optional[Dict[str, str]] = None,
@@ -671,17 +685,17 @@ 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)
-            do_before = CallableSerde.from_callable(before)
 
             @functools.wraps(func)
             def wrapper(*args, **kwargs):
                 if os.getenv("DATABRICKS_RUNTIME_VERSION") is not None:
                     return func(*args, **kwargs)
 
-                do_before()
-
                 return context.execute(
                     obj=serialized,
                     args=list(args),
@@ -854,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]
 
@@ -869,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.28 → ygg-0.1.30}/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.package_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.28 → ygg-0.1.30}/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.28 → ygg-0.1.30}/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.28 → ygg-0.1.30}/src/yggdrasil/databricks/sql/statement_result.py

@@ -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()