Big Query

This example demonstrates how to use the Big Query connector in ydata-sdk.

Example Code

"""Google Cloud Big Query example."""

from ydata.connectors import BigQueryConnector
from ydata.utils.formats import read_json


def get_token(token_path: str):
    "Utility to load a token from .secrets"
    return read_json(token_path)


if __name__ == "__main__":

    # Load the secret token
    token = get_token("gcs_credentials.json")

    # Instantiate the Connector
    connector = BigQueryConnector(
        project_id="ydatasynthetic",
        keyfile_dict=token,
    )

    # --- Read a full dataset ---
    # Data is streamed via the BigQuery Storage API (Arrow) for faster transfer.
    data = connector.query(
        query="SELECT * FROM `ydatasynthetic.dataset_test.cardio_data`"
    )
    nrows, ncols = data.shape()
    print(f"Full dataset: {nrows} rows, {ncols} columns.")

    # --- Read a sample ---
    # Uses INTERACTIVE priority for lower latency on smaller result sets.
    data_sample = connector.query_sample(
        query="SELECT * FROM `ydatasynthetic.dataset_test.cardio_data`",
        n_sample=10_000,
    )
    print(f"Sample dataset: {data_sample.nrows} rows, {data_sample.ncols} columns.")

    # --- Inspect available datasets and tables ---
    print("Available datasets:", connector.datasets)
    print("Tables in dataset_test:", connector.list_tables("dataset_test"))

    # --- Write data back to BigQuery ---
    # Uses the native BigQuery client (Parquet serialisation) — faster than pandas_gbq.
    connector.write_table(
        data=data_sample,
        database="dataset_test",
        table="cardio_sample",
    )

    # --- Export a query result directly to GCS ---
    # BigQuery does not support writing query results directly to GCS,
    # so write_query_to_gcs uses a temporary BQ table as an intermediate step.
    connector.write_query_to_gcs(
        query="SELECT age, gender, cardio FROM `ydatasynthetic.dataset_test.cardio_data`",
        path="gs://ydata_development/connectors/cardio_export.csv",
        clean_tmp=True,
    )