Run PySpark code in BigQuery Studio notebooks

This document shows you how to run PySpark code in a BigQuery Python notebook.

Before you begin

If you haven't already done so, create a Google Cloud project and a Cloud Storage bucket.

  1. Set up your project

    1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
    2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

      Roles required to select or create a project

      • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
      • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

      Go to project selector

    3. Enable the Managed Service for Apache Spark, BigQuery, and Cloud Storage APIs.

      Roles required to enable APIs

      To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

      Enable the APIs

    4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

      Roles required to select or create a project

      • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
      • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

      Go to project selector

    5. Enable the Managed Service for Apache Spark, BigQuery, and Cloud Storage APIs.

      Roles required to enable APIs

      To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

      Enable the APIs

  2. Create a Cloud Storage bucket in your project if you don't have one you can use.

  3. Set up your notebook

Pricing

For pricing information, see BigQuery Notebook runtime pricing.

Open a BigQuery Studio Python notebook

  1. In the Google Cloud console, go to the BigQuery page.

    Go to BigQuery

  2. In the tab bar of the details pane, click the arrow next to the + sign, and then click Notebook.

Create a Spark session in a BigQuery Studio notebook

You can use a BigQuery Studio Python notebook to create a Spark Connect interactive session. Each BigQuery Studio notebook can have only one active Spark session associated with it.

You can create a Spark session in a BigQuery Studio Python notebook in the following ways:

  • Configure and create a single session in the notebook.
  • Configure a Spark session in an interactive session template, then use the template to configure and create a session in the notebook. BigQuery provides a Query using Spark feature that helps you start coding the templated session as explained under the Templated Spark session tab.

Single session

To create a Spark session in a new notebook, do the following:

  1. In the tab bar of the editor pane, click the drop-down arrow next to the + sign, and then click Notebook.

    Screenshot showing the BigQuery interface with the '+' button for creating a new notebook.
  2. Copy and run the following code in a notebook cell to configure and create a basic Spark session.

from google.cloud.dataproc_spark_connect import DataprocSparkSession
from google.cloud.dataproc_v1 import Session

import pyspark.sql.functions as f

session = Session()

# Create the Spark session.
spark = (
   DataprocSparkSession.builder
     .appName("APP_NAME")
     .dataprocSessionConfig(session)
     .getOrCreate()
)

Replace the following:

  • APP_NAME: An optional name for your session.
  • Optional Session settings: You can add Managed Service for Apache Spark API Session settings to customize your session. Here are some examples:
    • RuntimeConfig:
      Code help showing session.runtime.config options.
      • session.runtime_config.properties={spark.property.key1:VALUE_1,...,spark.property.keyN:VALUE_N}
      • session.runtime_config.container_image = path/to/container/image
    • EnvironmentConfig:
      Code help showing session-environment-config-execution-config options.
      • session.environment_config.execution_config.subnetwork_uri = "SUBNET_NAME"
      • session.environment_config.execution_config.ttl = {"seconds": VALUE}
      • session.environment_config.execution_config.service_account = SERVICE_ACCOUNT

Templated Spark session

You can enter and run the code in a notebook cell to create a Spark session based on an existing session template. Any session configuration settings you provide in your notebook code will override any of the same settings that are set in the session template.

To get started quickly, use the Query using Spark template to pre-populate your notebook with Spark session template code:

  1. In the tab bar of the editor pane, click the drop-down arrow next to the + sign, and then click Notebook.
    Screenshot showing the BigQuery interface with the '+' button for creating a new notebook.
  2. Under Start with a template, click Query using Spark, then click Use template to insert the code in your notebook.
    BigQuery UI selections to start with a template
  3. Specify the variables as explained in the Notes.
  4. You can delete any additional sample code cells inserted in the notebook.
from google.cloud.dataproc_spark_connect import DataprocSparkSession
from google.cloud.dataproc_v1 import Session
session = Session()
project_id = "PROJECT_ID"
location = "LOCATION"
# Configure the session with an existing session template.
session_template = "SESSION_TEMPLATE"
session.session_template = f"projects/{project_id}/locations/{location}/sessionTemplates/{session_template}"
# Create the Spark session.
spark = (
   DataprocSparkSession.builder
     .appName("APP_NAME")
     .dataprocSessionConfig(session)
     .getOrCreate()
)

Replace the following:

  • PROJECT_ID: Your project ID, which is listed in the Project info section of the Google Cloud console dashboard.
  • LOCATION: The Compute Engine region where your notebook session will run. If not supplied, the region of the VM that creates the notebook will be used.
  • SESSION_TEMPLATE: The name of an existing interactive session template. Session configuration settings are obtained from the template. The template must also specify the following settings:

    • Runtime version 2.3+
    • Notebook type: Spark Connect

      Example:

      Screenshot showing the Spark Connect required settings.
  • APP_NAME: An optional name for your session.

Write and run PySpark code in your BigQuery Studio notebook

After you create a Spark session in your notebook, use the session to run Spark notebook code in the notebook.

Spark Connect PySpark API support: Your Spark Connect notebook session supports most PySpark APIs, including DataFrame, Functions, and Column, but does not support SparkContext and RDD and other PySpark APIs. For more information, see What is supported in Spark 3.5.

Spark Connect notebook direct writes: Spark sessions in a BigQuery Studio notebook pre-configure the Spark BigQuery connector to make DIRECT data writes. The DIRECT write method uses the BigQuery Storage Write API, which writes data directly into BigQuery; the INDIRECT write method, which is the default for Managed Service for Apache Spark batches, writes data to an intermediate Cloud Storage bucket, then writes the data to BigQuery (for more information on INDIRECT writes, see Read and write data from and to BigQuery).

Managed Service for Apache Spark specific APIs: Managed Service for Apache Spark simplifies adding PyPI packages dynamically to your Spark session by extending the addArtifacts method. You can specify the list in version-scheme format, (similar to pip install). This instructs the Spark Connect server to install packages and their dependencies on all cluster nodes, making them available to workers for your UDFs.

Example that installs specified textdistance version and latest compatible random2 libraries on the cluster to allow UDFs using textdistance and random2 to run on worker nodes.

spark.addArtifacts("textdistance==4.6.1", "random2", pypi=True)

Notebook code help: The BigQuery Studio notebook provides code help when you hold the pointer over a class or method name, and provides code completion help as you input code.

In the following example, entering DataprocSparkSession and holding the pointer over this class name displays code completion and documentation help.

Code documentation and code completion tip examples.

BigQuery Studio notebook PySpark examples

This section provides BigQuery Studio Python notebook examples with PySpark code to perform the following tasks:

  • Run a wordcount against a public Shakespeare dataset.
  • Create an Iceberg table with metadata saved in Lakehouse runtime catalog.

Wordcount

The following PySpark example creates a Spark session, then counts word occurrences in a public bigquery-public-data.samples.shakespeare dataset.

# Basic wordcount example
from google.cloud.dataproc_spark_connect import DataprocSparkSession
from google.cloud.dataproc_v1 import Session
import pyspark.sql.functions as f
session = Session()

# Create the Spark session.
spark = (
   DataprocSparkSession.builder
     .appName("APP_NAME")
     .dataprocSessionConfig(session)
     .getOrCreate()
)
# Run a wordcount on the public Shakespeare dataset.
df = spark.read.format("bigquery").option("table", "bigquery-public-data.samples.shakespeare").load()
words_df = df.select(f.explode(f.split(f.col("word"), " ")).alias("word"))
word_counts_df = words_df.filter(f.col("word") != "").groupBy("word").agg(f.count("*").alias("count")).orderBy("word")
word_counts_df.show()

Replace the following:

  • APP_NAME: An optional name for your session.

Output:

The cell output lists a sample of the wordcount output. To see session details in the Google Cloud console, click the Interactive Session Detail View link. To monitor your Spark session, click View Spark UI on the session details page.

View Spark UI button in session details page in console
Interactive Session Detail View: LINK
+------------+-----+
|        word|count|
+------------+-----+
|           '|   42|
|       ''All|    1|
|     ''Among|    1|
|       ''And|    1|
|       ''But|    1|
|    ''Gamut'|    1|
|       ''How|    1|
|        ''Lo|    1|
|      ''Look|    1|
|        ''My|    1|
|       ''Now|    1|
|         ''O|    1|
|      ''Od's|    1|
|       ''The|    1|
|       ''Tis|    4|
|      ''When|    1|
|       ''tis|    1|
|      ''twas|    1|
|          'A|   10|
|'ARTEMIDORUS|    1|
+------------+-----+
only showing top 20 rows

Iceberg table

Run PySpark code to create an Iceberg table with Lakehouse runtime catalog metadata

The following example code creates a sample_iceberg_table with table metadata stored in Lakehouse runtime catalog, and then queries the table.

from google.cloud.dataproc_spark_connect import DataprocSparkSession
from google.cloud.dataproc_v1 import Session
# Create the Dataproc Serverless session.
session = Session()
# Set the session configuration for BigLake Metastore with the Iceberg environment.
project_id = "PROJECT_ID"
region = "REGION"
subnet_name = "SUBNET_NAME"
location = "LOCATION"
session.environment_config.execution_config.subnetwork_uri = f"{subnet_name}"
warehouse_dir = "gs://BUCKET/WAREHOUSE_DIRECTORY"
catalog = "CATALOG"
namespace = "NAMESPACE"
session.runtime_config.properties[f"spark.sql.catalog.{catalog}"] = "org.apache.iceberg.spark.SparkCatalog"
session.runtime_config.properties[f"spark.sql.catalog.{catalog}.catalog-impl"] = "org.apache.iceberg.gcp.bigquery.BigQueryMetastoreCatalog"
session.runtime_config.properties[f"spark.sql.catalog.{catalog}.gcp_project"] = f"{project_id}"
session.runtime_config.properties[f"spark.sql.catalog.{catalog}.gcp_location"] = f"{location}"
session.runtime_config.properties[f"spark.sql.catalog.{catalog}.warehouse"] = f"{warehouse_dir}"
# Create the Spark Connect session.
spark = (
   DataprocSparkSession.builder
     .appName("APP_NAME")
     .dataprocSessionConfig(session)
     .getOrCreate()
)
# Create the namespace in BigQuery.
spark.sql(f"USE `{catalog}`;")
spark.sql(f"CREATE NAMESPACE IF NOT EXISTS `{namespace}`;")
spark.sql(f"USE `{namespace}`;")
# Create the Iceberg table.
spark.sql("DROP TABLE IF EXISTS `sample_iceberg_table`");
spark.sql("CREATE TABLE sample_iceberg_table (id int, data string) USING ICEBERG;")
spark.sql("DESCRIBE sample_iceberg_table;")
# Insert table data and query the table.
spark.sql("INSERT INTO sample_iceberg_table VALUES (1, \"first row\");")
# Alter table, then query and display table data and schema.
spark.