Skip to main content

Workflow composition, failure handlers, and nodes

Running queries​

flytekit lets you run BigQuery queries by writing regular Python functions decorated with @task and passing a BigQueryConfig to the task_config parameter of the decorator. The plugin handles conversion of parameters, execution of the query, and parsing of results.

from flytekit import task, workflow
from flytekit.types.structured import StructuredDataset
from flytekitplugins.bigquery import BigQueryConfig, BigQueryTask

@task(task_config=BigQueryConfig(ProjectID="my-project", Query="SELECT * FROM iris"))
def read_bq() -> StructuredDataset:
    ...
@task(task_config=BigQueryConfig(ProjectID=PROJECT, Query=QUERY), limits=Resources(mem="500Mi"))
def full_example() -> StructuredDataset:
    return StructuredDataset(...)

A dedicated task type: BigQueryTask​

Rather than using a generic @task with a BigQueryConfig, you can use BigQueryTask, a PythonTask from flytekitplugins.bigquery.task. It stores the query and project id and automatically handles the reading of results, and lets you parametrize the query with inputs:

bq_task = BigQueryTask(
    name="plugin.bigquery.wf.simple",
    inputs=kwtypes(percent_cluster=float),
    query_template="SELECT * FROM (SELECT {{.percent_cluster}} AS percent_cluster) AS X",
    project_id=PROJECT,
    output_schema_type=iris_schema,
)

Here, {{.percent_cluster}} is substituted with the input value at execution time.

Distinguishing compile-time from run-time values​

A common source of confusion: inside a workflow, the return value of a BigQuery task is a StructuredDataset promise, not a dataframe. You must pass it to another task rather than accessing it directly:

@workflow
def wf(percent_cluster: float):
    df = bq_task(percent_cluster=percent_cluster)
    write(df)   # ok
    # print(df)  # NOT ok — df is a promise at compile time

Accessing the query​

BigQueryTask validates that the query is not empty and is a string at construction time (BigQueryTask.__init__ raises ValueError("Query is not specified") if self._query_template is not set or not a string).

The query uses jinja2 for templating. The rendered query is stored in the task's _query attribute at execution time. In BigQueryTask.execute, self._query = get_query(self._query_template, **kwargs) builds the final SQL string from the template and the inputs.

BigQueryJobConfig​

Flytekit lets you pass any BigQueryJobConfig to the task. For instance, to specify the output dataset and destination table:

BigQueryConfig(
    Query="SELECT * FROM iris",
    ProjectID=PROJECT,
    job_config=bigquery.QueryJobConfig(
        destination=f"{PROJECT}.beam.iris",
    ),
)

Alternatively, on a BigQueryTask:

BigQueryTask(
    name="...",
    project_id=PROJECT,
    query_template=QUERY,
    output_result_type=StructuredDataset,
    job_config=bigquery.QueryJobConfig(...),
)

_get_job_config in flytekitplugins/bigquery/task.py merges the user-supplied job_config with the destination_table_name:

job_config = copy.deepcopy(self._job_config or bigquery.QueryJobConfig())
if self._destination_table_name is not None:
    job_config.destination = self._destination_table_name

Handling results​

BigQuery tasks must return a StructuredDataset. BigQueryTask.execute returns a StructuredDataset whose underlying dataframe is the result of running the query:

def execute(self, **kwargs) -> StructuredDataset:
    client = self._get_bigquery_client()
    self._query = get_query(self._query_template, **kwargs)
    job = client.query(self._query, job_config=self._get_job_config())
    result = job.result()
    ...
    return StructuredDataset(df=result.to_dataframe())

The result can be used by other tasks that accept a StructuredDataset, including pandas and parquet-based consumers.

How the query is built at run time​

get_query renders the query_template with jinja2:

def get_query(template: str, **kwargs) -> str:
    # Render the template with the kwargs
    return jinja2.Template(template).render(kwargs)

The kwargs come from the task invocation. Because BigQueryTask.execute calls get_query(self._query_template, **kwargs), only the declared task inputs are available in the template context.

Task metadata and resources​

Because BigQueryTask is a PythonTask, it participates in flytekit's metadata system. You can attach standard task metadata like requests, limits, retries, and cache:

bq_task = BigQueryTask(
    name="my-bq-task",
    query_template=QUERY,
    project_id=PROJECT,
    retries=2,
    cache=True,
    cache_version="1.0",
    requests=Resources(mem="512Mi"),
    limits=Resources(mem="1Gi"),
)

Error handling​

BigQueryTask validates inputs at execution time. If a rendered query is empty, execute raises:

if not self._query:
    raise ValueError("Query is empty")

Local execution​

When running a workflow that contains a BigQueryTask locally with pyflyte run or python my_wf.py, flytekit executes the query against BigQuery with the credentials from the ambient environment (Application Default Credentials, or the GOOGLE_APPLICATION_CREDENTIALS environment variable). This means a local run still requires real credentials and network access to GCP.

Summary​

Use BigQueryConfig with a plain @task when you need fine-grained control over the job configuration, and BigQueryTask when you want a reusable, parametrized query that returns a StructuredDataset. In both cases, the query is defined at task-definition time, parameter substitution happens through jinja2 at execution time, and results flow to downstream tasks as StructuredDataset promises.