airflow_databricks_submit_run_operator_spark_jar_task.py

1import datetime
2
3from airflow import DAG
4from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
5
6with DAG(
7    dag_id='example_databricks_operator',
8    schedule=None,
9    start_date=datetime.datetime(2021, 1, 1),
10    catchup=False,
11    tags=['example'],
12) as dag:
13    # [START howto_operator_databricks_json]
14    # Example of using the DatabricksSubmitRunOperator to submit a Spark JAR task
15    new_cluster = {
16        'spark_version': '7.3.x-scala2.12',
17        'node_type_id': 'r3.xlarge',
18        'num_workers': 2,
19    }
20
21    notebook_task = {
22        'notebook_path': '/Users/airflow@example.com/PrepareData',
23    }
24
25    spark_jar_task = {
26        'main_class_name': 'com.example.ProcessData',
27        'parameters': ['--input', 'dbfs:/input.json', '--output', 'dbfs:/output.json'],
28    }
29
30    submit_run = DatabricksSubmitRunOperator(
31        task_id='submit_run',
32        new_cluster=new_cluster,
33        spark_jar_task=spark_jar_task,
34        libraries=[
35            {'jar': 'dbfs:/lib/example-job.jar'},
36        ],
37    )
38    # [END howto_operator_databricks_json]