Generative AI Google Brain

gemini-use-cases-retrieval-augmented-generation-scalable-rag-with-bigframes

4 min read

Production & Scalable RAG Pipeline Using BigFrames

Source notebook

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

Builds a scalable BigFrames RAG pipeline over Stack Overflow data with BigQuery, Vertex AI, and LangChain.

Summary

The notebook teaches how to load Stack Overflow Q&A data from BigQuery into BigFrames, clean HTML into Markdown, chunk text, and generate embeddings with Vertex AI. It saves embeddings into partitioned BigQuery tables with incremental append and deduplication, then uses BigQueryVectorStore and LangChain retrieval chains to answer questions with Gemini. It also shows how to move custom chunking into a BigFrames remote function for larger datasets.

Key code patterns

Incremental BigQuery load

query = f"""
SELECT creation_date, last_edit_date, question_id, question_title,
       question_body AS question_text, answers
FROM `production-ai-template.stackoverflow_qa.stackoverflow_python_questions_and_answers`
WHERE TRUE
  {date_filter if IS_INCREMENTAL else ""}
"""
df = bpd.read_gbq(query)

Loads source data into BigFrames and gates scheduled runs by a date window.

HTML to Markdown document assembly

df["question_title_md"] = "# " + df["question_title"] + "\n"
df["question_text_md"] = df["question_text"].to_pandas().apply(convert_html_to_markdown) + "\n"
df["answers_md"] = df["answers"].to_pandas().apply(create_answers_markdown)
df["full_text_md"] = df["question_title_md"] + df["question_text_md"] + df["answers_md"]

Creates a readable question-plus-answers document for downstream chunking and retrieval.

Chunk and explode text

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1500,
    chunk_overlap=20,
    length_function=len,
)
df["text_chunk"] = df["full_text_md"].to_pandas().astype(object).swifter.apply(text_splitter.split_text)
df = df.explode("text_chunk").reset_index(drop=True)
df["chunk_id"] = df["question_id"].astype("string") + "__" + chunk_ids

Turns each Stack Overflow document into retrievable chunks with stable chunk IDs.

BigFrames Vertex AI embeddings

embedder = llm.TextEmbeddingGenerator(model_name="text-embedding-005")
embeddings_df = embedder.predict(df["text_chunk"])
df = df.assign(
    embedding_result=embeddings_df["ml_generate_embedding_result"],
    embedding_statistics=embeddings_df["ml_generate_embedding_statistics"],
    embedding_status=embeddings_df["ml_generate_embedding_status"],
)

Uses BigFrames integration with Vertex AI batch scoring to embed chunks in the data pipeline.

Partitioned BigQuery write

table = bigquery.Table(f"{project_id}.{dataset_id}.{table_id}", schema=table_schema)
table.time_partitioning = bigquery.TimePartitioning(
    type_=bigquery.TimePartitioningType.DAY,
    field=partition_column,
)
bq_client.create_dataset(dataset, exists_ok=True)
bq_client.create_table(table=table, exists_ok=True)
df.to_gbq(destination_table=f"{dataset_id}.{table_id}", if_exists=if_exists_mode)

Stores embeddings in BigQuery with daily partitioning and append-or-replace behavior.

LangChain RAG over BigQuery

embedding_model = VertexAIEmbeddings(model_name="text-embedding-005", project=PROJECT_ID)
bq_store = BigQueryVectorStore(
    project_id=PROJECT_ID, location="US", dataset_name=DESTINATION_DATASET_ID,
    table_name=DESTINATION_DEDUPED_QUESTIONS_TABLE_ID,
    embedding=embedding_model, embedding_field="embedding_result", content_field="text_chunk",
)
llm = ChatVertexAI(model_name="gemini-2.0-flash")
rag_chain = create_retrieval_chain(bq_store.as_retriever(), combine_docs_chain)

Connects BigQuery retrieval to a Gemini-backed LangChain answer generation chain.

Remote function chunking

@bigframes.pandas.remote_function(packages=["langchain"], reuse=True)
def chunk_text_udf(text: str) -> str:
    return json.dumps([
        chunk.page_content for chunk in text_splitter.create_documents([text])
    ])
 
df_udf["full_text_chunk"] = df_udf["full_text_md"].apply(chunk_text_udf)
df_udf["full_text_chunk"] = bbq.json_extract_string_array(df_udf["full_text_chunk"])

Shows how custom Python processing can run remotely instead of loading large data into memory.

Models & APIs used

  • Models: text-embedding-005, gemini-2.0-flash
  • APIs / services: Vertex AI, BigQuery, BigQuery DataFrames, Cloud Functions, BigQuery Connection API
  • SDKs / libraries: google-cloud-aiplatform, vertexai, bigframes, google-cloud-bigquery, google-cloud-bigquery-connection, langchain, langchain-google-community, langchain-google-vertexai, markdownify, swifter

When to use this

Use this pattern for scheduled, scalable RAG pipelines over BigQuery data that need Vertex AI embeddings, BigQuery storage, and LangChain retrieval.

Gotchas & caveats

  • Colab users must authenticate with google.colab.auth.authenticate_user().
  • The notebook installs packages and explicitly restarts the runtime before continuing.
  • Vertex AI API must be enabled for the selected Google Cloud project.
  • GOOGLE_CLOUD_REGION must be in a US region because the source dataset is in US.
  • BigQuery client and BigFrames are configured for location US while Vertex AI uses LOCATION.
  • Calling to_pandas() loads data into memory, so large datasets should use remote functions for custom processing.
  • Embedding generation may take a few minutes.
  • Incremental appends can create duplicate questions and require deduplication by latest creation_timestamp.
  • BigFrames remote functions expect simple input and output types, so the chunk list is returned as a JSON string and parsed back with json_extract_string_array.
  • Cleanup deletes the BigQuery dataset, the generated Cloud Function, and the BigQuery external connection.

Best practices

  • Use BigFrames to process BigQuery data with pandas-like syntax without moving terabyte-scale data out of BigQuery.
  • Parameterize scheduled runs with RUN_DATE, IS_INCREMENTAL, LOOK_BACK_DAYS, START_DATE, and END_DATE.
  • Sort by last_edit_date and drop duplicate question_id values before embedding.
  • Convert HTML questions and answers to Markdown before chunking and generation.
  • Use chunk overlap and consider preserving paragraphs, sections, markdown hierarchy, code blocks, lists, and question-answer grouping in production.
  • Partition the embeddings table by creation_timestamp.
  • Append during incremental runs and replace during full refreshes.
  • Deduplicate incremental results by question_id while keeping the latest creation_timestamp.
  • Use BigFrames remote functions for custom Python transformations that are too large for local memory.
  • Clean up created BigQuery, Cloud Function, and connection resources after the notebook.

Graph View

Backlinks