NOTE: This notebook has been tested in the following environment:

• Python version = 3.10

!pip install bigframes --upgrade --quiet

PROJECT_ID = ""  # @param {type:"string"}

# Set the project id
! gcloud config set project {PROJECT_ID}

ERROR: (gcloud.config.set) argument VALUE: Must be specified.
Usage: gcloud config set SECTION/PROPERTY VALUE [optional flags]
  optional flags may be  --help | --installation

For detailed information on this command and its flags, run:
  gcloud config set --help

REGION = "US"  # @param {type: "string"}

# from google.colab import auth
# auth.authenticate_user()

import bigframes.pandas as bf
from google.cloud import bigquery
from google.cloud import bigquery_connection_v1 as bq_connection

# Note: The project option is not required in all environments.
# On BigQuery Studio, the project ID is automatically detected.
bf.options.bigquery.project = PROJECT_ID

# Note: The location option is not required.
# It defaults to the location of the first table or query
# passed to read_gbq(). For APIs where a location can't be
# auto-detected, the location defaults to the "US" location.
bf.options.bigquery.location = REGION

Define the LLM model#

BigQuery DataFrames provides integration with Gemini Models via Vertex AI.

This section walks through a few steps required in order to use the model in your notebook.

from bigframes.ml.llm import GeminiTextGenerator

model = GeminiTextGenerator(model_name="gemini-2.0-flash-001")

Read data from Cloud Storage into BigQuery DataFrames#

You can create a BigQuery DataFrames DataFrame by reading data from any of the following locations:

• A local data file

• Data stored in a BigQuery table

• A data file stored in Cloud Storage

• An in-memory pandas DataFrame

In this tutorial, you create BigQuery DataFrames DataFrames by reading two CSV files stored in Cloud Storage, one containing a list of DataFrame API names and one containing a list of Series API names.

df_api = bf.read_csv("gs://cloud-samples-data/vertex-ai/bigframe/df.csv")
series_api = bf.read_csv("gs://cloud-samples-data/vertex-ai/bigframe/series.csv")

Take a peek at a few rows of data for each file:

Generate code using the LLM model#

Prepare the prompts and send them to the LLM model for prediction.

df_prompt_prefix = "Generate Pandas sample code for DataFrame."
series_prompt_prefix = "Generate Pandas sample code for Series."

df_prompt = (df_prompt_prefix + df_api['API'])
series_prompt = (series_prompt_prefix + series_api['API'])

df_prompt.head(2)

0    Generate Pandas sample code for DataFrame.values
1    Generate Pandas sample code for DataFrame.dtypes
Name: API, dtype: string

df_pred = model.predict(df_prompt.to_frame(), max_output_tokens=1024)
series_pred = model.predict(series_prompt.to_frame(), max_output_tokens=1024)

print(df_pred['ml_generate_text_llm_result'].iloc[0])

```python
import pandas as pd

# Create a DataFrame
df = pd.DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Get the values as a NumPy array
values = df.values

# Print the values
print(values)
```

Manipulate LLM output using a remote function#

The output that the LLM provides often contains additional text beyond the code sample itself. Using BigQuery DataFrames, you can deploy custom Python functions that process and transform this output.

Running the cell below creates a custom function that you can use to process the LLM output data in two ways:

1. Strip the LLM text output to include only the code block.

2. Substitute import pandas as pd with import bigframes.pandas as bf so that the resulting code block works with BigQuery DataFrames.

@bf.remote_function(cloud_function_service_account="default")
def extract_code(text: str) -> str:
  try:
    res = text[text.find('\n')+1:text.find('```', 3)]
    res = res.replace("import pandas as pd", "import bigframes.pandas as bf")
    if "import bigframes.pandas as bf" not in res:
      res = "import bigframes.pandas as bf\n" + res
    return res
  except:
    return ""

The custom function is deployed as a Cloud Function, and then integrated with BigQuery as a remote function. Save both of the function names so that you can clean them up at the end of this notebook.

CLOUD_FUNCTION_NAME = format(extract_code.bigframes_cloud_function)
print("Cloud Function Name " + CLOUD_FUNCTION_NAME)
REMOTE_FUNCTION_NAME = format(extract_code.bigframes_remote_function)
print("Remote Function Name " + REMOTE_FUNCTION_NAME)

Cloud Function Name projects/swast-scratch/locations/us-central1/functions/bigframes-6e7606963c3f06b8181b3cb9449a4363
Remote Function Name swast-scratch._63cfa399614a54153cc386c27d6c0c6fdb249f9e.bigframes_6e7606963c3f06b8181b3cb9449a4363

Apply the custom function to each LLM output DataFrame to get the processed results:

df_code = df_pred.assign(code=df_pred['ml_generate_text_llm_result'].apply(extract_code))
series_code = series_pred.assign(code=series_pred['ml_generate_text_llm_result'].apply(extract_code))

You can see the differences by inspecting the first row of data:

print(df_code['code'].iloc[0])

import bigframes.pandas as bf

# Create a DataFrame
df = pd.DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Get the values as a NumPy array
values = df.values

# Print the values
print(values)

Save the results to Cloud Storage#

BigQuery DataFrames lets you save a BigQuery DataFrames DataFrame as a CSV file in Cloud Storage for further use. Try that now with your processed LLM output data.

Create a new Cloud Storage bucket with a unique name:

import uuid
BUCKET_ID = "code-samples-" + str(uuid.uuid1())

!gcloud storage buckets create gs://{BUCKET_ID}

Creating gs://code-samples-773ee0f2-e302-11ee-8298-4201c0a8181f/...

Use to_csv to write each BigQuery DataFrames DataFrame as a CSV file in the Cloud Storage bucket:

df_code[["code"]].to_csv(f"gs://{BUCKET_ID}/df_code*.csv")
series_code[["code"]].to_csv(f"gs://{BUCKET_ID}/series_code*.csv")

You can navigate to the Cloud Storage bucket browser to download the two files and view them.

Run the following cell, and then follow the link to your Cloud Storage bucket browser:

print(f'https://console.developers.google.com/storage/browser/{BUCKET_ID}/')

https://console.developers.google.com/storage/browser/code-samples-773ee0f2-e302-11ee-8298-4201c0a8181f/

Summary and next steps#

You’ve used BigQuery DataFrames’ integration with LLM models (bigframes.ml.llm) to generate code samples, and have tranformed LLM output by creating and using a custom function in BigQuery DataFrames.

Learn more about BigQuery DataFrames in the documentation and find more sample notebooks in the GitHub repo.

# # Delete the BigQuery Connection
# from google.cloud import bigquery_connection_v1 as bq_connection
# client = bq_connection.ConnectionServiceClient()
# CONNECTION_ID = f"projects/{PROJECT_ID}/locations/{REGION}/connections/{CONN_NAME}"
# client.delete_connection(name=CONNECTION_ID)
# print(f"Deleted connection '{CONNECTION_ID}'.")

# # Delete the Cloud Function
# ! gcloud functions delete {CLOUD_FUNCTION_NAME} --quiet
# # Delete the Remote Function
# REMOTE_FUNCTION_NAME = REMOTE_FUNCTION_NAME.replace(PROJECT_ID + ".", "")
# ! bq rm --routine --force=true {REMOTE_FUNCTION_NAME}

# # Delete the Google Cloud Storage bucket and files
# ! gcloud storage rm gs://{BUCKET_ID} --recursive
# print(f"Deleted bucket '{BUCKET_ID}'.")