DEV Community

Cover image for How I stopped cloning datasets for Vertex AI pipelines
Aniket Abhishek Soni
Aniket Abhishek Soni

Posted on

How I stopped cloning datasets for Vertex AI pipelines

You ship the job. It passes CI. Then, the SRE team pings you on Slack because your pipeline just swallowed 400GB of storage in an hour, and your cost projection for the quarter looks like a down payment on a house. You’re duplicating data again.

Most of us treat the Lakehouse as a write-once repository. We ETL, we transform, and then we export the "features" into a brand-new table or GCS bucket so Vertex AI can ingest them. It’s a habit born of convenience, but it’s a death sentence for your budget and your data consistency. In regulated environments like healthcare or finance, having two versions of a truth—one in your production warehouse and one in your training set—is a compliance nightmare waiting to happen.

The secret is simple: stop moving the data. If your features live in BigQuery or a Delta Lake, let them stay there. Use BigQuery as your source of truth and connect Vertex AI Pipelines directly to your raw or processed tables.

How it actually works

The goal is to treat your existing warehouse as the feature store. We do this by leveraging BigQuery DataFrames (BQF) or direct SQL execution within the Kubeflow Pipelines (KFP) components.

Instead of writing a Spark job to dump features into a CSV, you write a KFP component that executes a CREATE OR REPLACE TABLE or, better yet, a CREATE VIEW that points directly to your refined features.

Here is what the pattern looks like in a standard kfp.dsl.component using the Google Cloud AI Platform SDK:

from kfp.dsl import component, Input, Output, Dataset

@component(base_image="python:3.9")
def train_model_op(project_id: str, query: str):
    from google.cloud import bigquery

    # Do not materialize. Query the existing lakehouse table directly.
    client = bigquery.Client(project=project_id)
    job = client.query(query)
    df = job.to_dataframe()

    # Proceed to training
    # ...
Enter fullscreen mode Exit fullscreen mode

The real magic happens when you move from to_dataframe() to BigQuery ML (BQML) or Vertex AI’s BigqueryTabularDataset. By using aiplatform.TabularDataset.create(display_name=..., bq_source=...), you create a pointer. You are not creating a copy; you are creating a URI that tells Vertex AI: "Go look at this BigQuery table when you need to train."

When your pipeline runs, it doesn't wait for a data export job to finish. It simply validates that the schema exists. If you are using BigQuery Omni, this even works across clouds. Your pipeline triggers, reads the metadata, and maps the features into the training container. Your cost stays centered on compute, not storage.

Photo by Wilhelm Gunkel on Unsplash
Photo by Wilhelm Gunkel on Unsplash

The tradeoffs nobody mentions

If this were a silver bullet, everyone would do it. But there are jagged edges here that vendors love to gloss over.

First, performance is non-deterministic. If your feature engineering logic involves complex joins across massive partition-pruned tables, running that query during the training loop means you’re fighting the warehouse for concurrency slots. If a DBA runs a heavy ANALYZE command while your pipeline is trying to materialize features, your training job will hang. I’ve seen 4-hour training windows stretch to 12 hours because of resource contention in the underlying warehouse.

Second, the "Data Snapshot" problem. If you don't materialize, you are training on live data. This is usually what we want, but it’s dangerous. If your upstream pipeline updates a feature table mid-run, you’ve just introduced non-deterministic bias into your model. You need to implement strict partitioning—if you aren't using _PARTITIONTIME or a specific run_id as a partition key in your BigQuery tables, you are effectively training on a moving target.

Third, the Vertex AI SDK versioning is... volatile. If you are using google-cloud-aiplatform>=1.30.0, the way it handles bq_source in the Dataset resource is different from older versions. I’ve spent entire afternoons debugging an InvalidArgument error that turned out to be a deprecation of a specific BigQuery metadata mapping in the SDK. Always pin your SDK versions in your requirements.txt inside your container image. Do not rely on "latest".

Photo by Markus Spiske on Unsplash
Photo by Markus Spiske on Unsplash

When to reach for it (and when not to)

Use the "no-duplication" approach when your feature engineering is idempotent and your warehouse is the primary source of truth. If your BigQuery setup is already using dbt (data build tool) to manage dependencies, your Vertex AI pipeline should simply be an orchestrator that triggers a model build after the dbt run completes. This keeps the lineage clean and ensures you aren't reinventing the wheel.

Don't reach for this when you need extremely low-latency feature retrieval. If you are serving a model that requires sub-millisecond lookups for features, the overhead of querying a lakehouse—even with BigQuery’s caching—is too high. For that, you need a dedicated feature store like Feast or the Vertex AI Feature Store. The "don't duplicate" pattern is for training/batch inference, not for real-time inference.

Also, avoid this if your data scientists are still in the "exploration" phase. If they are constantly dropping columns, renaming features, and blowing up schemas, you don't want them messing with your production BigQuery tables. Let them work in a sandbox, duplicate all they want, and only transition to the "no-duplication" pipeline once the feature logic is hardened into a production dbt model.

Conclusion

Data duplication is a tax on your engineering team and your cloud bill. By treating BigQuery or your Lakehouse as an API for your training pipeline rather than a storage locker, you regain control over the lineage. It forces you to be disciplined about schemas and partitioning, which is a good thing for your sanity.

Ship the pipeline. Keep the data in place. Stop paying for the same bytes twice.

Cover photo by Albert Stoynov on Unsplash.

Top comments (1)

Collapse
 
koda2026 profile image
Harun - solo dev •

aniket, this is a brutally honest and highly pragmatic breakdown of the vertex ai + bigquery integration.

the "400gb storage swallow" is a rite of passage, and moving to bq_source pointers is the exact right architectural pivot for maintaining clean batch training lineage.

your point about the "data snapshot" problem is critical. training on a moving target is a silent killer. for anyone reading, the perfect countermeasure to this in bq is using time travel: CREATE OR REPLACE TABLE training_snapshot AS SELECT * FROM feature_table FOR SYSTEM_TIME AS OF '<pipeline_start_timestamp>'. it gives you a deterministic, zero-copy logical snapshot without the storage tax of a full materialized export.

regarding the resource contention you mentioned (the dba running heavy analyze commands starving the training loop): have you found success in setting up dedicated bigquery reservations/slots specifically for the vertex ai training service account? or do you rely on zero-copy clones as a middle ground to isolate the training read workload from the production warehouse concurrency?

also, +100 on pinning the google-cloud-aiplatform sdk version. relying on "latest" in ml ops is a guaranteed way to wake up to a broken pipeline. would love to hear how you've tuned the slot allocation for this! 🐯