DatabricksTaskOperator¶

Use the DatabricksTaskOperator to launch and monitor task runs on Databricks as Airflow tasks. This can be used as a standalone operator in a Dag and as well as part of a Databricks Workflow by using it as an operator(task) within the DatabricksWorkflowTaskGroup.

Examples¶

Running a notebook in Databricks using DatabricksTaskOperator¶

tests/system/databricks/example_databricks.py[source]

    task_operator_nb_1 = DatabricksTaskOperator(
        task_id="nb_1",
        databricks_conn_id="databricks_conn",
        job_cluster_key="Shared_job_cluster",
        task_config={
            "notebook_task": {
                "notebook_path": "/Shared/Notebook_1",
                "source": "WORKSPACE",
            },
            "libraries": [
                {"pypi": {"package": "Faker"}},
                {"pypi": {"package": "simplejson"}},
            ],
        },
    )

Running a SQL query in Databricks using DatabricksTaskOperator¶

tests/system/databricks/example_databricks.py[source]

    task_operator_sql_query = DatabricksTaskOperator(
        task_id="sql_query",
        databricks_conn_id="databricks_conn",
        task_config={
            "sql_task": {
                "query": {
                    "query_id": QUERY_ID,
                },
                "warehouse_id": WAREHOUSE_ID,
            }
        },
    )

Configuring Databricks-native task retries¶

Use max_retries, min_retry_interval_millis and retry_on_timeout to configure Databricks-native task retries. Databricks reruns failed task attempts within the same job run, so Airflow sees only the final result. Set max_retries to -1 to retry indefinitely, or 0 to disable retries.

These settings are independent of the Airflow task-level retries parameter, which retries the whole Airflow task. You can set the same fields directly in task_config. When both are set, the operator parameter takes precedence. If a field is unset, Databricks uses its default.

Airflow retries behaves differently depending on where the operator runs. For a standalone operator, each retry submits a new Databricks run. Inside a DatabricksWorkflowTaskGroup, the Airflow task monitors a sub-run that was already submitted by the workflow launch task, so a retry only re-polls the terminal sub-run. Use max_retries to retry Databricks work inside a workflow task group.

Inside a DatabricksWorkflowTaskGroup, a task that exhausts a finite max_retries is reported as failed as soon as its final failed attempt is observed, so downstream failure handling is not delayed by long-running sibling tasks. Only unlimited retries (max_retries=-1) keep the Airflow task waiting (or deferring) until the parent workflow run reaches a terminal state, because Databricks may still launch a retry attempt under the same task_key until then. Sibling tasks in the run continue independently.