dataproc-spark-connect 1.0.0rc1__py2.py3-none-any.whl → 1.0.0rc3__py2.py3-none-any.whl

{dataproc_spark_connect-1.0.0rc1.dist-info → dataproc_spark_connect-1.0.0rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataproc-spark-connect
-Version: 1.0.0rc1
+Version: 1.0.0rc3
 Summary: Dataproc client library for Spark Connect
 Home-page: https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python
 Author: Google LLC

{dataproc_spark_connect-1.0.0rc1.dist-info → dataproc_spark_connect-1.0.0rc3.dist-info}/RECORD

@@ -1,13 +1,13 @@
-dataproc_spark_connect-1.0.0rc1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+dataproc_spark_connect-1.0.0rc3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 google/cloud/dataproc_spark_connect/__init__.py,sha256=dIqHNWVWWrSuRf26x11kX5e9yMKSHCtmI_GBj1-FDdE,1101
-google/cloud/dataproc_spark_connect/environment.py,sha256=UICy9XyqAxL-cryVWx7GZPRAxoir5LKk0dtqqY_l--c,2307
+google/cloud/dataproc_spark_connect/environment.py,sha256=l1wWiHMHtBQ9YonE-kHTpaZlN9vLE4fyJSTn7RZP6kA,2503
 google/cloud/dataproc_spark_connect/exceptions.py,sha256=WF-qdzgdofRwILCriIkjjsmjObZfF0P3Ecg4lv-Hmec,968
 google/cloud/dataproc_spark_connect/pypi_artifacts.py,sha256=gd-VMwiVP-EJuPp9Vf9Shx8pqps3oSKp0hBcSSZQS-A,1575
-google/cloud/dataproc_spark_connect/session.py,sha256=vD9kJXJbkzfKlYt5TFp0umHi6A7ZdheSJNjrqVyL0oo,37432
+google/cloud/dataproc_spark_connect/session.py,sha256=u2QxRLn84EMdpFnv9pI6Y7F8OFQ2mqCllb6AsLd73eo,42402
 google/cloud/dataproc_spark_connect/client/__init__.py,sha256=6hCNSsgYlie6GuVpc5gjFsPnyeMTScTpXSPYqp1fplY,615
 google/cloud/dataproc_spark_connect/client/core.py,sha256=GRc4OCTBvIvdagjxOPoDO22vLtt8xDSerdREMRDeUBY,4659
 google/cloud/dataproc_spark_connect/client/proxy.py,sha256=qUZXvVY1yn934vE6nlO495XUZ53AUx9O74a9ozkGI9U,8976
-dataproc_spark_connect-1.0.0rc1.dist-info/METADATA,sha256=D1e6sjZ8-hVccXMjHwkoX0OPOOW876hzsQi25WxMbMI,3468
-dataproc_spark_connect-1.0.0rc1.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-dataproc_spark_connect-1.0.0rc1.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
-dataproc_spark_connect-1.0.0rc1.dist-info/RECORD,,
+dataproc_spark_connect-1.0.0rc3.dist-info/METADATA,sha256=yUMy-S__bGisjsxduDg_VmJzF3ZKqCybaNUwn5DMymo,3468
+dataproc_spark_connect-1.0.0rc3.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+dataproc_spark_connect-1.0.0rc3.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
+dataproc_spark_connect-1.0.0rc3.dist-info/RECORD,,

google/cloud/dataproc_spark_connect/environment.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import sys
 from typing import Callable, Tuple, List
 
 
@@ -46,6 +47,18 @@ def is_jetbrains_ide() -> bool:
     return "jetbrains" in os.getenv("TERMINAL_EMULATOR", "").lower()
 
 
+def is_interactive():
+    return hasattr(sys, "ps1")
+
+
+def is_terminal():
+    return sys.stdin.isatty()
+
+
+def is_interactive_terminal():
+    return is_interactive() and is_terminal()
+
+
 def get_client_environment_label() -> str:
     """
     Map current environment to a standardized client label.

google/cloud/dataproc_spark_connect/session.py

@@ -24,8 +24,11 @@ import threading
 import time
 import uuid
 import tqdm
+from packaging import version
+from tqdm import tqdm as cli_tqdm
+from tqdm.notebook import tqdm as notebook_tqdm
 from types import MethodType
-from typing import Any, cast, ClassVar, Dict, Optional, Union
+from typing import Any, cast, ClassVar, Dict, Iterable, Optional, Union
 
 from google.api_core import retry
 from google.api_core.client_options import ClientOptions
@@ -104,12 +107,14 @@ class DataprocSparkSession(SparkSession):
     """
 
     _DEFAULT_RUNTIME_VERSION = "3.0"
+    _MIN_RUNTIME_VERSION = "3.0"
 
     _active_s8s_session_uuid: ClassVar[Optional[str]] = None
     _project_id = None
     _region = None
     _client_options = None
     _active_s8s_session_id: ClassVar[Optional[str]] = None
+    _execution_progress_bar = dict()
 
     class Builder(SparkSession.Builder):
 
@@ -247,7 +252,11 @@ class DataprocSparkSession(SparkSession):
             assert self._channel_builder is not None
             session = DataprocSparkSession(connection=self._channel_builder)
 
+            # Register handler for Cell Execution Progress bar
+            session._register_progress_execution_handler()
+
             DataprocSparkSession._set_default_and_active_session(session)
+
             return session
 
         def __create(self) -> "DataprocSparkSession":
@@ -262,6 +271,9 @@ class DataprocSparkSession(SparkSession):
 
                 dataproc_config: Session = self._get_dataproc_config()
 
+                # Check runtime version compatibility before creating session
+                self._check_runtime_compatibility(dataproc_config)
+
                 session_id = self.generate_dataproc_session_id()
                 dataproc_config.name = f"projects/{self._project_id}/locations/{self._region}/sessions/{session_id}"
                 logger.debug(
@@ -593,6 +605,43 @@ class DataprocSparkSession(SparkSession):
                         stacklevel=3,
                     )
 
+        def _check_runtime_compatibility(self, dataproc_config):
+            """Check if runtime version 3.0 client is compatible with older runtime versions.
+
+            Runtime version 3.0 clients do not support older runtime versions (pre-3.0).
+            There is no backward or forward compatibility between different runtime versions.
+
+            Args:
+                dataproc_config: The Session configuration containing runtime version
+
+            Raises:
+                DataprocSparkConnectException: If server is using pre-3.0 runtime version
+            """
+            runtime_version = dataproc_config.runtime_config.version
+
+            if not runtime_version:
+                return
+
+            logger.debug(f"Detected server runtime version: {runtime_version}")
+
+            # Parse runtime version to check if it's below minimum supported version
+            try:
+                server_version = version.parse(runtime_version)
+                min_version = version.parse(
+                    DataprocSparkSession._MIN_RUNTIME_VERSION
+                )
+
+                if server_version < min_version:
+                    raise DataprocSparkConnectException(
+                        f"Specified {runtime_version} Dataproc Runtime version is not supported, "
+                        f"use {DataprocSparkSession._MIN_RUNTIME_VERSION} version or higher."
+                    )
+            except version.InvalidVersion:
+                # If we can't parse the version, log a warning but continue
+                logger.warning(
+                    f"Could not parse runtime version: {runtime_version}"
+                )
+
         def _display_view_session_details_button(self, session_id):
             try:
                 session_url = f"https://console.cloud.google.com/dataproc/interactive/sessions/{session_id}/locations/{self._region}?project={self._project_id}"
@@ -682,6 +731,78 @@ class DataprocSparkSession(SparkSession):
             execute_and_fetch_as_iterator_wrapped_method, self.client
         )
 
+        # Patching clearProgressHandlers method to not remove Dataproc Progress Handler
+        clearProgressHandlers_base_method = self.clearProgressHandlers
+
+        def clearProgressHandlers_wrapper_method(_, *args, **kwargs):
+            clearProgressHandlers_base_method(*args, **kwargs)
+
+            self._register_progress_execution_handler()
+
+        self.clearProgressHandlers = MethodType(
+            clearProgressHandlers_wrapper_method, self
+        )
+
+    def _register_progress_execution_handler(self):
+        from pyspark.sql.connect.shell.progress import StageInfo
+
+        def handler(
+            stages: Optional[Iterable[StageInfo]],
+            inflight_tasks: int,
+            operation_id: Optional[str],
+            done: bool,
+        ):
+            if operation_id is None:
+                return
+
+            # Don't build / render progress bar for non-interactive (despite
+            # Ipython or non-IPython)
+            if not environment.is_interactive():
+                return
+
+            total_tasks = 0
+            completed_tasks = 0
+
+            for stage in stages or []:
+                total_tasks += stage.num_tasks
+                completed_tasks += stage.num_completed_tasks
+
+            tqdm_pbar = notebook_tqdm
+            if environment.is_interactive_terminal():
+                tqdm_pbar = cli_tqdm
+
+            # Use a lock to ensure only one thread can access and modify
+            # the shared dictionaries at a time.
+            with self._lock:
+                if operation_id in self._execution_progress_bar:
+                    pbar = self._execution_progress_bar[operation_id]
+                    if pbar.total != total_tasks:
+                        pbar.reset(
+                            total=total_tasks
+                        )  # This force resets the progress bar % too on next refresh
+                else:
+                    pbar = tqdm_pbar(
+                        total=total_tasks,
+                        leave=True,
+                        dynamic_ncols=True,
+                        bar_format="{l_bar}{bar} {n_fmt}/{total_fmt} Tasks",
+                    )
+                    self._execution_progress_bar[operation_id] = pbar
+
+                # To handle skipped or failed tasks.
+                # StageInfo proto doesn't have skipped and failed tasks information to process.
+                if done and completed_tasks < total_tasks:
+                    completed_tasks = total_tasks
+
+                pbar.n = completed_tasks
+                pbar.refresh()
+
+                if done:
+                    pbar.close()
+                    self._execution_progress_bar.pop(operation_id, None)
+
+        self.registerProgressHandler(handler)
+
     @staticmethod
     def _sql_lazy_transformation(req):
         # Select SQL command
@@ -711,6 +832,11 @@ class DataprocSparkSession(SparkSession):
         """
 
     def _display_operation_link(self, operation_id: str):
+        # Don't print per-operation Spark UI link for non-interactive (despite
+        # Ipython or non-IPython)
+        if not environment.is_interactive():
+            return
+
         assert all(
             [
                 operation_id is not None,
@@ -726,12 +852,13 @@ class DataprocSparkSession(SparkSession):
             f"associatedSqlOperationId={operation_id}?project={self._project_id}"
         )
 
+        if environment.is_interactive_terminal():
+            print(f"Spark Query: {url}")
+            return
+
         try:
             from IPython.display import display, HTML
-            from IPython.core.interactiveshell import InteractiveShell
 
-            if not InteractiveShell.initialized():
-                return
             html_element = f"""
               <div>
                   <p><a href="{url}">Spark Query</a> (Operation: {operation_id})</p>