Generative AI Google Brain

gemini-use-cases-applying-llms-to-data-semantic-search-in-bigquery-stackoverflow-questions-semantic-search

3 min read

Performing Semantic Search in BigQuery

Source notebook

Repo path: gemini/use-cases/applying-llms-to-data/semantic-search-in-bigquery/stackoverflow_questions_semantic_search.ipynb · Open on GitHub · intermediate

Builds semantic search over Stack Overflow questions in BigQuery with Vertex AI text embeddings.

Summary

It shows how to enable Vertex AI and BigQuery Connection APIs, create a BigQuery Cloud resource connection, grant Vertex AI User to the connection service account, and define a BigQuery ML remote model for text-embedding-005. It embeds 10,000 iOS-tagged Stack Overflow questions with ML.GENERATE_EMBEDDING, creates an IVF cosine vector index, verifies async index readiness, and runs VECTOR_SEARCH to return the top 5 similar questions.

Key code patterns

Cloud resource connection

!gcloud services enable aiplatform.googleapis.com bigqueryconnection.googleapis.com
!bq mk --connection --location=us --connection_type=CLOUD_RESOURCE vertex_conn
!gcloud projects add-iam-policy-binding $PROJECT_ID \
  --member=serviceAccount:$SERVICE_ACCT_EMAIL \
  --role=roles/aiplatform.user

BigQuery uses the connection service account to call Vertex AI remote models.

Remote embedding model

CREATE OR REPLACE MODEL `bigquery_demo.text_embedding_005`
REMOTE WITH CONNECTION `us.vertex_conn`
OPTIONS (endpoint = 'text-embedding-005')

Registers the Vertex AI embedding endpoint as a BigQuery ML remote model.

Generate table embeddings

CREATE OR REPLACE TABLE `bigquery_demo.posts_questions_embedding` AS
SELECT *
FROM ML.GENERATE_EMBEDDING(
  MODEL `bigquery_demo.text_embedding_005`,
  (SELECT id, title, body, CONCAT(title, body) AS CONTENT
   FROM `bigquery-public-data.stackoverflow.posts_questions`
   WHERE tags LIKE '%ios%'
   ORDER BY view_Count DESC
   LIMIT 10000),
  STRUCT(TRUE AS flatten_json_output, 'SEMANTIC_SIMILARITY' AS task_type));

Embeds selected Stack Overflow title and body text for semantic similarity search.

Vector index

CREATE OR REPLACE VECTOR INDEX ix_posts_questions
ON `bigquery_demo.posts_questions_embedding` (ml_generate_embedding_result)
OPTIONS(index_type = 'IVF', distance_type = 'COSINE',
  ivf_options = '{"num_lists":500}');

Creates an IVF cosine index so VECTOR_SEARCH can search embeddings efficiently.

Query by embedding

SELECT query.query, base.id, base.title, distance
FROM VECTOR_SEARCH(
  TABLE `bigquery_demo.posts_questions_embedding`,
  'ml_generate_embedding_result',
  (SELECT ml_generate_embedding_result, content AS query
   FROM ML.GENERATE_EMBEDDING(MODEL `bigquery_demo.text_embedding_005`,
     (SELECT 'Why does my iOS app crash with a low memory warning despite minimal memory usage?' AS content))),
  top_k => 5,
  OPTIONS => '{"fraction_lists_to_search": 0.10}')
ORDER BY distance ASC;

Embeds a new question and retrieves the closest stored questions by vector distance.

Models & APIs used

  • Models: text-embedding-005
  • APIs / services: BigQuery, BigQuery ML, Vertex AI API, BigQuery Connection API

When to use this

Use this pattern for semantic search over BigQuery text tables using Vertex AI embeddings and BigQuery vector indexes.

Gotchas & caveats

  • Vertex AI API and BigQuery Connection API must be enabled before setup.
  • The Cloud resource connection service account needs roles/aiplatform.user.
  • The notebook assumes PROJECT_ID comes from the GOOGLE_CLOUD_PROJECT environment variable.
  • Resources are created in the US location using bigquery_demo and us.vertex_conn.
  • Remote model permission errors may require waiting a minute and retrying.
  • Vector indexes populate asynchronously; check coverage_percentage > 0 and last_refresh_time is not NULL.
  • BigQuery, BigQuery ML, and Vertex AI API are billable; the demo limits input to 10,000 iOS posts.

Best practices

  • Create a BigQuery Cloud resource connection before defining the remote model.
  • Grant Vertex AI User to the connection service account used by BigQuery.
  • Use SEMANTIC_SIMILARITY as the embedding task type for search embeddings.
  • Concatenate question title and body before embedding to include both fields.
  • Create an IVF cosine vector index on ml_generate_embedding_result before VECTOR_SEARCH.
  • Query INFORMATION_SCHEMA.VECTOR_INDEXES to verify index readiness.
  • Order VECTOR_SEARCH results by ascending distance.
  • Delete the demo dataset and connection when cleaning up.

Graph View

Backlinks