import torch
from torch_data_bigquery import BigQueryStorageDataset

dataset = BigQueryStorageDataset(
    billing_project=PROJECT,
    selected_fields=[
        "PassengerId",
        "Survived",
        "Pclass",
    ],
    location=f"bq://your-gcp-project.kaggle_titanic.train",
)

dataloader = torch.utils.data.DataLoader(dataset)

import torch
from torch_data_bigquery import BigQueryDataset

dataset = BigQueryDataset(
    billing_project=PROJECT,
    query=f"""
        SELECT
            Survived                                AS survived,    # 1
            Pclass                                  AS pclass,      # 2
            DENSE_RANK() OVER(ORDER by Sex)         AS sex,         # 3   
            COALESCE(Age, AVG(Age) OVER())          AS age,         # 4
            SibSp                                   AS siblings,    # 5
            Parch                                   AS parents,     # 6
            Fare                                    AS fare,        # 7
            DENSE_RANK() OVER(ORDER by Embarked)    AS embarked,    # 8
        FROM `your-gcp-project.kaggle_titanic.train`
    """,
)

dataloader = torch.utils.data.DataLoader(dataset)