Generative AI Google Brain

gemini-use-cases-retrieval-augmented-generation-rag-with-bigquery

4 min read

Run RAG Pipelines in BigQuery with BQML and Vector Search

Source notebook

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

Builds a BigQuery RAG pipeline over a PDF using Document AI, embeddings, vector search, and Gemini.

Summary

This notebook teaches an end-to-end retrieval-augmented generation workflow implemented mostly in BigQuery SQL. It loads a Federal Reserve PDF from Cloud Storage into a BigQuery object table, parses it with Document AI through BigQuery ML, chunks and embeds the text, retrieves relevant chunks with VECTOR_SEARCH, and asks Gemini to generate a concise answer from the retrieved context.

Key code patterns

Cloud resource connection

!bq mk --connection --connection_type=CLOUD_RESOURCE --location=us --project_id={PROJECT_ID} "demo_conn"
!bq show --location=us --connection --project_id={PROJECT_ID} "demo_conn"

Lets BigQuery access Cloud Storage, Vertex AI, and Document AI through a connection service account.

Object table over PDF

CREATE OR REPLACE EXTERNAL TABLE `docai_demo.object_table`
WITH CONNECTION `us.demo_conn`
OPTIONS (
  uris = ['gs://{bucket_name}/scf23.pdf'],
  object_metadata = 'DIRECTORY'
);

Exposes the PDF in Cloud Storage as an object table that BigQuery can pass to ML.PROCESS_DOCUMENT.

Document AI remote model

CREATE OR REPLACE MODEL `docai_demo.layout_parser`
REMOTE WITH CONNECTION `us.demo_conn`
OPTIONS(remote_service_type="CLOUD_AI_DOCUMENT_V1", document_processor="{processor_id}")

Registers the Document AI Layout Parser processor so BigQuery ML can parse PDF content.

Chunk PDF with ML.PROCESS_DOCUMENT

CREATE OR REPLACE TABLE docai_demo.demo_result AS (
  SELECT * FROM ML.PROCESS_DOCUMENT(
    MODEL docai_demo.layout_parser,
    TABLE docai_demo.object_table,
    PROCESS_OPTIONS => (JSON '{"layout_config": {"chunking_config": {"chunk_size": 250}}}')
  )
);

Calls Document AI from BigQuery and produces chunked document output for downstream retrieval.

Generate embeddings

CREATE OR REPLACE MODEL `docai_demo.embedding_model`
REMOTE WITH CONNECTION `us.demo_conn`
OPTIONS(endpoint="text-embedding-005");
 
CREATE OR REPLACE TABLE `docai_demo.embeddings` AS
SELECT * FROM ML.GENERATE_EMBEDDING(
  MODEL `docai_demo.embedding_model`,
  TABLE `docai_demo.demo_result_parsed`
);

Creates vector representations for parsed chunks directly in BigQuery.

RAG answer with Gemini

SELECT ml_generate_text_llm_result AS generated
FROM ML.GENERATE_TEXT(
  MODEL `docai_demo.gemini_flash`,
  (SELECT CONCAT(question, STRING_AGG(FORMAT("context: %s and reference: %s", base.content, base.uri), ',\n')) AS prompt
   FROM VECTOR_SEARCH(TABLE `docai_demo.embeddings`, 'ml_generate_embedding_result', query_embedding, top_k => 10)),
  STRUCT(512 AS max_output_tokens, TRUE AS flatten_json_output)
);

Combines retrieval, prompt construction, and Gemini text generation inside BigQuery.

Models & APIs used

  • Models: text-embedding-005, gemini-2.0-flash
  • APIs / services: BigQuery, BigQuery ML, Vertex AI, Document AI, Cloud Storage
  • SDKs / libraries: google-cloud-documentai, google-cloud-bigquery

When to use this

Use this pattern when unstructured PDFs in Cloud Storage need SQL-native RAG with BigQuery, Document AI, embeddings, and Gemini.

Gotchas & caveats

  • Requires a Google Cloud project ID and project number.
  • The notebook installs google-cloud-documentai==2.31.0 and restarts the runtime.
  • Colab users must authenticate with google.colab.auth.authenticate_user().
  • The BigQuery connection service account needs documentai.viewer, storage.objectViewer, and aiplatform.user roles.
  • Resources are created in the us location, including the dataset, connection, bucket, and Document AI processor.
  • Document AI processor permissions may take a minute to propagate after IAM binding.
  • BigQuery, Vertex AI Generative AI models, Document AI, and Cloud Storage are billable.

Best practices

  • Use a Cloud resource connection so BigQuery can call Cloud Storage, Vertex AI, and Document AI.
  • Grant required IAM roles to the BigQuery connection service account before creating remote models.
  • Store source PDFs in Cloud Storage and expose them through a BigQuery object table.
  • Parse Document AI JSON into chunk content and metadata before embedding.
  • Use ML.GENERATE_EMBEDDING for both document chunks and user queries to keep embeddings consistent.
  • Retrieve top chunks with VECTOR_SEARCH before calling ML.GENERATE_TEXT.
  • Include source URI references in the context passed to Gemini.
  • Clean up BigQuery assets, the connection, Cloud Storage bucket, and Document AI processor after the tutorial.

Graph View

Backlinks