Optimizing BigQuery and Gemini API for Real-Time Enterprise AI Pipelines

Formare AI designs and optimizes enterprise data pipelines that integrate directly with Gemini API and Vertex AI. By applying partition-aware query planning, materialized view pre-computation, and Gemini context caching, these pipelines reduce token costs by up to 70% and cut P95 query latency below 1 second for datasets exceeding 10 billion rows.

The most common cost explosion in production AI pipelines: an agent that calls a BigQuery tool on every conversation turn, with no awareness of partition pruning. At 50,000 tokens per query result and 1,000 daily active users, that is €2,400/month in Gemini API costs alone — before the BigQuery compute bill.

The Six Optimizations

Partition-aware query generation — agents only scan the date ranges they need
Materialized views for agent "hot paths" — pre-compute common aggregations
Gemini context caching — reuse large system prompts across 1,000s of requests
Batched embedding generation — vector embeddings in parallel, not serial
Result set compression — compress BigQuery results before passing to Gemini
Adaptive token budgeting — limit result rows dynamically based on available context window

Optimization 1: Partition-Aware Query Generation

BigQuery tables at enterprise scale are almost always partitioned by date or ingestion time. An agent that generates SQL without explicit partition filters will perform a full table scan — 10 billion rows instead of 50 million. Inject partition awareness into the agent's system prompt with the exact column name and valid date range:

python

1from datetime import datetime, timedelta
2from google.cloud import bigquery
3
4def build_partition_context(table_id: str) -> str:
5    """
6    Introspects the table to find its partition column and recent date range.
7    Injects this directly into the agent system prompt.
8    """
9    client = bigquery.Client()
10    table = client.get_table(table_id)
11
12    partition_col = None
13    if table.time_partitioning:
14        partition_col = table.time_partitioning.field or "_PARTITIONTIME"
15    elif table.range_partitioning:
16        partition_col = table.range_partitioning.field
17
18    if not partition_col:
19        return ""  # Unpartitioned table — no hint to inject
20
21    today = datetime.utcnow().date()
22    ninety_days_ago = today - timedelta(days=90)
23
24    return f"""
25CRITICAL PERFORMANCE RULE:
26Table `{table_id}` is partitioned on `{partition_col}`.
27ALWAYS include a filter on `{partition_col}` to avoid full table scans.
28Default date range unless user specifies otherwise: {ninety_days_ago} to {today}.
29Example: WHERE DATE({partition_col}) BETWEEN '{ninety_days_ago}' AND '{today}'
30A query without this filter will scan the ENTIRE table history and will be rejected.
31"""
32
33# Inject into agent system prompt at initialization
34SYSTEM_PROMPT = base_prompt + build_partition_context("project.dataset.sales")

Optimization 2: Materialized Views for Agent Hot Paths

Identify the 5–10 query patterns your agents run most frequently and pre-compute them as materialized views. BigQuery refreshes these automatically and serves results from cache when the underlying data hasn't changed. For our demand forecasting agents, the daily sales-by-SKU-by-store aggregation was being recomputed on every agent call — moving it to a materialized view cut that specific query's cost to near-zero:

sql

1-- Create materialized view for the agent's most common aggregation
2-- BigQuery auto-refreshes within 5 minutes of the source table updating
3CREATE MATERIALIZED VIEW `project.analytics.mv_daily_sales_by_sku`
4OPTIONS (
5  enable_refresh = TRUE,
6  refresh_interval_minutes = 60,
7  max_staleness = INTERVAL "4" HOUR
8)
9AS
10SELECT
11  DATE(sale_timestamp)          AS sale_date,
12  store_id,
13  sku_id,
14  product_category,
15  SUM(quantity_sold)            AS units_sold,
16  SUM(net_revenue_eur)          AS revenue_eur,
17  AVG(unit_price_eur)           AS avg_price_eur,
18  COUNT(DISTINCT transaction_id) AS transaction_count
19FROM `project.raw.pos_transactions`
20WHERE sale_timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 180 DAY)
21GROUP BY 1, 2, 3, 4;
22
23-- Agents now query this view instead of the raw table
24-- 10B row scan → 45M row pre-aggregated view
25-- P95 latency: 4.2s → 380ms

Optimization 3: Gemini Context Caching

Gemini's context caching feature (available for Gemini 1.5 Pro and 2.x models) allows you to cache a large system prompt or document corpus and reference it by ID instead of re-sending it on every request. For agents that carry a large schema context or RAG document set in every prompt, this is the single highest-impact cost optimization:

python

1import vertexai
2from vertexai.generative_models import GenerativeModel, Content, Part
3from vertexai.preview import caching
4import datetime
5
6# Cache the large schema context once per day
7# (Minimum cache TTL: 1 hour. Billed at reduced rate.)
8def create_schema_cache(schema_text: str) -> str:
9    """
10    Creates a cached context for the BigQuery schema.
11    Returns the cache name to use in subsequent requests.
12    """
13    cached_content = caching.CachedContent.create(
14        model_name="gemini-2.0-flash-001",
15        system_instruction=f"""You are a BigQuery SQL specialist.
16You have access to the following schema. Only reference these exact columns.
17
18{schema_text}
19
20Rules: Always partition-filter, always qualify table names, always LIMIT 1000.""",
21        contents=[],  # Schema is in system_instruction; no additional content
22        ttl=datetime.timedelta(hours=24),
23        display_name="bigquery-schema-cache",
24    )
25    return cached_content.name
26
27# Use the cached context for every agent request
28CACHE_NAME = create_schema_cache(SCHEMA_TEXT)  # Called once at startup
29
30model = GenerativeModel.from_cached_content(
31    cached_content=caching.CachedContent(CACHE_NAME)
32)
33
34# Subsequent requests reference the cache — no re-tokenization of schema
35# Cost reduction: 40,000 cached input tokens → ~€0.0002/request (vs €0.04 uncached)

Reduction in Gemini API token costs after context caching: 70%
P95 query latency after materialized view optimization (was 4.2s): 800ms
Rows in source tables — partition filtering is non-negotiable: 10B+
Monthly API cost avoided per 1,000 daily active agent users: €3,000+

Optimization 4: Batched Embedding Generation

If your pipeline generates vector embeddings for RAG — product descriptions, document chunks, user query embeddings — never generate them serially in a loop. The Vertex AI text-embedding API accepts batches of up to 250 texts per request. Batching reduces API call overhead by 50–100× and is the difference between an overnight embedding job and a 10-minute one:

python

1from vertexai.language_models import TextEmbeddingModel
2from concurrent.futures import ThreadPoolExecutor
3import math
4
5model = TextEmbeddingModel.from_pretrained("text-embedding-005")
6
7def embed_in_batches(texts: list[str], batch_size: int = 250) -> list[list[float]]:
8    """
9    Generates embeddings in parallel batches of 250.
10    For 10,000 texts: ~40 API calls (parallel) instead of 10,000 serial calls.
11    """
12    batches = [texts[i:i+batch_size] for i in range(0, len(texts), batch_size)]
13
14    def embed_batch(batch: list[str]) -> list[list[float]]:
15        embeddings = model.get_embeddings(batch)
16        return [e.values for e in embeddings]
17
18    with ThreadPoolExecutor(max_workers=8) as executor:
19        results = list(executor.map(embed_batch, batches))
20
21    # Flatten batch results
22    return [emb for batch_result in results for emb in batch_result]
23
24# Example: embed 50,000 product descriptions
25# Serial:  50,000 API calls × ~200ms = ~2.8 hours
26# Batched: 200 batches × ~800ms (parallel, 8 workers) = ~3 minutes

Optimization 5: Result Set Compression Before Gemini Context

A BigQuery query that returns 1,000 rows of full transactional data can easily consume 30,000–50,000 tokens when serialized as JSON and injected into a Gemini prompt. The agent usually only needs summary statistics, not raw rows. Add a compression layer in the tool function that summarizes the result set before returning it to the agent:

python

1import json
2from google.cloud import bigquery
3
4MAX_ROWS_FOR_AGENT = 50   # Above this, compress the result set
5
6def compress_bq_result(rows: list[dict], query_intent: str) -> str:
7    """
8    If result has >50 rows, summarize it statistically instead of returning raw JSON.
9    This cuts token usage by 60-80% for large result sets.
10    """
11    if len(rows) <= MAX_ROWS_FOR_AGENT:
12        return json.dumps(rows)
13
14    # Compute summary statistics for numeric columns
15    if not rows:
16        return json.dumps([])
17
18    numeric_cols = [k for k, v in rows[0].items() if isinstance(v, (int, float))]
19    summary = {
20        "total_rows": len(rows),
21        "note": f"Result compressed from {len(rows)} rows to summary statistics",
22        "statistics": {},
23    }
24    for col in numeric_cols:
25        values = [r[col] for r in rows if r.get(col) is not None]
26        if values:
27            summary["statistics"][col] = {
28                "sum":  round(sum(values), 2),
29                "mean": round(sum(values) / len(values), 2),
30                "min":  round(min(values), 2),
31                "max":  round(max(values), 2),
32            }
33
34    # Include first 5 and last 5 rows as samples
35    summary["sample_rows"] = rows[:5] + rows[-5:]
36    return json.dumps(summary)

Optimization 6: Adaptive Token Budgeting

As a conversation grows longer, the available context window shrinks. A tool that returns 5,000 tokens of data on turn 1 may cause the model to exceed its context limit by turn 8. Implement adaptive budgeting: query the current token count from the session and reduce the MAX_ROWS_FOR_AGENT limit proportionally. This guarantees the conversation can continue without hitting context limits regardless of how long it runs:

python

1def get_adaptive_row_limit(session_token_count: int, model_context_limit: int = 1_000_000) -> int:
2    """
3    Returns a safe row limit based on remaining context window.
4    Ensures the agent always has room for at least 10 more turns.
5    """
6    remaining_tokens = model_context_limit - session_token_count
7    # Reserve 50% of remaining for future turns + model response
8    available_for_tool = remaining_tokens * 0.5
9
10    # Assume ~100 tokens per result row (conservative estimate)
11    tokens_per_row = 100
12    safe_limit = int(available_for_tool / tokens_per_row)
13
14    # Clamp between 5 (always return something) and 200 (max useful)
15    return max(5, min(200, safe_limit))
16
17# Usage in a tool function
18def query_sales_data(query: str, ctx) -> str:
19    session_tokens = ctx.session.token_count or 0
20    row_limit = get_adaptive_row_limit(session_tokens)
21    rows = run_bigquery(query, limit=row_limit)
22    return compress_bq_result(rows, query)

Putting It Together: The Optimized Pipeline Architecture

Agent Conversation Turn N
  ↓
Tool Call: query_sales_data(user_question)
  │
  ├── [1] Partition filter injected from schema context (pre-loaded at startup)
  ├── [2] Check Redis cache → HIT: return cached result (0ms, 0 tokens)
  │                        → MISS: continue
  ├── [3] Route to materialized view if query matches known hot-path pattern
  ├── [4] Execute BigQuery query (P95: <800ms after optimization)
  ├── [5] Compress result if >50 rows → ~70% token reduction
  ├── [6] Apply adaptive row limit based on session token count
  └── [7] Store result in Redis (TTL: 5 min)
  ↓
Agent receives compressed, token-budgeted result
  ↓
Gemini API call — system prompt served from context cache (40K tokens → $0.0002)
  ↓
Response streamed back to user

These six optimizations are not theoretical. They are in production in an enterprise retail environment, running across multiple agent workloads that collectively process millions of BigQuery rows and thousands of Gemini API calls daily. The compounding effect of applying all six simultaneously is roughly a 70% cost reduction and a 5× latency improvement over the naive implementation. Start with materialized views and context caching — those two alone typically pay for the engineering time within the first billing cycle.