dataproc-spark-connect 0.8.1__tar.gz → 0.8.3__tar.gz

{dataproc_spark_connect-0.8.1 → dataproc_spark_connect-0.8.3}/PKG-INFO

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

{dataproc_spark_connect-0.8.1 → dataproc_spark_connect-0.8.3}/dataproc_spark_connect.egg-info/PKG-INFO

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

{dataproc_spark_connect-0.8.1 → dataproc_spark_connect-0.8.3}/google/cloud/dataproc_spark_connect/session.py

@@ -22,9 +22,10 @@ import re
 import string
 import threading
 import time
-from typing import Any, cast, ClassVar, Dict, Optional, Union
 import uuid
 import tqdm
+from types import MethodType
+from typing import Any, cast, ClassVar, Dict, Optional, Union
 
 from google.api_core import retry
 from google.api_core.client_options import ClientOptions
@@ -234,6 +235,9 @@ class DataprocSparkSession(SparkSession):
                     target=create_session_pbar
                 )
 
+                # Activate Spark Connect mode for Spark client
+                os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1"
+
                 try:
                     if (
                         os.getenv(
@@ -538,19 +542,57 @@ class DataprocSparkSession(SparkSession):
 
         super().__init__(connection, user_id)
 
-        base_method = self.client._execute_plan_request_with_metadata
+        execute_plan_request_base_method = (
+            self.client._execute_plan_request_with_metadata
+        )
+        execute_base_method = self.client._execute
+        execute_and_fetch_as_iterator_base_method = (
+            self.client._execute_and_fetch_as_iterator
+        )
 
-        def wrapped_method(*args, **kwargs):
-            req = base_method(*args, **kwargs)
+        def execute_plan_request_wrapped_method(*args, **kwargs):
+            req = execute_plan_request_base_method(*args, **kwargs)
             if not req.operation_id:
                 req.operation_id = str(uuid.uuid4())
                 logger.debug(
                     f"No operation_id found. Setting operation_id: {req.operation_id}"
                 )
-            self._display_operation_link(req.operation_id)
             return req
 
-        self.client._execute_plan_request_with_metadata = wrapped_method
+        self.client._execute_plan_request_with_metadata = (
+            execute_plan_request_wrapped_method
+        )
+
+        def execute_wrapped_method(client_self, req, *args, **kwargs):
+            if not self._sql_lazy_transformation(req):
+                self._display_operation_link(req.operation_id)
+            execute_base_method(req, *args, **kwargs)
+
+        self.client._execute = MethodType(execute_wrapped_method, self.client)
+
+        def execute_and_fetch_as_iterator_wrapped_method(
+            client_self, req, *args, **kwargs
+        ):
+            if not self._sql_lazy_transformation(req):
+                self._display_operation_link(req.operation_id)
+            return execute_and_fetch_as_iterator_base_method(
+                req, *args, **kwargs
+            )
+
+        self.client._execute_and_fetch_as_iterator = MethodType(
+            execute_and_fetch_as_iterator_wrapped_method, self.client
+        )
+
+    @staticmethod
+    def _sql_lazy_transformation(req):
+        # Select SQL command
+        if req.plan and req.plan.command and req.plan.command.sql_command:
+            return (
+                "select"
+                in req.plan.command.sql_command.sql.strip().lower().split()
+            )
+
+        return False
 
     def _repr_html_(self) -> str:
         if not self._active_s8s_session_id:
@@ -593,7 +635,7 @@ class DataprocSparkSession(SparkSession):
                 return
             html_element = f"""
               <div>
-                  <p><a href="{url}">Spark UI</a> (Operation: {operation_id})</p>
+                  <p><a href="{url}">Spark Query</a> (Operation: {operation_id})</p>
               </div>
               """
             display(HTML(html_element))

{dataproc_spark_connect-0.8.1 → dataproc_spark_connect-0.8.3}/setup.py

@@ -20,7 +20,7 @@ long_description = (this_directory / "README.md").read_text()
 
 setup(
     name="dataproc-spark-connect",
-    version="0.8.1",
+    version="0.8.3",
     description="Dataproc client library for Spark Connect",
     long_description=long_description,
     author="Google LLC",