Spark Up Your MongoDB and BigQuery Using BigQuery Spark Stored Procedures

1
from pyspark.sql import SparkSession
2
from google.cloud import secretmanager
3
4
def access_secret_version(secret_id, project_id):
5
    client = secretmanager.SecretManagerServiceClient()
6
    name = f"projects/{project_id}/secrets/{secret_id}/versions/1"
7
    response = client.access_secret_version(request={"name": name})
8
    payload = response.payload.data.decode("UTF-8")
9
    return payload
10
# Update project_number, username_secret_id and password_secret_id, comment them out if you did not create the secrets earlier 
11
12
project_id = "<Your project number, 12 digit number>"
13
username_secret_id = "<Your username secret id>"
14
password_secret_id = "<Your password secret id>"
15
16
username = access_secret_version(username_secret_id, project_id)
17
password = access_secret_version(password_secret_id, project_id)
18
19
  # Update the mongodb_uri directly if with your username and password if you did not create a secret from Step 7, update the hostname with your hostname
20
mongodb_uri = "mongodb+srv://"+username+":"+password+"@<hostname>/sample_mflix.comments"
21
22
my_spark = SparkSession \
23
  .builder \
24
  .appName("myApp") \
25
  .config("spark.mongodb.read.connection.uri", mongodb_uri) \
26
  .config("spark.mongodb.write.connection.uri", mongodb_uri) \
27
  .getOrCreate()
28
29
30
dataFrame = my_spark.read.format("mongodb").option("database", "sample_mflix").option("collection", "comments").load()
31
32
dataFrame.show()
33
34
# Saving the data to BigQuery
35
dataFrame.write.format("bigquery") \
36
  .option("writeMethod", "direct") \
37
    .save("spark_run.sample_mflix_comments")

1
!bq mk \
2
  --connection \
3
  --location=US \
4
  --project_id=<GCP Project id> \
5
  --connection_type=CLOUD_RESOURCE gentext-conn 

1
CREATE OR REPLACE MODEL `gcp-pov.spark_run.llm_model`
2
REMOTE WITH CONNECTION `us.gentext-conn`
3
OPTIONS (ENDPOINT = 'gemini-pro');

1
SELECT prompt,ml_generate_text_result
2
FROM
3
ML.GENERATE_TEXT( MODEL `gcp-pov.spark_run.llm_model`,
4
  (
5
    SELECT CONCAT('Extract the host name from the email: ', email) AS prompt,
6
      *  FROM  `gcp-pov.spark_run.sample_mflix_comments`
7
    LIMIT 5),
8
  STRUCT(
9
    0.9 AS temperature,
10
    100 AS max_output_tokens
11
    )
12
  );