Preparing Data for ML, Embeddings and RAG

What it is

Machine learning models are only as good as the data that feeds them. A data engineer prepares features for training and serving, often with BigQuery ML so models are built where the data lives. For generative AI, the engineer turns unstructured documents into chunks and embeddings so a model can retrieve relevant facts: retrieval-augmented generation (RAG).

How it works

Preparing data for feature engineering

Data leakage happens when a feature contains information that would not exist at prediction time, such as a "cancellation date" used to predict churn. Leaky models look excellent in evaluation and fail in production.

Training and serving with BigQuery ML

The TRANSFORM clause stores preprocessing with the model. ML.PREDICT then applies the same steps automatically, which prevents training-serving skew.

CREATE OR REPLACE MODEL `ml.churn_model`
TRANSFORM (
  ML.STANDARD_SCALER(tenure_months) OVER () AS tenure_scaled,
  ML.QUANTILE_BUCKETIZE(monthly_spend, 4) OVER () AS spend_bucket,
  plan_type,
  churned
)
OPTIONS (model_type = 'LOGISTIC_REG', input_label_cols = ['churned']) AS
SELECT tenure_months, monthly_spend, plan_type, churned
FROM ml.customer_features
WHERE snapshot_date < '2026-07-01';

SELECT * FROM ML.EVALUATE(MODEL `ml.churn_model`);

SELECT customer_id, predicted_churned
FROM ML.PREDICT(MODEL `ml.churn_model`,
  (SELECT * FROM ml.customer_features WHERE snapshot_date = '2026-10-01'));

Preparing unstructured data for embeddings

CREATE MODEL `rag.embedder`
REMOTE WITH CONNECTION `prod-ai.us.ai-conn`
OPTIONS (ENDPOINT = 'text-embedding-005');

CREATE TABLE rag.chunk_vectors AS
SELECT * FROM ML.GENERATE_EMBEDDING(
  MODEL `rag.embedder`,
  (SELECT chunk_id, chunk_text AS content FROM rag.chunks),
  STRUCT(TRUE AS flatten_json_output));

SELECT base.chunk_id, base.content, distance
FROM VECTOR_SEARCH(
  TABLE rag.chunk_vectors, 'ml_generate_embedding_result',
  (SELECT ml_generate_embedding_result
   FROM ML.GENERATE_EMBEDDING(MODEL `rag.embedder`,
     (SELECT 'How do I reset the router?' AS content))),
  top_k => 5, distance_type => 'COSINE');

Use the same embedding model for documents and questions; vectors from different models are not comparable. Re-embed chunks when source documents change, and keep metadata such as product and date to filter results.

A worked example

A telecom company wants a support assistant that answers from 20,000 PDF manuals and cites the source. The engineer creates an object table over the PDFs and extracts text. Text is split into overlapping chunks tagged with manual ID and section, and embeddings go into a BigQuery table with a vector index.

At question time, the question is embedded, VECTOR_SEARCH returns the closest chunks, and the LLM answers using only those chunks with their IDs as citations.

How the exam asks about this

• "Preprocessing in training differs from what the application sends at prediction time." Put preprocessing in the BigQuery ML TRANSFORM clause so ML.PREDICT applies it automatically.
• "A churn model scores 99% in evaluation but fails in production." Check for data leakage from features unknown at prediction time.
• "An LLM answers questions about internal documents incorrectly." Implement RAG: chunk and embed the documents, retrieve with vector search, ground the prompt.