How to skip Airflow SparkSubmitOperator task based on exit code that my Spark programm returns?

Question

My Spark submit application is doing some query and returning different exit code depends on the dataset state.

Is it possible to skip downstream tasks right after my spark-submit operator? I am thinking about skip_exit_code feature of BashOperator, which is surprisingly missing in all other operators.

def spark_job(task_id: str, cfg: ArgList, main_class: str, dag: DAG) -> BaseOperator:
    copy_args = cfg.to_arg_list()

    return SparkSubmitOperator(
        task_id=task_id,
        conn_id='spark_default',
        java_class=main_class,
        application=SPARK_JOBS_JAR,
        application_args=copy_args,
        total_executor_cores='2',
        executor_cores='1',
        executor_memory='1g',
        num_executors='1',
        name=task_id,
        verbose=False,
        driver_memory='1g',
        dag=dag
    )
cfg = CheckDataCfg(...)
check_data_task = spark_job('check-data', cfg, 'etljobs.spark.CheckDataRecieved', dag)

check_data_task >> '<my next task which I need to skip sometimes>'

UPDATE: Current SparkSubmitHook implementation does throw an exception if returncode is not 0. So there are only two workarounds which I found later:

1. Create custom SparkSubmitHook and SparkSubmitOperator classes to ignore user-defined non-zero exit codes in way that either AirflowSkipException exception is going to be thrown or return code value will be pushed to XCom for further use.
2. Use BashOperator instead. It already support skip_exit_code feature. You would need to construct all CLI Spark args manually in Python which is not a big deal.

Answer 1

The SparkSubmitHook has _spark_exit_code that can be used here. We can create a custom operator that inherits all SparkSubmitOperator functionality with addition of returning the _spark_exit_code value.

I didn't test it but I think the following code should work for you:

from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
from airflow.operators.python_operator import ShortCircuitOperator

def shortcircuit_fn(**context):
    exit_code = context['ti'].xcom_pull(task_ids='check-data')
    if exit_code=='125': # Place the codes here
        return True
    return False


class MySparkSubmitOperator(SparkSubmitOperator):

    def execute(self, context):
        super().execute(context)
        return self._hook._spark_exit_code

with DAG(dag_id='spar',
         default_args=default_args,
         schedule_interval=None,
         ) as dag:
    spark_op = MySparkSubmitOperator(task_id='check-data',..., do_xcom_push=True)
    short_op = ShortCircuitOperator(task_id='short_circuit', python_callable=shortcircuit_fn)
    next_op = AnyOperator()
    spark_op >> short_op >> next_op

This is how it works: MySparkSubmitOperator will push to xcom the value of _spark_exit_code then ShortCircuitOperator will verify it against the expected codes if the condition is met the workflow will continue if not it will mark all downstream tasks as skipped.