Skip to content

Airflow with AWS compute

Apache Airflow can schedule Polars Nanook queries. The compute-intensive work runs on Polars Nanook compute, so the Airflow instance stays focused on orchestration.

Prerequisites

The integration requires:

  • Airflow 3.1 or later
  • The polars and polars-nanook packages installed in your Airflow environment (installation guide)
  • A Polars Nanook service account with client ID and secret

Credentials

An Airflow Connection holds the Polars Nanook credentials. This is the recommended approach for managing secrets in Airflow workflows.

In the Airflow UI, navigate to Admin > Connections > Add Connection and configure:

  • Connection ID: polars_nanook
  • Connection Type: Generic
  • Login: Your service account client ID
  • Password: Your service account client secret

Alternative credential storage

Credentials can also be stored in a Secrets Backend for enhanced security. This prevents storing secrets in the Airflow database and integrates with external secret managers like AWS Secrets Manager, HashiCorp Vault, or Google Secret Manager.

Task authentication

A task retrieves the connection credentials and authenticates to Polars Nanook:

from airflow.sdk import BaseHook
import polars_nanook as pn

conn = BaseHook.get_connection("polars_nanook")
pn.authenticate(client_id=conn.login, client_secret=conn.password)

For DAGs with multiple tasks, a reusable authentication decorator avoids repeating this code:

from functools import wraps
from airflow.sdk import BaseHook
import polars_nanook as pn


def authenticate(fn):
    @wraps(fn)
    def authenticated_fn(*args, **kwargs):
        conn = BaseHook.get_connection("polars_nanook")
        pn.authenticate(client_id=conn.login, client_secret=conn.password)
        return fn(*args, **kwargs)

    return authenticated_fn

Apply this decorator to any task requiring Polars Nanook access:

@task()
@authenticate
def my_query_task():
    # Task code with authenticated Polars Nanook access
    ...

A single query in a DAG

A basic DAG that executes a single Polars query on Polars Nanook:

from datetime import datetime

import polars as pl
import polars_nanook as pn

from airflow.sdk import BaseHook, dag, task


@dag(start_date=datetime(2026, 1, 1), schedule="@daily")
def single_query_dag():
    @task()
    def aggregate_sales():
        conn = BaseHook.get_connection("polars_nanook")
        pn.authenticate(client_id=conn.login, client_secret=conn.password)

        ctx = pn.ComputeContext(
            workspace="My workspace", cpus=8, memory=16, cluster_size=1
        )

        (
            pl.scan_parquet(
                "s3://<BUCKET>/input-data/*.parquet"
            )
            .group_by("status")
            .agg(pl.len())
            .remote(ctx)
            .sink_parquet("s3://<BUCKET>/output-data/")
        )

    aggregate_sales()


single_query_dag()

This example defines a DAG that runs daily, creates a ComputeContext with specific hardware requirements, and executes a query that writes results to S3.

Fire-and-forget execution

The .sink_parquet() method returns as soon as the query is submitted (after starting the cluster if it is not running yet), so the Airflow task completes while the query continues executing on the remote cluster. The cluster shuts down automatically after the configured idle timeout (default: 1 hour).

Query failure as task failure

The fragments below assume the imports and the authenticate decorator from the sections above (from datetime import datetime, import polars as pl, import polars_nanook as pn, from airflow.sdk import BaseHook, TriggerRule, dag, task). The complete example at the end of the page is self-contained.

To ensure query completion and show failures in Airflow, wait for the query result and check its status:

@task()
@authenticate
def aggregate_sales_with_check():
    ctx = pn.ComputeContext(
        workspace="My workspace", cpus=8, memory=16, cluster_size=1
    )

    query_result = (
        pl.scan_parquet("s3://<BUCKET>/input-data/*.parquet")
        .group_by("status")
        .agg(pl.len())
        .remote(ctx)
        .sink_parquet("s3://<BUCKET>/output-data/")
    )

    query_result.await_result()

    query_status = query_result.get_status()
    if query_status != pn.QueryStatus.SUCCESS:
        raise ValueError(f"Query failed with status: {query_status}")

The .await_result() method blocks until the query completes and raises the query's error if it failed, and the status check raises an exception that marks the Airflow task as failed. On success the task finishes normally. On any other status, such as a canceled query, it raises ValueError and Airflow marks the task failed.

Cluster lifecycle

Several patterns manage the compute context lifecycle (the cluster started by a ComputeContext).

Context wrapper for immediate shutdown

ComputeContext as a context manager shuts down the cluster immediately after the scope exits, which avoids idle time and its cost:

@task()
@authenticate
def query_with_immediate_shutdown():
    with pn.ComputeContext(
        workspace="My workspace", cpus=8, memory=16, cluster_size=1
    ) as ctx:
        query_result = (
            pl.scan_parquet("s3://<BUCKET>/input-data/*.parquet")
            .group_by("status")
            .agg(pl.len())
            .remote(ctx)
            .sink_parquet("s3://<BUCKET>/output-data/")
        )

        query_result.await_result()
        if query_result.get_status() != pn.QueryStatus.SUCCESS:
            raise ValueError("Query failed")
    # Cluster shuts down here when exiting the context

Context manager behavior

The cluster shuts down when the context exits, regardless of running queries. Wait for query completion (.await_result()) before exiting the context.

Named manifests for reusable clusters

Instead of specifying hardware requirements in every task, a named manifest defines a reusable cluster configuration.

A manifest is created programmatically, from an environment with the same Polars and Python versions as the Airflow workers:

pn.ComputeContext(
    workspace="My workspace", cpus=16, memory=32, cluster_size=3
).register("airflow-production")

or through the portal under Compute > Manifests > Create new manifest.

Tasks then reference the manifest by name:

@task()
@authenticate
def query_using_manifest():
    ctx = pn.ComputeContext(workspace="My workspace", name="airflow-production")

    (
        pl.scan_parquet("s3://<BUCKET>/input-data/*.parquet")
        .group_by("status")
        .agg(pl.len())
        .remote(ctx)
        .sink_parquet("s3://<BUCKET>/output-data/")
    )

Manifests provide several benefits:

  • Hardware configuration defined once, reused across tasks and users
  • Multiple queries can share the same cluster, reducing startup overhead
  • Queries are automatically queued when the cluster is busy
  • Simplifies DAG code by removing hardware specification details

Manual cluster shutdown

To manually shut down a cluster immediately after all tasks complete, add a dedicated shutdown task:

@dag(schedule="@daily", start_date=datetime(2026, 1, 1))
def dag_with_manual_shutdown():
    @task()
    @authenticate
    def query_1():
        ...

    @task()
    @authenticate
    def query_2():
        ...

    # requires: from airflow.sdk import TriggerRule
    @task(trigger_rule=TriggerRule.ALL_DONE)
    @authenticate
    def shutdown_cluster():
        ctx = pn.ComputeContext(workspace="My workspace", name="airflow-production")
        ctx.stop()

    [query_1(), query_2()] >> shutdown_cluster()


dag_with_manual_shutdown()

The trigger_rule=TriggerRule.ALL_DONE ensures the shutdown task runs after all upstream tasks complete, regardless of success or failure. Query tasks that sink must end with .await_result(), or the shutdown task stops the cluster while their queries still run.

Configuring idle timeout

The ComputeContext accepts an idle_timeout_mins parameter (default: 60 minutes, minimum: 10 minutes) that controls automatic shutdown after inactivity. For short-running workflows, reduce this value to minimize costs: ComputeContext(..., idle_timeout_mins=10).

Concurrent queries on one cluster

Multiple queries run concurrently on the same cluster when they are marked as single-node queries. This uses the resources efficiently when individual queries do not require the full cluster:

@dag(schedule="@daily", start_date=datetime(2026, 1, 1))
def parallel_queries_dag():
    @task()
    @authenticate
    def query_1():
        ctx = pn.ComputeContext(workspace="My workspace", name="airflow-production")
        (
            pl.scan_parquet("s3://<BUCKET>/dataset-1/*.parquet")
            .group_by("category")
            .agg(pl.sum("amount"))
            .remote(ctx)
            .single_node()
            .sink_parquet("s3://<BUCKET>/output-1/")
        )

    @task()
    @authenticate
    def query_2():
        ctx = pn.ComputeContext(workspace="My workspace", name="airflow-production")
        (
            pl.scan_parquet("s3://<BUCKET>/dataset-2/*.parquet")
            .group_by("region")
            .agg(pl.mean("value"))
            .remote(ctx)
            .single_node()
            .sink_parquet("s3://<BUCKET>/output-2/")
        )

    @task()
    @authenticate
    def query_3():
        ctx = pn.ComputeContext(workspace="My workspace", name="airflow-production")
        (
            pl.scan_parquet("s3://<BUCKET>/dataset-3/*.parquet")
            .group_by("date")
            .agg(pl.len())
            .remote(ctx)
            .single_node()
            .sink_parquet("s3://<BUCKET>/output-3/")
        )

    query_1()
    query_2()
    query_3()


parallel_queries_dag()

The .single_node() method indicates the query runs on a single worker, enabling the scheduler to execute multiple queries concurrently on a multi-node cluster. Without this, queries use .distributed() by default, which reserves the entire cluster and queues subsequent queries.

Results across tasks

A pipeline whose later tasks depend on results from earlier tasks passes intermediate result locations between stages:

@dag(schedule="@daily", start_date=datetime(2026, 1, 1))
def multistage_pipeline():
    @task()
    @authenticate
    def stage_1_transform() -> list[str]:
        ctx = pn.ComputeContext(workspace="My workspace", name="airflow-production")
        query_result = (
            pl.scan_parquet("s3://<BUCKET>/raw-data/*.parquet")
            .filter(pl.col("status") == "active")
            .with_columns(pl.col("amount").cast(pl.Float64))
            .remote(ctx)
            .execute()
        )

        if query_result.location is None:
            raise ValueError("Query result location is None")

        return query_result.location

    @task()
    @authenticate
    def stage_2_aggregate(result_locations: list[str]):
        lf = pl.scan_parquet(result_locations)
        ctx = pn.ComputeContext(workspace="My workspace", name="airflow-production")

        (
            lf.group_by("category")
            .agg(pl.sum("amount"), pl.len())
            .remote(ctx)
            .sink_parquet("s3://<BUCKET>/aggregated-data/")
        )

    stage_1_results = stage_1_transform()
    stage_2_aggregate(stage_1_results)


multistage_pipeline()

The first stage executes a query without .sink_parquet(), instead using .execute() to retrieve the query_result.location. This contains a list of pre-signed URLs to the intermediate results Polars Nanook stores, valid for 8 hours. These URLs are serialized and passed to the next stage, which scans them with pl.scan_parquet() to continue processing. Anyone holding a URL can read the result, and a stage that runs more than 8 hours later can no longer read it.

Intermediate result retention

Intermediate results stored by Polars Nanook are automatically deleted after one day. For persistent storage, use .sink_parquet() with your own S3 path. See the remote query documentation for more details.

Summary

  • Credentials: an Airflow Connection (polars_nanook, type Generic, login = client ID, password = client secret) or a Secrets Backend.
  • Waiting for the query: not required. .sink_parquet() returns once the query is submitted. .await_result() followed by a check of get_status() == pn.QueryStatus.SUCCESS surfaces failures.
  • Cluster stop: after idle_timeout_mins (default 60, minimum 10), when a ComputeContext context manager exits, or when ctx.stop() is called.
  • Shared cluster: reference a named manifest and mark queries .single_node(), because .distributed() (the default) reserves the whole cluster and queues others.
  • Chained stages: return query_result.location (pre-signed URLs, valid for 8 hours) from one task and pl.scan_parquet() it in the next. Intermediate results are deleted after one day.

Complete example

The following example combines all patterns: authentication decorator, named manifest, parallel execution, multi-stage pipeline, and manual cluster shutdown:

from datetime import datetime
from functools import wraps

import polars as pl
import polars_nanook as pn

from airflow.sdk import BaseHook, TriggerRule, dag, task


def authenticate(fn):
    @wraps(fn)
    def authenticated_fn(*args, **kwargs):
        conn = BaseHook.get_connection("polars_nanook")
        pn.authenticate(client_id=conn.login, client_secret=conn.password)
        return fn(*args, **kwargs)

    return authenticated_fn


WORKSPACE = "My workspace"
MANIFEST_NAME = "airflow-production"


@dag(schedule="@daily", start_date=datetime(2026, 1, 1))
def complete_pipeline():
    @task()
    @authenticate
    def transform_sales() -> list[str]:
        ctx = pn.ComputeContext(workspace=WORKSPACE, name=MANIFEST_NAME)
        query_result = (
            pl.scan_parquet("s3://<BUCKET>/sales/*.parquet")
            .with_columns((pl.col("price") * pl.col("quantity")).alias("revenue"))
            .remote(ctx)
            .single_node()
            .execute()
        )
        query_status = query_result.status
        if query_status != pn.QueryStatus.SUCCESS:
            raise ValueError(f"Query failed with status: {query_status}")
        if query_result.location is None:
            raise ValueError("Sales transformation failed")
        return query_result.location

    @task()
    @authenticate
    def transform_inventory() -> list[str]:
        ctx = pn.ComputeContext(workspace=WORKSPACE, name=MANIFEST_NAME)
        query_result = (
            pl.scan_parquet("s3://<BUCKET>/inventory/*.parquet")
            .with_columns(
                (pl.col("stock_level") < pl.col("min_threshold")).alias("low_stock")
            )
            .remote(ctx)
            .single_node()
            .execute()
        )
        query_status = query_result.status
        if query_status != pn.QueryStatus.SUCCESS:
            raise ValueError(f"Query failed with status: {query_status}")
        if query_result.location is None:
            raise ValueError("Inventory transformation failed")
        return query_result.location

    @task()
    @authenticate
    def transform_returns() -> list[str]:
        ctx = pn.ComputeContext(workspace=WORKSPACE, name=MANIFEST_NAME)
        query_result = (
            pl.scan_parquet("s3://<BUCKET>/returns/*.parquet")
            .with_columns((pl.col("return_date") - pl.col("purchase_date")).alias("return_days"))
            .remote(ctx)
            .single_node()
            .execute()
        )
        query_status = query_result.status
        if query_status != pn.QueryStatus.SUCCESS:
            raise ValueError(f"Query failed with status: {query_status}")
        if query_result.location is None:
            raise ValueError("Returns transformation failed")
        return query_result.location

    @task()
    @authenticate
    def generate_report(
        sales_locations: list[str],
        inventory_locations: list[str],
        returns_locations: list[str],
    ):
        sales = pl.scan_parquet(sales_locations)
        inventory = pl.scan_parquet(inventory_locations)
        returns = pl.scan_parquet(returns_locations)

        ctx = pn.ComputeContext(workspace=WORKSPACE, name=MANIFEST_NAME)
        query = (
            sales.join(inventory, on="product_id")
            .join(returns, on="product_id", how="left")
            .group_by("product_category")
            .agg(
                pl.sum("revenue").alias("total_revenue"),
                pl.sum("low_stock").alias("low_stock_count"),
                pl.len().alias("transaction_count"),
            )
            .remote(ctx)
            .distributed()
            .sink_parquet("s3://<BUCKET>/daily-report/")   
        )
        query_result = query.await_result()
        query_status = query_result.status
        if query_status != pn.QueryStatus.SUCCESS:
            raise ValueError(f"Query failed with status: {query_status}")

    @task(trigger_rule=TriggerRule.ALL_DONE)
    @authenticate
    def shutdown_cluster():
        ctx = pn.ComputeContext(workspace=WORKSPACE, name=MANIFEST_NAME)
        ctx.stop()

    sales_results = transform_sales()
    inventory_results = transform_inventory()
    returns_results = transform_returns()
    report_task = generate_report(sales_results, inventory_results, returns_results)
    report_task >> shutdown_cluster()


complete_pipeline()