Generative AI Google Brain

gemini-use-cases-retrieval-augmented-generation-rag-vector-embedding-in-bigquery

3 min read

Augment Gemini Output with Vector Embeddings from BigQuery

Source notebook

Repo path: gemini/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.ipynb · Open on GitHub · intermediate

Builds BigQuery vector-search RAG over patent abstracts and uses Gemini to generate project ideas.

Summary

The notebook teaches how to create BigQuery remote models for Vertex AI text embeddings and Gemini text generation. It generates embeddings for Singapore patent abstracts, stores them in BigQuery, creates an IVF cosine vector index, runs VECTOR_SEARCH, and feeds retrieved patent context into ML.GENERATE_TEXT. It also shows IAM setup, connection creation, index readiness checks, and cleanup.

Key code patterns

BigQuery query wrapper

client = bigquery.Client(project=PROJECT_ID)
 
def run_bq_query(sql: str):
    query_job = client.query(sql)
    result = query_job.result()
    print(f"JOB ID: {query_job.job_id} STATUS: {query_job.state}")
    return result

Centralizes SQL execution and status reporting for BigQuery jobs.

Create remote embedding model

CREATE OR REPLACE MODEL
  `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_MODEL_ID}`
REMOTE WITH CONNECTION
  `{PROJECT_ID}.{REGION}.{CONN_NAME}`
OPTIONS (ENDPOINT = '{EMBEDDINGS_ENDPOINT_TYPE}');

Maps a BigQuery ML remote model to a hosted Vertex AI embedding endpoint.

Generate and store embeddings

CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}` AS
SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(
  MODEL `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_MODEL_ID}`,
  (SELECT *, abstract AS content
   FROM `patents-public-data.google_patents_research.publications`
   WHERE LENGTH(abstract) > 0 AND LENGTH(title) > 0 AND country = 'Singapore')
)
WHERE ARRAY_LENGTH(text_embedding) > 0;

Creates reusable vector data inside BigQuery from patent abstracts.

Create vector index

CREATE OR REPLACE VECTOR INDEX my_index
ON `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}`(text_embedding)
OPTIONS(index_type = 'IVF', distance_type = 'COSINE', ivf_options = '{"num_lists":500}')

Enables more efficient approximate nearest-neighbor search over embeddings.

RAG prompt with ML.GENERATE_TEXT

SELECT ml_generate_text_llm_result AS generated, prompt
FROM ML.GENERATE_TEXT(
  MODEL `{PROJECT_ID}.{DATASET_ID}.{LLM_MODEL_ID}`,
  (SELECT CONCAT('Propose some project ideas...', STRING_AGG(FORMAT(
    "patent title: %s, patent abstract: %s", base.title, base.abstract))) AS prompt
   FROM VECTOR_SEARCH(TABLE `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}`, 'text_embedding', query, top_k => 5)),
  STRUCT(600 AS max_output_tokens, TRUE AS flatten_json_output));

Feeds vector-search context into Gemini generation directly from SQL.

Models & APIs used

  • Models: text-embedding-005, gemini-2.0-flash
  • APIs / services: Vertex AI, BigQuery, BigQuery Connection API
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-bigquery, google-cloud-bigquery-connection

When to use this

Use this pattern when BigQuery already holds text data and you want SQL-native embedding search plus Gemini-based RAG generation.

Gotchas & caveats

  • Requires roles/bigquery.connectionAdmin to create a connection.
  • Requires resourcemanager.projects.setIamPolicy to grant connection service account permissions.
  • Remaining BigQuery operations need roles/bigquery.dataEditor and roles/bigquery.user.
  • Connection service account is granted serviceUsageConsumer, bigquery.connectionUser, and aiplatform.user.
  • Workbench service account may not have sufficient permissions to add IAM policy bindings.
  • IAM updates may need propagation time; the notebook waits 60 seconds.
  • Embedding generation query might take up to 10 minutes.
  • Vector index creation might take up to 5 minutes and populates asynchronously.
  • Vector search accuracy depends on using the same embedding model for query and stored embeddings.
  • Approximate nearest-neighbor search improves performance with reduced recall.

Best practices

  • Create BigQuery remote models through a Cloud resource connection for Vertex AI access.
  • Filter out rows with empty abstract and title before generating embeddings.
  • Store generated embeddings in a BigQuery table for reuse.
  • Check INFORMATION_SCHEMA.VECTOR_INDEXES for coverage_percentage and last_refresh_time before relying on the index.
  • Use the same embedding model for indexed embeddings and query embeddings.
  • Clean up vector index, models, table, connection, dataset, and client resources after the tutorial.

Graph View

Backlinks