dataproc-spark-connect 0.7.2__py2.py3-none-any.whl → 0.7.4__py2.py3-none-any.whl

{dataproc_spark_connect-0.7.2.dist-info → dataproc_spark_connect-0.7.4.dist-info}/METADATA

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

{dataproc_spark_connect-0.7.2.dist-info → dataproc_spark_connect-0.7.4.dist-info}/RECORD

@@ -1,12 +1,12 @@
 google/cloud/dataproc_spark_connect/__init__.py,sha256=dIqHNWVWWrSuRf26x11kX5e9yMKSHCtmI_GBj1-FDdE,1101
 google/cloud/dataproc_spark_connect/exceptions.py,sha256=ilGyHD5M_yBQ3IC58-Y5miRGIQVJsLaNKvEGcHuk_BE,969
 google/cloud/dataproc_spark_connect/pypi_artifacts.py,sha256=gd-VMwiVP-EJuPp9Vf9Shx8pqps3oSKp0hBcSSZQS-A,1575
-google/cloud/dataproc_spark_connect/session.py,sha256=7CgRmYVp8M4oVosvxz08jp72BU4ZtAKvxjtd0101GHs,24952
+google/cloud/dataproc_spark_connect/session.py,sha256=kMCZWmi_-ScJy9NO7NFrHaHDTXKxMwaCSDbdqGxEngk,26390
 google/cloud/dataproc_spark_connect/client/__init__.py,sha256=6hCNSsgYlie6GuVpc5gjFsPnyeMTScTpXSPYqp1fplY,615
 google/cloud/dataproc_spark_connect/client/core.py,sha256=m3oXTKBm3sBy6jhDu9GRecrxLb5CdEM53SgMlnJb6ag,4616
 google/cloud/dataproc_spark_connect/client/proxy.py,sha256=qUZXvVY1yn934vE6nlO495XUZ53AUx9O74a9ozkGI9U,8976
-dataproc_spark_connect-0.7.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-dataproc_spark_connect-0.7.2.dist-info/METADATA,sha256=EtPp009uOCvR1I3tiT-x1FJ4bWpPHhEqZxvJO4rfwKk,3328
-dataproc_spark_connect-0.7.2.dist-info/WHEEL,sha256=OpXWERl2xLPRHTvd2ZXo_iluPEQd8uSbYkJ53NAER_Y,109
-dataproc_spark_connect-0.7.2.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
-dataproc_spark_connect-0.7.2.dist-info/RECORD,,
+dataproc_spark_connect-0.7.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+dataproc_spark_connect-0.7.4.dist-info/METADATA,sha256=viQwCWio0b-xja72qtR447f9Ol7nl0k5d6bx1j2BAEk,3328
+dataproc_spark_connect-0.7.4.dist-info/WHEEL,sha256=OpXWERl2xLPRHTvd2ZXo_iluPEQd8uSbYkJ53NAER_Y,109
+dataproc_spark_connect-0.7.4.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
+dataproc_spark_connect-0.7.4.dist-info/RECORD,,

google/cloud/dataproc_spark_connect/session.py

@@ -67,7 +67,7 @@ class DataprocSparkSession(SparkSession):
     ... ) # doctest: +SKIP
     """
 
-    _DEFAULT_RUNTIME_VERSION = "2.2"
+    _DEFAULT_RUNTIME_VERSION = "2.3"
 
     _active_s8s_session_uuid: ClassVar[Optional[str]] = None
     _project_id = None
@@ -198,7 +198,7 @@ class DataprocSparkSession(SparkSession):
                 DataprocSparkSession._active_s8s_session_id = session_id
                 s8s_creation_start_time = time.time()
 
-                stop_create_session_pbar = False
+                stop_create_session_pbar_event = threading.Event()
 
                 def create_session_pbar():
                     iterations = 150
@@ -208,12 +208,12 @@ class DataprocSparkSession(SparkSession):
                         ncols=80,
                     )
                     for i in pbar:
-                        if stop_create_session_pbar:
+                        if stop_create_session_pbar_event.is_set():
                             break
                         # Last iteration
                         if i >= iterations - 1:
                             # Sleep until session created
-                            while not stop_create_session_pbar:
+                            while not stop_create_session_pbar_event.is_set():
                                 time.sleep(1)
                         else:
                             time.sleep(1)
@@ -258,7 +258,7 @@ class DataprocSparkSession(SparkSession):
                             timeout=600,  # seconds
                         )
                     )
-                    stop_create_session_pbar = True
+                    stop_create_session_pbar_event.set()
                     create_session_pbar_thread.join()
                     print("Dataproc Session was successfully created")
                     file_path = (
@@ -280,7 +280,7 @@ class DataprocSparkSession(SparkSession):
                                 f"Exception while writing active session to file {file_path}, {e}"
                             )
                 except (InvalidArgument, PermissionDenied) as e:
-                    stop_create_session_pbar = True
+                    stop_create_session_pbar_event.set()
                     if create_session_pbar_thread.is_alive():
                         create_session_pbar_thread.join()
                     DataprocSparkSession._active_s8s_session_id = None
@@ -288,7 +288,7 @@ class DataprocSparkSession(SparkSession):
                         f"Error while creating Dataproc Session: {e.message}"
                     )
                 except Exception as e:
-                    stop_create_session_pbar = True
+                    stop_create_session_pbar_event.set()
                     if create_session_pbar_thread.is_alive():
                         create_session_pbar_thread.join()
                     DataprocSparkSession._active_s8s_session_id = None
@@ -296,7 +296,7 @@ class DataprocSparkSession(SparkSession):
                         f"Error while creating Dataproc Session"
                     ) from e
                 finally:
-                    stop_create_session_pbar = True
+                    stop_create_session_pbar_event.set()
 
                 logger.debug(
                     f"Dataproc Session created: {session_id} in {int(time.time() - s8s_creation_start_time)} seconds"
@@ -408,6 +408,30 @@ class DataprocSparkSession(SparkSession):
                 dataproc_config.labels["colab-notebook-kernel-id"] = os.environ[
                     "COLAB_NOTEBOOK_KERNEL_ID"
                 ]
+            default_datasource = os.getenv(
+                "DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE"
+            )
+            if (
+                default_datasource
+                and dataproc_config.runtime_config.version == "2.3"
+            ):
+                if default_datasource == "bigquery":
+                    bq_datasource_properties = {
+                        "spark.datasource.bigquery.viewsEnabled": "true",
+                        "spark.datasource.bigquery.writeMethod": "direct",
+                        "spark.sql.catalog.spark_catalog": "com.google.cloud.spark.bigquery.BigQuerySparkSessionCatalog",
+                        "spark.sql.legacy.createHiveTableByDefault": "false",
+                        "spark.sql.sources.default": "bigquery",
+                    }
+                    # Merge default configs with existing properties, user configs take precedence
+                    for k, v in bq_datasource_properties.items():
+                        if k not in dataproc_config.runtime_config.properties:
+                            dataproc_config.runtime_config.properties[k] = v
+                else:
+                    logger.warning(
+                        f"DATAPROC_SPARK_CONNECT_DEFAULT_DATASOURCE is set to an invalid value:"
+                        f" {default_datasource}. Supported value is 'bigquery'."
+                    )
             return dataproc_config
 
         @staticmethod