Generative AI Google Brain

embeddings-bigquery-import

3 min read

Import from BigQuery into Vector Search

Source notebook

Repo path: embeddings/bigquery-import.ipynb · Open on GitHub · intermediate

Imports BigQuery embedding rows into a Vertex AI Vector Search index using the import REST API.

Summary

This notebook teaches how to load vector embedding data from a BigQuery table into a Vector Search index. It creates sample BigQuery data with embeddings, restrict columns, numeric restricts, and metadata, then creates a tree-AH Matching Engine index. It calls the regional aiplatform import endpoint with a BigQuery source mapping and checks the returned long-running operation status.

Key code patterns

Initialize project and region

PROJECT_ID = "your-project-id"
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
 
from google.cloud import aiplatform
aiplatform.init(project=PROJECT_ID, location=LOCATION)

Sets the Google Cloud project and location before creating Vector Search resources.

Create sample BigQuery source

CREATE SCHEMA import_example_dataset;
CREATE TABLE import_example_dataset.test_table (
  id INTEGER,
  embedding ARRAY <FLOAT64>,
  allow_column STRING,
  deny_column STRING,
  int_column INTEGER,
  float_column FLOAT64,
  metadata_column STRING
);

Shows the BigQuery schema expected by the import mapping: id, embedding, restricts, numeric restricts, and metadata.

Create tree-AH index

my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name="import_test_index_name",
    dimensions=3,
    approximate_neighbors_count=10,
    index_update_method="BATCH_UPDATE",
)

Creates a Vector Search index whose dimensions match the sample embedding length.

Import BigQuery rows

url = f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{my_index.resource_name}:import"
request = {
    "is_complete_overwrite": True,
    "config": {"big_query_source_config": {
        "table_path": f"bq://{PROJECT_ID}.import_example_dataset.test_table",
        "datapoint_field_mapping": {"id_column": "id", "embedding_column": "embedding"}
    }}
}
response = requests.post(url, headers=headers, json=request)

Calls the import endpoint with a BigQuery table path and datapoint field mapping.

Check import LRO

operation = response.json()["name"]
response = requests.get(
    f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{operation}",
    headers=headers,
)
if "done" in response.json():
    print("Import succeeded!" if "error" not in response.json() else "Import failed")

Tracks the long-running import operation returned by the REST request.

Models & APIs used

  • APIs / services: Vertex AI, BigQuery, Vector Search
  • SDKs / libraries: google.cloud.aiplatform, requests

When to use this

Use this pattern when embeddings already live in BigQuery and need to be batch-imported into a Vertex AI Vector Search index.

Gotchas & caveats

  • Colab requires google.colab.auth.authenticate_user(); Agent Platform Workbench does not.
  • The Google Cloud project must exist and have the Agent Platform API and BigQuery API enabled.
  • The index dimensions must match the embedding array length; the sample uses dimensions=3.
  • The import request uses gcloud auth print-access-token for the REST Authorization header.
  • The import returns a long-running operation and must be polled for completion.
  • Embedding metadata import is commented out and noted as requiring allow-listing for the Vector Search metadata preview.

Best practices

  • Initialize aiplatform with project and location before creating the index.
  • Use BATCH_UPDATE when creating the index for BigQuery import.
  • Map BigQuery columns explicitly to datapoint fields including id_column and embedding_column.
  • Include restricts and numeric_restricts mappings when filterable fields are present.
  • Check the REST response status code and print error text on failure.
  • Poll the long-running operation and inspect both done and error fields.

Graph View

Backlinks